@@ -6,7 +6,6 @@ from django.contrib import admin from django.contrib.admin import BooleanFieldListFilter, DateFieldListFilter from django.contrib.admin.models import LogEntry from django.contrib.auth.admin import UserAdmin -from django.contrib.contenttypes.models import ContentType from django.db.models import Count, QuerySet from django.http import HttpRequest, HttpResponse from django.utils import timezone @@ -401,15 +400,13 @@ class CompanyIPWhitelistAdmin(admin.ModelAdmin): inlines = [CompanyIPInline] def log_change(self, request: HttpRequest, object: Any, message: list[dict[str, any]]) -> LogEntry: - ct = ContentType.objects.get_for_model(object, for_concrete_model=False) return [ - LogEntry.objects.log_action( - request.user.pk, - ct.pk, - object.pk, - str(object), - 1, - [message_entry], + LogEntry.objects.log_actions( + user_id=request.user.pk, + queryset=[object], + action_flag=1, + change_message=[message_entry], + single_object=True, ) for message_entry in message if not message_entry.get('changed', None) @@ -424,5 +421,5 @@ class LogEntryAdmin(admin.ModelAdmin): def log_change(self, *args, **kwargs) -> LogEntry: return None - def log_deletion(self, *args, **kwargs) -> LogEntry: - return None + def log_deletions(self, *args, **kwargs) -> list[LogEntry]: + return [] @@ -9,6 +9,3 @@ class AuthenticationConfig(AppConfig): default_auto_field = 'django.db.models.BigAutoField' name = 'authentication' verbose_name = 'Пользователи' - - def ready(self): - from .signals import invalidate_user_cache @@ -1,19 +0,0 @@ -from cacheops import cache -from cacheops.getset import dnfs_to_conj_keys - -from authentication.models import BusinessAccount, CustomUserModel - -from django.db.models.signals import post_save, post_delete -from django.dispatch import receiver - - -@receiver([post_save, post_delete], sender=BusinessAccount) -def invalidate_user_cache(sender, instance, signal, **kwargs): - cache_keys = cache.conn.smembers( - dnfs_to_conj_keys('', {'authentication_customusermodel': [{'uid': instance.user_id}]})[0] - ) - for key in cache_keys: - data = cache.get(key.decode()) - if isinstance(data, list) and isinstance((user := data[0]), CustomUserModel): - user.business_account = instance if signal == post_save else None - cache.set(key.decode(), [user]) @@ -40,6 +40,7 @@ SYSTEM_APPS = [ 'django.contrib.contenttypes', 'django.contrib.sessions', 'django.contrib.messages', + 'django.contrib.postgres', 'django.contrib.staticfiles', ] @@ -54,12 +55,12 @@ EXTERNAL_APPS = [ 'drf_social_oauth2', 'dj_rest_auth', 'dj_rest_auth.registration', - 'django_minio_backend', + 'django_minio_backend.apps.DjangoMinioBackendConfig', 'drf_spectacular', 'drf_spectacular_sidecar', 'ordered_model', 'import_export', - 'cacheops', + 'cachalot', 'django_celery_beat', ] @@ -212,20 +213,31 @@ JWT_SETTINGS = { DEFAULT_AUTO_FIELD = 'django.db.models.BigAutoField' -STORAGES = { - 'default': {'BACKEND': 'django_minio_backend.models.MinioBackend'}, - 'staticfiles': {'BACKEND': 'django.contrib.staticfiles.storage.StaticFilesStorage'}, -} +# MinIO +MINIO_BUCKET_CHECK_ON_SAVE = env.bool( + 'MINIO_BUCKET_CHECK_ON_SAVE', + default=False, +) -MINIO_BUCKET_CHECK_ON_SAVE = env.bool('MINIO_BUCKET_CHECK_ON_SAVE', default=False) MINIO_ENDPOINT = env.str('MINIO_ENDPOINT') MINIO_USE_HTTPS = env.bool('MINIO_USE_HTTPS', default=False) -MINIO_EXTERNAL_ENDPOINT = env.str('MINIO_EXTERNAL_ENDPOINT', default='localhost:9000') -MINIO_EXTERNAL_ENDPOINT_USE_HTTPS = env.bool('MINIO_EXTERNAL_ENDPOINT_USE_HTTPS', default=False) + +MINIO_EXTERNAL_ENDPOINT = env.str( + 'MINIO_EXTERNAL_ENDPOINT', + default='localhost:9000', +) +MINIO_EXTERNAL_ENDPOINT_USE_HTTPS = env.bool( + 'MINIO_EXTERNAL_ENDPOINT_USE_HTTPS', + default=False, +) MINIO_ACCESS_KEY = env.str('MINIO_ACCESS_KEY') MINIO_SECRET_KEY = env.str('MINIO_SECRET_KEY') + +MINIO_STATIC_FILES_BUCKET = 'air-static' +MINIO_DEFAULT_BUCKET = 'air-media' + MINIO_PRIVATE_BUCKETS = [ 'air-messages', 'air-achievements', @@ -235,12 +247,38 @@ MINIO_PRIVATE_BUCKETS = [ 'air-models', 'air-media-presets', 'air-voices', + MINIO_STATIC_FILES_BUCKET, + MINIO_DEFAULT_BUCKET, ] -MINIO_STATIC_FILES_BUCKET = 'air-static' -MINIO_PRIVATE_BUCKETS.append(MINIO_STATIC_FILES_BUCKET) -MINIO_MEDIA_FILES_BUCKET = 'air-media' -MINIO_PRIVATE_BUCKETS.append(MINIO_MEDIA_FILES_BUCKET) +MINIO_PUBLIC_BUCKETS = [] + +STORAGES = { + 'default': { + 'BACKEND': 'django_minio_backend.models.MinioBackend', + 'OPTIONS': { + 'MINIO_ENDPOINT': MINIO_ENDPOINT, + 'MINIO_EXTERNAL_ENDPOINT': MINIO_EXTERNAL_ENDPOINT, + 'MINIO_EXTERNAL_ENDPOINT_USE_HTTPS': ( + MINIO_EXTERNAL_ENDPOINT_USE_HTTPS + ), + 'MINIO_ACCESS_KEY': MINIO_ACCESS_KEY, + 'MINIO_SECRET_KEY': MINIO_SECRET_KEY, + 'MINIO_USE_HTTPS': MINIO_USE_HTTPS, + 'MINIO_PRIVATE_BUCKETS': MINIO_PRIVATE_BUCKETS, + 'MINIO_PUBLIC_BUCKETS': MINIO_PUBLIC_BUCKETS, + 'MINIO_DEFAULT_BUCKET': MINIO_DEFAULT_BUCKET, + 'MINIO_STATIC_FILES_BUCKET': MINIO_STATIC_FILES_BUCKET, + 'MINIO_BUCKET_CHECK_ON_SAVE': MINIO_BUCKET_CHECK_ON_SAVE, + 'MINIO_CONSISTENCY_CHECK_ON_START': False, + }, + }, + 'staticfiles': { + 'BACKEND': ( + 'django.contrib.staticfiles.storage.StaticFilesStorage' + ), + }, +} SPECTACULAR_SETTINGS = { 'TITLE': 'AIR', @@ -467,22 +505,8 @@ if (SENTRY_URL := env.str('SENTRY_URL', '')) and RELEASE and ENVIRONMENT: ], ) -CACHEOPS_REDIS = env.str('CACHEOPS_REDIS', CACHES['default']['LOCATION']) -CACHEOPS_DEGRADE_ON_FAILURE = True - -if CACHEOPS_REDIS: - CACHEOPS = { - # 'authentication.*': {'ops': 'all', 'timeout': 60 * 60}, - 'authentication.companyipwhitelist': {'ops': 'all', 'timeout': 60 * 60}, - 'ml_model.*': {'ops': 'all', 'timeout': 60 * 60}, - 'tools.chats.*': {'ops': 'all', 'timeout': 60 * 60}, - 'tools.media.*': {'ops': 'all', 'timeout': 60 * 60}, - 'payments.paymentplan': {'ops': 'all', 'timeout': 60 * 60}, - 'payments.invoice': {'ops': 'all', 'timeout': 60 * 60 * 24 * 7}, - 'messages.*': {'ops': 'all', 'timeout': 60 * 60}, - 'reports.*': {'ops': 'all', 'timeout': 60 * 60}, - 'token_blacklist.outstandingtoken': {'ops': 'get', 'timeout': 60 * 60 * 24}, - } +CACHALOT_CACHE = 'default' +CACHALOT_TIMEOUT = 60 * 60 # Unleash settings UNLEASH_API_URL = env.str('UNLEASH_API_URL', 'https://example.com') @@ -0,0 +1,21 @@ +from collections.abc import Sequence + +import pytest + + +class PytestTestRunner: + def __init__(self, verbosity: int = 1, **kwargs) -> None: + self.verbosity = verbosity + + def run_tests( + self, + test_labels: Sequence[str] | None = None, + extra_args: Sequence[str] | None = None, + **kwargs, + ) -> int: + args = list(test_labels or ('tests',)) + args.extend(extra_args or ()) + if self.verbosity > 1: + args.insert(0, f'-{"v" * min(self.verbosity, 3)}') + + return pytest.main(args) @@ -1,7 +1,7 @@ from abc import abstractmethod from typing import Any -from cacheops import invalidate_all +from cachalot.api import invalidate from django.test import TestCase from ninja.testing import TestClient from rest_framework_simplejwt.tokens import RefreshToken @@ -21,7 +21,7 @@ class BaseAPITest(TestCase): def test_unauthorized_status_code(self) -> None: ... def setUp(self): - invalidate_all() + invalidate() super().setUp() @@ -0,0 +1,48 @@ +from django.core.management.base import BaseCommand, CommandError + +from core.test_runner import PytestTestRunner + + +class Command(BaseCommand): + help = 'Run the pytest suite through the Django management interface.' + + def add_arguments(self, parser) -> None: + parser.add_argument('test_labels', nargs='*') + parser.add_argument( + '--model-slug', + action='append', + default=[], + help='Run ML model API contracts only for this slug. May be repeated.', + ) + parser.add_argument( + '--profile-resources', + action='store_true', + help='Report wall time, CPU time, and peak RSS for each test.', + ) + parser.add_argument( + '--provider-smoke', + action='store_true', + help='Run one explicitly selected model against its real provider.', + ) + parser.add_argument( + '--pytest-arg', + action='append', + default=[], + help='Forward an additional argument to pytest. May be repeated.', + ) + + def handle(self, *args, **options) -> None: + pytest_args = list(options['pytest_arg']) + test_labels = options['test_labels'] + for slug in options['model_slug']: + pytest_args.extend(('--model-slug', slug)) + if options['profile_resources']: + pytest_args.append('--profile-resources') + if options['provider_smoke']: + pytest_args.extend(('--provider-smoke', '-s')) + test_labels = ('tests/ml_models/test_provider_smoke.py',) + + runner = PytestTestRunner(verbosity=options['verbosity']) + exit_code = runner.run_tests(test_labels, extra_args=pytest_args) + if exit_code: + raise CommandError(f'pytest exited with status {exit_code}') @@ -79,7 +79,12 @@ class PaymentPlanUserInfo(BaseModel): update_fields=None, ): self.last_payment_at = datetime.now().date() - return super().save(force_insert, force_update, using, update_fields) + return super().save( + force_insert=force_insert, + force_update=force_update, + using=using, + update_fields=update_fields, + ) @property def primary_method(self): @@ -136,6 +136,8 @@ class PaymentPlanUserInfoAdmin(admin.ModelAdmin): class PaymentPlanFeatureAdmin(OrderedModelAdmin): list_display = ('plan', 'model', 'move_up_down_links') list_filter = ('plan__tokens_per_plan', 'model__category') + search_fields = ('model__title', 'model__slug') + search_help_text = _('Search by model title or slug') class PaymentAttemptInline(admin.TabularInline): @@ -0,0 +1,217 @@ +import json +from pathlib import Path + +import pytest +from django.core.files.storage import FileSystemStorage +from django.core.files.uploadedfile import SimpleUploadedFile +from rest_framework.test import APIClient + +from authentication.services.token import TokenService +from messages.models import Message +from ml_model.models import ModelCategory +from payments.models import Invoice +from tests.factories import NeuronModelFactory +from tests.ml_models.cases import EndpointKind, MEDIA_MODEL_CASES, get_model_case +from tests.ml_models.provider_fakes import install_provider_fakes +from tools.media.models import Audio, Image, Preset, Video, Voice, VoiceClone + + +MEDIA_MANAGERS = { + EndpointKind.AUDIO: Audio, + EndpointKind.IMAGE: Image, + EndpointKind.VIDEO: Video, + EndpointKind.VOICE: VoiceClone, +} +MEDIA_CASES = tuple( + pytest.param(case.endpoint, case.slug, MEDIA_MANAGERS[case.endpoint], id=case.slug) + for case in MEDIA_MODEL_CASES +) + + +@pytest.mark.django_db +@pytest.mark.usefixtures('fake_provider_proxy') +@pytest.mark.parametrize(('media_kind', 'model_slug', 'manager'), MEDIA_CASES) +def test_media_generation_api_contract( + media_kind, + model_slug, + manager, + authenticated_client, + user, + local_message_storage, + monkeypatch, +) -> None: + case = get_model_case(model_slug) + model = NeuronModelFactory(slug=model_slug) + install_provider_fakes(monkeypatch, case, model) + request_data = { + 'content': 'Generic media API contract', + 'info': dict(case.info), + } + if media_kind == EndpointKind.VOICE: + request_data['info']['voice_id'] = 'test-voice' + input_file = case.make_input_file() + if input_file: + request_data['file'] = input_file + request_data['info'] = json.dumps(request_data['info']) + + balance_before = user.payment_plan.current_token_balance + response = authenticated_client.post( + f'/api/v1/media/{media_kind.value}/{model_slug}', + request_data, + format='multipart' if input_file else 'json', + ) + + assert response.status_code == 201, response.json() + payload = response.json() + assert payload + assert all(message['from_model'] is True for message in payload) + + gallery = manager.objects.get(user=user, model=model) + output_messages = gallery.output_messages + assert output_messages.exists() + assert all(message.file for message in output_messages) + assert all(local_message_storage.exists(message.file.name) for message in output_messages) + + invoice = Invoice.objects.get(user=user, model=model) + user.payment_plan.refresh_from_db() + assert balance_before - user.payment_plan.current_token_balance == invoice.cost + + model_response = authenticated_client.get(f'/api/v1/media/{media_kind.value}/{model_slug}') + assert model_response.status_code == 200 + assert {message['uid'] for message in model_response.json()} == { + str(message.uid) for message in output_messages + } + + other_model = NeuronModelFactory(slug=f'{model_slug}_other') + other_gallery = manager.objects.create(user=user, model=other_model) + other_message = Message.objects.create( + content='Another media store message', + content_object=other_gallery, + from_model=True, + ) + + gallery_strategy = 'voice' if media_kind == EndpointKind.VOICE else f'{media_kind.value}s' + gallery_response = authenticated_client.get(f'/api/v1/media/gallery/{gallery_strategy}') + assert gallery_response.status_code == 200 + assert {message['uid'] for message in gallery_response.json()} == { + str(message.uid) for message in output_messages + } | {str(other_message.uid)} + + +@pytest.mark.django_db +def test_voice_crud_api(user, monkeypatch, tmp_path) -> None: + storage = FileSystemStorage(location=Path(tmp_path) / 'voices') + file_field = Voice._meta.get_field('file') + monkeypatch.setattr(file_field, 'storage', storage) + + token = TokenService._encode(payload={'uid': str(user.uid)}, token_type='access') + client = APIClient() + client.credentials(HTTP_AUTHORIZATION=f'Bearer {token}') + response = client.post( + '/api/v1/api/media/voices/', + { + 'file': SimpleUploadedFile('reference.mp3', b'mocked audio input', content_type='audio/mpeg'), + 'title': 'Reference voice', + 'transcription': 'Initial transcription', + }, + format='multipart', + ) + + assert response.status_code == 201 + voice = Voice.objects.get(user=user) + assert voice.title == 'Reference voice' + assert storage.exists(voice.file.name) + + response = client.get('/api/v1/api/media/voices/') + assert response.status_code == 200 + assert response.json()[0]['id'] == voice.pk + + response = client.patch( + f'/api/v1/api/media/voices/{voice.pk}/', + {'title': 'Renamed voice', 'transcription': 'Updated transcription'}, + format='json', + ) + + assert response.status_code == 200 + voice.refresh_from_db() + assert voice.title == 'Renamed voice' + assert voice.transcription == 'Updated transcription' + + response = client.delete(f'/api/v1/api/media/voices/{voice.pk}/') + + assert response.status_code == 204 + assert not Voice.objects.filter(pk=voice.pk).exists() + + +@pytest.mark.django_db +def test_media_links_include_only_active_models(authenticated_client) -> None: + images_category = ModelCategory.objects.create(slug='images', title='Images') + active_model = NeuronModelFactory(slug='active-image-model', category=images_category) + inactive_model = NeuronModelFactory(slug='inactive-image-model', category=images_category) + inactive_model.model_settings.is_active = False + inactive_model.model_settings.save(update_fields=('is_active',)) + + response = authenticated_client.get('/api/v1/api/media/images/links/') + + assert response.status_code == 200 + assert {item['slug'] for item in response.json()} == {active_model.slug} + + +@pytest.mark.django_db +@pytest.mark.usefixtures('fake_provider_proxy') +@pytest.mark.usefixtures('local_message_storage') +def test_media_generation_accepts_multipart_input( + authenticated_client, + user, + monkeypatch, +) -> None: + case = get_model_case('fluxpulid') + model = NeuronModelFactory(slug=case.slug) + install_provider_fakes(monkeypatch, case, model) + request_data = { + 'content': 'Upscale this image', + 'info': json.dumps(case.info), + 'file': case.make_input_file(), + } + + response = authenticated_client.post( + f'/api/v1/media/image/{case.slug}', + request_data, + format='multipart', + ) + + assert response.status_code == 201, response.json() + assert Image.objects.get(user=user, model=model).output_messages.exists() + + +@pytest.mark.django_db +def test_presets_api(authenticated_client) -> None: + voice_preset = Preset.objects.create( + title='Voice preset', + slug='voice-preset', + kind='voice', + metadata={'transcription': 'Preset transcription'}, + ) + instrumental_preset = Preset.objects.create( + title='Instrumental preset', + slug='instrumental-preset', + kind='instrumental', + metadata={'duration': 10}, + ) + + response = authenticated_client.get('/api/v1/api/media/presets/') + + assert response.status_code == 200 + payload_by_uid = {item['uid']: item for item in response.json()} + assert set(payload_by_uid) == {str(voice_preset.uid), str(instrumental_preset.uid)} + assert payload_by_uid[str(voice_preset.uid)] == { + 'uid': str(voice_preset.uid), + 'title': voice_preset.title, + 'file': None, + 'metadata': voice_preset.metadata, + } + + response = authenticated_client.get('/api/v1/api/media/presets/?kind=voice') + + assert response.status_code == 200 + assert [item['uid'] for item in response.json()] == [str(voice_preset.uid)] @@ -0,0 +1,233 @@ +import base64 +from dataclasses import dataclass, field +from decimal import Decimal +from enum import StrEnum + +from django.core.files.uploadedfile import SimpleUploadedFile + + +class ResultKind(StrEnum): + TEXT = 'text' + FILE = 'file' + + +class EndpointKind(StrEnum): + CHAT = 'chat' + IMAGE = 'image' + VIDEO = 'video' + AUDIO = 'audio' + VOICE = 'voice' + + +@dataclass(frozen=True) +class InputFileFixture: + name: str + content: bytes + content_type: str + + +@dataclass(frozen=True) +class ModelCase: + slug: str + result_kind: ResultKind + endpoint: EndpointKind + file_suffix: str = '' + input_file: InputFileFixture | None = None + info: dict[str, object] = field(default_factory=dict) + + def make_input_file(self) -> SimpleUploadedFile | None: + if not self.input_file: + return None + + return SimpleUploadedFile( + self.input_file.name, + self.input_file.content, + content_type=self.input_file.content_type, + ) + + +PNG_INPUT_FILE = InputFileFixture( + name='input.png', + content=base64.b64decode( + 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII=' + ), + content_type='image/png', +) +MP3_INPUT_FILE = InputFileFixture( + name='input.mp3', + content=b'mocked audio input', + content_type='audio/mpeg', +) + + +def text_model( + slug: str, + info: dict[str, object] | None = None, + input_file: InputFileFixture | None = None, +) -> ModelCase: + return ModelCase( + slug=slug, + result_kind=ResultKind.TEXT, + endpoint=EndpointKind.CHAT, + input_file=input_file, + info=info or {}, + ) + + +def file_model( + slug: str, + file_suffix: str, + info: dict[str, object] | None = None, + input_file: InputFileFixture | None = None, + endpoint: EndpointKind | None = None, +) -> ModelCase: + endpoint_by_suffix = { + '.mp3': EndpointKind.AUDIO, + '.mp4': EndpointKind.VIDEO, + '.png': EndpointKind.IMAGE, + } + + return ModelCase( + slug=slug, + result_kind=ResultKind.FILE, + endpoint=endpoint or endpoint_by_suffix[file_suffix], + file_suffix=file_suffix, + input_file=input_file, + info=info or {}, + ) + + +TEXT_MODEL_CASES = ( + text_model('chatgpt_4', {'version': 'o3-mini'}), + text_model('chatgpt', {'version': 'gpt-5.5'}), + text_model('chatgpt_5', {'version': 'gpt-5'}), + text_model('chatgpt_5_4', {'version': 'gpt-5.4'}), + text_model('claude', {'version': 'claude-sonnet-4.6'}), + text_model('codellama'), + text_model('deepl', {'source_lang': 'ru', 'target_lang': 'en'}), + text_model('deepseek', {'version': 'deepseek/deepseek-v4-flash-0731'}), + text_model('dola_seed', {'version': 'seed-2-0-pro'}), + text_model('gemini', {'version': 'gemini-2.0-flash-001'}), + text_model('gemini_3_1', {'version': 'gemini-3.1-pro-preview'}), + text_model('gemma'), + text_model('glm_4_7'), + text_model('granite'), + text_model('grok', {'version': 'grok-4.3'}), + text_model('grok_4_1_fast'), + text_model('llama', {'version': 'llama-3.3-70b-instruct'}), + text_model('mistral'), + text_model('perplexity', {'version': 'sonar'}), + text_model('qwen', {'version': 'qwq-32b'}), + text_model('qwen_235B', {'version': 'qwen3-235b-a22b-thinking-2507'}), + text_model('qwen_3_6', {'version': 'qwen3.6-flash'}), + text_model('qwen_3_7', {'version': 'qwen3.7-max'}), + text_model('qwen_3_max_thinking'), + text_model('raifgpt'), + text_model('vicuna'), + text_model('whisper', input_file=MP3_INPUT_FILE), +) + +MEDIA_MODEL_CASES = ( + file_model('dalle', '.png'), + file_model('djourney', '.png'), + file_model('elevenlabs', '.mp3', endpoint=EndpointKind.VOICE), + file_model('elevenlabs_music', '.mp3', {'duration': 5}), + file_model('epicphotogasm', '.png'), + file_model('flux', '.png'), + file_model('flux_2', '.png', {'version': 'flux-2-pro'}), + file_model('flux_3', '.png'), + file_model('fluxkrea', '.png'), + file_model('fluxlorafast', '.png'), + file_model('fluxproultra', '.png', {'version': 'flux-dev'}), + file_model('fluxpulid', '.png', input_file=PNG_INPUT_FILE), + file_model('geminiimage', '.png'), + file_model('gptimage', '.png'), + file_model('grok_image', '.png'), + file_model('grok_image_ultra', '.png'), + file_model('grok_imagine_video', '.mp4'), + file_model('hailuo', '.mp4', {'version': 'hailuo-2.3', 'resolution': '768p'}), + file_model('hunyuan', '.mp4', {'version': 'hunyuan-video'}), + file_model('iconic', '.png'), + file_model('ideogram', '.png'), + file_model('imagen', '.png'), + file_model('kandinsky', '.png'), + file_model('kling', '.mp4', {'mode': 'standard', 'duration': 5}, PNG_INPUT_FILE), + file_model('leonardo', '.png', {'version': 'lucid-origin'}), + file_model('lightning', '.png'), + file_model('logoai', '.png'), + file_model('ltx', '.mp4', {'resolution': '1080p', 'duration': 6}), + file_model('lyria', '.mp3', {'version': 'lyria-3'}), + file_model('midjourney', '.png'), + file_model('minimaxmusic', '.mp3', {'style': 'ambient'}), + file_model('minimaxmusic_lite', '.mp3'), + file_model('minimaxvideo', '.mp4'), + file_model('musicgen', '.mp3'), + file_model('nanobanana', '.png', {'version': 'nano-banana'}), + file_model('nanobanana_2', '.png', {'resolution': '2K'}), + file_model('photon', '.png'), + file_model('pixverse', '.mp4', {'quality': '1080p', 'duration': 5}), + file_model( + 'pruna_v', + '.mp4', + {'resolution': '720p', 'generation_mode': 'standard', 'duration': 5}, + ), + file_model('prunaai', '.mp4', {'version': 'p-image'}), + file_model('pulid', '.png'), + file_model('ray', '.mp4', {'version': 'ray-2-720p', 'duration': 5}), + file_model('recraft', '.png', {'version': 'recraft-v3', 'style': 'любой'}), + file_model('reve', '.png'), + file_model('runway', '.mp4', {'duration': 5}, PNG_INPUT_FILE), + file_model('sdxlemoji', '.png'), + file_model( + 'seedance', + '.mp4', + {'version': 'seedance-2.0-fast', 'resolution': '720p', 'duration': 5}, + ), + file_model( + 'seedance_2_dreamina', + '.mp4', + {'version': 'dreamina-seedance-2-0', 'resolution': '720p', 'duration': 5}, + ), + file_model('seedream', '.png', {'version': 'seedream-boosted', 'size': '2K'}), + file_model('sora', '.mp4', {'version': 'sora-2', 'seconds': 4}), + file_model('stablediffusion', '.png', {'version': 'sd3'}), + file_model('stablemusic', '.mp3'), + file_model('suno', '.mp3', {'style': 'ambient'}), + file_model('upscaleai', '.png', input_file=PNG_INPUT_FILE), + file_model('veo', '.mp4', {'version': 'veo-3'}), + file_model('wan', '.mp4', {'resolution': '720p', 'duration': 5}, PNG_INPUT_FILE), + file_model('wan_lite', '.mp4', {'resolution': '720p'}), +) + +MODEL_CASES = TEXT_MODEL_CASES + MEDIA_MODEL_CASES + +MODEL_CASES_BY_SLUG = {case.slug: case for case in MODEL_CASES} +NON_CONTRACT_MODEL_SLUGS = frozenset( + { + 'audio_test_model', + 'image_test_model', + 'text_test_model', + 'video_test_model', + } +) + + +def get_model_case(slug: str) -> ModelCase: + try: + return MODEL_CASES_BY_SLUG[slug] + except KeyError as error: + available_slugs = ', '.join(MODEL_CASES_BY_SLUG) + + raise ValueError(f'Unknown ML model slug: {slug}. Available slugs: {available_slugs}') from error + + +def get_model_cases(result_kind: ResultKind | None = None) -> tuple[ModelCase, ...]: + if result_kind is None: + return MODEL_CASES + + return tuple(case for case in MODEL_CASES if case.result_kind == result_kind) + + +MOCK_OUTPUT = 'Mocked provider response' +MOCK_MEDIA_BYTES = b'mocked media bytes' +STARTING_BALANCE = Decimal('1000000000') @@ -0,0 +1,322 @@ +import base64 +import importlib +import inspect +from dataclasses import dataclass +from decimal import Decimal +from io import BytesIO +from types import ModuleType +from types import SimpleNamespace +from unittest.mock import Mock + +from langchain_core.messages import AIMessage + +from ml_model.adapters.models import ModelResponse +from ml_model.services.base import SimpleService +from tests.ml_models.cases import MOCK_MEDIA_BYTES, MOCK_OUTPUT, ModelCase, ResultKind + + +class SmartMediaURL(str): + def __new__(cls): + return super().__new__(cls, 'https://provider.test/generated-file') + + def __iter__(self): + yield self + + @property + def url(self): + return self + + def __getitem__(self, item): + if item == 0: + return self + + return super().__getitem__(item) + + +class SmartStatus: + SUCCESS_VALUES = {'COMPLETED', 'completed', 'succeeded'} + + def __eq__(self, other) -> bool: + return other in self.SUCCESS_VALUES + + def __hash__(self) -> int: + return hash('completed') + + +class SmartResponse: + status_code = 200 + content = MOCK_MEDIA_BYTES + text = MOCK_OUTPUT + + def json(self) -> dict: + encoded_media = base64.b64encode(MOCK_MEDIA_BYTES).decode() + + return { + 'id': 'mock-generation', + 'status': SmartStatus(), + 'logs': '', + 'input_tokens': 100, + 'output': [SmartMediaURL()], + 'response_url': SmartMediaURL(), + 'status_url': SmartMediaURL(), + 'urls': {'get': SmartMediaURL()}, + 'images': [{'url': SmartMediaURL()}], + 'data': [{'b64_json': encoded_media}], + 'artifacts': [{'base64': encoded_media}], + 'choices': [{'message': {'content': MOCK_OUTPUT, 'annotations': []}}], + 'usage': { + 'prompt_tokens': 100, + 'completion_tokens': 40, + 'input_tokens': 100, + 'output_tokens': 40, + 'total_tokens': 140, + 'input_tokens_details': {'image_tokens': 0, 'text_tokens': 100}, + }, + } + + def raise_for_status(self) -> None: + return None + + +class SmartClient: + def __init__(self, *args, **kwargs) -> None: + self.response = SmartResponse() + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + return None + + def post(self, *args, **kwargs) -> SmartResponse: + return self.response + + def get(self, *args, **kwargs) -> SmartResponse: + return self.response + + def close(self) -> None: + return None + + +class FakeLLM: + model_name = 'gpt-4o' + + def __init__(self, *args, **kwargs) -> None: + self.model_name = kwargs.get('model', self.model_name) + + +class FakeConversation: + def __init__(self, *args, **kwargs) -> None: + pass + + def invoke(self, *args, **kwargs) -> AIMessage: + return AIMessage(content=MOCK_OUTPUT) + + +class FakeAudioSegment: + @classmethod + def empty(cls): + return cls() + + @classmethod + def from_file(cls, *args, **kwargs): + return cls() + + def __iadd__(self, other): + return self + + def export(self, output: BytesIO, format: str) -> None: + output.write(MOCK_MEDIA_BYTES) + + +@dataclass +class ProviderMockTracker: + call_count: int = 0 + + def record(self) -> None: + self.call_count += 1 + + +class SuccessfulTask: + def __init__(self, result) -> None: + self.result = result + + def get(self): + return self.result + + def successful(self) -> bool: + return True + + +def install_provider_fakes(monkeypatch, case: ModelCase, model) -> ProviderMockTracker: + tracker = ProviderMockTracker() + media_url = SmartMediaURL() + + monkeypatch.setattr(SimpleService, 'neuron_model', property(lambda service: model)) + monkeypatch.setattr(SimpleService, 'translate_prompt', lambda service, prompt, to='en': prompt) + service_class = model.service + if not hasattr(service_class, 'price'): + monkeypatch.setattr(service_class, 'price', Decimal('1'), raising=False) + + def provider_result(*args, **kwargs): + tracker.record() + if case.result_kind == ResultKind.TEXT: + return MOCK_OUTPUT + + return media_url + + def token_provider_result(*args, **kwargs): + tracker.record() + + return MOCK_OUTPUT, 100, 40 + + def bytedance_result(*args, **kwargs): + tracker.record() + content_type = str(kwargs.get('content_type', '')).lower() + if 'chat' in content_type: + return MOCK_OUTPUT, 100, 40 + if 'video' in content_type: + return media_url, 40 + + return media_url + + def streaming_result(*args, **kwargs) -> ModelResponse: + tracker.record() + + return ModelResponse(MOCK_OUTPUT, 100, 40, 0.001) + + service_modules = _service_modules(service_class) + + for module in service_modules: + for name, replacement in ( + ('replicate_run', provider_result), + ('upscale_run', lambda *args, **kwargs: [media_url]), + ('openrouter_run', token_provider_result), + ('bytedance_model_ark_run', bytedance_result), + ): + if hasattr(module, name): + monkeypatch.setattr(module, name, replacement) + + if hasattr(module, 'requests'): + monkeypatch.setattr(module.requests, 'get', Mock(return_value=SmartResponse())) + monkeypatch.setattr(module.requests, 'post', Mock(return_value=SmartResponse())) + if hasattr(module, 'httpx'): + monkeypatch.setattr(module.httpx, 'Client', SmartClient) + monkeypatch.setattr(module.httpx, 'get', Mock(return_value=SmartResponse())) + monkeypatch.setattr(module.httpx, 'post', Mock(return_value=SmartResponse())) + if hasattr(module, 'client'): + monkeypatch.setattr(module, 'client', SmartClient()) + + _install_adapter_fakes(monkeypatch, service_class, service_modules, streaming_result) + _install_task_fakes(monkeypatch, case.slug) + _install_model_specific_fakes(monkeypatch, case.slug, tracker, media_url) + + return tracker + + +def _service_modules(service_class) -> tuple[ModuleType, ...]: + modules = [] + for base_class in service_class.__mro__: + module = inspect.getmodule(base_class) + if module and module.__name__.startswith('ml_model.services.') and module not in modules: + modules.append(module) + + return tuple(modules) + + +def _install_adapter_fakes(monkeypatch, service_class, service_modules, streaming_result) -> None: + if any(hasattr(module, 'OpenrouterAdapter') for module in service_modules): + from ml_model.adapters.openrouter import OpenrouterAdapter + + monkeypatch.setattr(OpenrouterAdapter, 'collect_streaming_api', streaming_result) + + if any(hasattr(module, 'BytedanceModelArkAdapter') for module in service_modules): + from ml_model.adapters.bytedance_model_ark import BytedanceModelArkAdapter + + monkeypatch.setattr(BytedanceModelArkAdapter, 'batch_tokenize', lambda model, texts: 100) + + from ml_model.services.chatgpt import Chatgpt + from ml_model.services.chatgpt_4 import Chatgpt_4 + + if issubclass(service_class, Chatgpt_4): + monkeypatch.setattr( + Chatgpt_4, + 'call_openai_api', + lambda service, proxy, endpoint, json_data: (100, 40, AIMessage(content=MOCK_OUTPUT)), + ) + monkeypatch.setattr(Chatgpt_4, 'count_text_tokens', lambda service, messages: 100) + + if issubclass(service_class, Chatgpt): + monkeypatch.setattr(Chatgpt, '_count_responses_input_tokens', lambda service, proxy, payload: 100) + monkeypatch.setattr( + Chatgpt, + '_stream_openai_responses', + lambda service, proxy, json_data, model_name: (100, 40, AIMessage(content=MOCK_OUTPUT)), + ) + + +def _install_task_fakes(monkeypatch, slug: str) -> None: + if slug == 'deepl': + module = importlib.import_module('ml_model.services.deepl') + monkeypatch.setattr( + module.translate, + 'delay', + lambda callback_data: SuccessfulTask(MOCK_OUTPUT), + ) + + if slug == 'whisper': + module = importlib.import_module('ml_model.services.whisper') + monkeypatch.setattr( + module.transcript_audio, + 'delay', + lambda file: SuccessfulTask({'text': MOCK_OUTPUT}), + ) + audio_info = SimpleNamespace(info=SimpleNamespace(length=1)) + monkeypatch.setattr(module, 'MP3', lambda audio: audio_info) + monkeypatch.setattr(module, 'WAVE', lambda audio: audio_info) + + +def _install_model_specific_fakes( + monkeypatch, + slug: str, + tracker: ProviderMockTracker, + media_url: SmartMediaURL, +) -> None: + if slug == 'minimaxmusic_lite': + module = importlib.import_module('ml_model.services.minimaxmusic_lite') + monkeypatch.setattr( + module.Preset.objects, + 'get', + lambda **kwargs: SimpleNamespace(file=SimpleNamespace(url=media_url)), + ) + + if slug == 'granite': + module = importlib.import_module('ml_model.services.granite') + monkeypatch.setattr( + module.Granite, + '_call_api', + lambda service, payload: (tracker.record() or {'output': [MOCK_OUTPUT]}), + ) + + if slug in {'chatgpt_4', 'raifgpt'}: + module = importlib.import_module(f'ml_model.services.{slug}') + monkeypatch.setattr(module, 'ChatOpenAI', FakeLLM) + if hasattr(module, 'RunnableWithMessageHistory'): + monkeypatch.setattr(module, 'RunnableWithMessageHistory', FakeConversation) + + if slug == 'elevenlabs': + module = importlib.import_module('ml_model.services.elevenlabs') + monkeypatch.setattr(module, 'AudioSegment', FakeAudioSegment) + monkeypatch.setattr( + module.FileProcessingService, + 'get_voice_file', + classmethod(lambda cls, voice_id, preset_id, user: SimpleNamespace(url=media_url)), + ) + + if slug == 'gptimage': + module = importlib.import_module('ml_model.services.gptimage') + monkeypatch.setattr( + module.Gptimage, + 'count_predict_tokens', + classmethod(lambda cls, text, width, height, size, quality: (100, 0, 196)), + ) @@ -0,0 +1,88 @@ +import ast +import json +from pathlib import Path + +import pytest + +from messages.models import Message +from payments.models import Invoice +from tests.factories import ChatFactory, NeuronModelFactory +from tests.ml_models.cases import ( + NON_CONTRACT_MODEL_SLUGS, + ResultKind, + EndpointKind, + get_model_cases, +) +from tests.ml_models.provider_fakes import install_provider_fakes + + +MODEL_CASES = get_model_cases() +CHAT_MODEL_PARAMS = tuple( + pytest.param(case, id=case.slug, marks=pytest.mark.ml_model(case.slug)) + for case in MODEL_CASES + if case.endpoint == EndpointKind.CHAT +) + + +@pytest.mark.django_db +@pytest.mark.usefixtures('fake_provider_proxy') +@pytest.mark.parametrize('case', CHAT_MODEL_PARAMS) +def test_model_api_contract( + case, + authenticated_client, + user, + local_message_storage, + monkeypatch, +) -> None: + model = NeuronModelFactory(slug=case.slug) + chat = ChatFactory(user=user, model=model) + install_provider_fakes(monkeypatch, case, model) + request_data = { + 'content': 'Generic ML model API contract', + 'info': dict(case.info), + } + if input_file := case.make_input_file(): + request_data['file'] = input_file + request_data['info'] = json.dumps(request_data['info']) + + balance_before = user.payment_plan.current_token_balance + response = authenticated_client.post( + f'/api/v1/chats/{chat.uid}/messages/', + request_data, + format='multipart' if input_file else 'json', + ) + + assert response.status_code == 201, response.json() + payload = response.json() + assert len(payload) >= 2 + assert payload[0]['from_model'] is False + assert all(message['from_model'] is True for message in payload[1:]) + + output_messages = Message.objects.filter(object_id=chat.uid, from_model=True) + assert output_messages.exists() + if case.result_kind == ResultKind.TEXT: + assert any(message.content for message in output_messages) + else: + assert all(message.file for message in output_messages) + for output_message in output_messages: + assert local_message_storage.exists(output_message.file.name) + with local_message_storage.open(output_message.file.name, 'rb') as saved_file: + assert saved_file.read() + + invoice = Invoice.objects.get(user=user, model=model) + user.payment_plan.refresh_from_db() + charged_tokens = balance_before - user.payment_plan.current_token_balance + assert charged_tokens == invoice.cost + + +def test_every_exported_model_has_api_case() -> None: + services_init = Path('ml_model/services/__init__.py') + module = ast.parse(services_init.read_text()) + exported_slugs = { + node.module.rsplit('.', 1)[1] + for node in module.body + if isinstance(node, ast.ImportFrom) and node.module and node.module.startswith('ml_model.services.') + } + case_slugs = {case.slug for case in MODEL_CASES} + + assert case_slugs == exported_slugs - NON_CONTRACT_MODEL_SLUGS @@ -0,0 +1,85 @@ +import json +import os + +import pytest + +from messages.models import Message +from payments.models import Invoice +from poller.models import Proxy +from tests.factories import ChatFactory, NeuronModelFactory +from tests.ml_models.cases import ModelCase, ResultKind, get_model_case + + +def _configure_provider_proxy() -> None: + proxy_address = os.getenv('PROVIDER_SMOKE_PROXY_ADDRESS', '').strip() + if not proxy_address: + return + + proxy_protocol = os.getenv('PROVIDER_SMOKE_PROXY_PROTOCOL', 'http').strip().lower() + allowed_protocols = {choice[0] for choice in Proxy.ProtocolChoices.choices} + if proxy_protocol not in allowed_protocols: + raise pytest.UsageError( + 'PROVIDER_SMOKE_PROXY_PROTOCOL must be http, https, or socks.' + ) + if '://' in proxy_address: + raise pytest.UsageError('PROVIDER_SMOKE_PROXY_ADDRESS must not contain a protocol.') + + Proxy.objects.create(address=proxy_address, protocol=proxy_protocol) + + +def pytest_generate_tests(metafunc: pytest.Metafunc) -> None: + selected_slugs = metafunc.config.getoption('--model-slug') + if len(selected_slugs) != 1: + raise pytest.UsageError('--provider-smoke requires exactly one --model-slug.') + + try: + case = get_model_case(selected_slugs[0]) + except ValueError as error: + raise pytest.UsageError(str(error)) from error + + parameter = pytest.param(case, id=case.slug, marks=pytest.mark.ml_model(case.slug)) + metafunc.parametrize('case', (parameter,)) + + +@pytest.mark.provider_smoke +@pytest.mark.django_db +def test_real_provider_api_flow( + case: ModelCase, + authenticated_client, + user, +) -> None: + _configure_provider_proxy() + model = NeuronModelFactory(slug=case.slug) + chat = ChatFactory(user=user, model=model) + request_data = { + 'content': 'Short provider smoke test. Reply or generate a minimal result.', + 'info': dict(case.info), + } + if input_file := case.make_input_file(): + request_data['file'] = input_file + request_data['info'] = json.dumps(request_data['info']) + + balance_before = user.payment_plan.current_token_balance + response = authenticated_client.post( + f'/api/v1/chats/{chat.uid}/messages/', + request_data, + format='multipart' if input_file else 'json', + ) + + response_payload = response.json() + + assert response.status_code == 201, response_payload + output_messages = Message.objects.filter(object_id=chat.uid, from_model=True) + assert output_messages.exists() + if case.result_kind == ResultKind.TEXT: + assert any(message.content for message in output_messages) + else: + assert all(message.file for message in output_messages) + assert all( + message.file.storage.exists(message.file.name) for message in output_messages + ) + + invoice = Invoice.objects.get(user=user, model=model) + user.payment_plan.refresh_from_db() + charged_tokens = balance_before - user.payment_plan.current_token_balance + assert charged_tokens == invoice.cost @@ -0,0 +1,85 @@ +import pytest +from django.utils.translation import gettext as _ + +from ml_model.exceptions import FileNotProvided, InvalidParameterError +from ml_model.models import ModelInput, ModelVersion +from ml_model.validators import ModelInputValidator +from tests.factories import NeuronModelFactory + + +@pytest.mark.django_db +@pytest.mark.parametrize('content', [None, '', ' ', '\n\t']) +def test_model_input_validator_rejects_empty_text(content) -> None: + model = NeuronModelFactory() + + with pytest.raises(InvalidParameterError) as error: + ModelInputValidator(model, content=content, file=None).validate() + + assert error.value.error_text == _('The request must not be empty') + assert str(error.value) == str(_('The request must not be empty')) + + +@pytest.mark.django_db +def test_model_input_validator_requires_file() -> None: + model = NeuronModelFactory() + model.inputs.all().delete() + ModelInput.objects.create(model=model, type=ModelInput.TypeChoices.IMAGE, required=True) + + with pytest.raises(FileNotProvided) as error: + ModelInputValidator(model, content='Describe this image', file=None).validate() + + assert error.value.file_type == ModelInput.TypeChoices.IMAGE.label + assert str(error.value) == str( + _('The %(file_type)s is not attached') % {'file_type': _(ModelInput.TypeChoices.IMAGE.label).lower()} + ) + + +@pytest.mark.django_db +def test_model_input_validator_accepts_required_inputs() -> None: + text_model = NeuronModelFactory() + ModelInputValidator(text_model, content='A valid prompt', file=None).validate() + + file_model = NeuronModelFactory() + file_model.inputs.all().delete() + ModelInput.objects.create(model=file_model, type=ModelInput.TypeChoices.IMAGE, required=True) + + ModelInputValidator(file_model, content=None, file=object()).validate() + + +@pytest.mark.django_db +def test_model_input_validator_ignores_optional_file() -> None: + model = NeuronModelFactory() + ModelInput.objects.create(model=model, type=ModelInput.TypeChoices.IMAGE, required=False) + + ModelInputValidator(model, content='A valid prompt', file=None).validate() + + +@pytest.mark.django_db +def test_model_input_validator_uses_inputs_for_requested_version() -> None: + model = NeuronModelFactory() + first_version = ModelVersion.objects.create(model=model, name='First', slug='first') + second_version = ModelVersion.objects.create(model=model, name='Second', slug='second') + model.inputs.all().delete() + text_input = ModelInput.objects.create( + model=model, + type=ModelInput.TypeChoices.TEXT, + required=True, + ) + text_input.versions.add(first_version) + file_input = ModelInput.objects.create( + model=model, + type=ModelInput.TypeChoices.IMAGE, + required=True, + ) + file_input.versions.add(second_version) + + ModelInputValidator( + model, content='A valid prompt', file=None, info={'version': first_version.slug} + ).validate() + + with pytest.raises(FileNotProvided) as error: + ModelInputValidator( + model, content='A valid prompt', file=None, info={'version': second_version.slug} + ).validate() + + assert error.value.file_type == ModelInput.TypeChoices.IMAGE.label @@ -0,0 +1,45 @@ +import pytest +from django.core.files.storage import FileSystemStorage +from rest_framework.test import APIClient +from rest_framework_simplejwt.tokens import RefreshToken + +from messages.models import Message +from poller.models import Proxy +from tests.factories import UserFactory +from tests.ml_models.cases import STARTING_BALANCE + + +@pytest.fixture +def api_client() -> APIClient: + return APIClient() + + +@pytest.fixture +def user(db): + user = UserFactory() + user.payment_plan.current_token_balance = STARTING_BALANCE + user.payment_plan.save() + + return user + + +@pytest.fixture +def authenticated_client(api_client: APIClient, user) -> APIClient: + access_token = RefreshToken.for_user(user).access_token + api_client.credentials(HTTP_AUTHORIZATION=f'Bearer {access_token}') + + return api_client + + +@pytest.fixture +def local_message_storage(tmp_path, monkeypatch) -> FileSystemStorage: + storage = FileSystemStorage(location=tmp_path) + file_field = Message._meta.get_field('file') + monkeypatch.setattr(file_field, 'storage', storage) + + return storage + + +@pytest.fixture +def fake_provider_proxy() -> Proxy: + return Proxy.objects.create(address='proxy.test:8080', protocol=Proxy.ProtocolChoices.HTTP) @@ -0,0 +1,34 @@ +# Тесты ML-моделей + +Основные тесты проходят полный API-флоу чата: проверяют сообщения, файлы, счета и +списание баланса. Внешние провайдеры в них замоканы. Для каждой экспортированной +модели должен быть кейс в `tests/ml_models/cases.py`. + +```bash +# Все тесты +python manage.py runtests + +# Одна или несколько моделей +python manage.py runtests --model-slug gemma --model-slug qwen_3_6 + +# С замерами времени, CPU и памяти +python manage.py runtests --profile-resources +``` + +GitLab CI сохраняет JUnit-отчёт в `test-results/junit.xml` и показывает статистику +тестов в pipeline и merge request. + +## Проверка реального провайдера + +Ручной smoke-тест запускает ровно одну модель без моков и печатает полученный ответ: + +```bash +python manage.py runtests --provider-smoke --model-slug grok +``` + +Он использует ключи провайдера из окружения. Для моделей, которым нужен прокси, +задайте `PROVIDER_SMOKE_PROXY_ADDRESS` в формате поля `Proxy.address` (без протокола) +и `PROVIDER_SMOKE_PROXY_PROTOCOL`: `http`, `https` или `socks`. + +Этот тест может расходовать реальные токены. Обычный `runtests` и GitLab CI его не +собирают. @@ -0,0 +1,124 @@ +import resource +import time +from collections.abc import Iterator + +import pytest + +pytest_plugins = ('tests.support.fixtures',) + +RESOURCE_PROFILES: list[dict[str, str | float | int]] = [] +PROVIDER_SMOKE_TEST_FILE = 'test_provider_smoke.py' + + +def pytest_ignore_collect(collection_path, config) -> bool: + is_provider_smoke_test = collection_path.name == PROVIDER_SMOKE_TEST_FILE + is_manual_provider_smoke_run = config.getoption('--provider-smoke') + + return is_provider_smoke_test and not is_manual_provider_smoke_run + + +def pytest_addoption(parser) -> None: + group = parser.getgroup('AIR ML model contracts') + group.addoption( + '--model-slug', + action='append', + default=[], + help='Run ML model API contracts only for the selected slug.', + ) + group.addoption( + '--profile-resources', + action='store_true', + help='Report wall time, CPU time, and peak RSS for each test.', + ) + group.addoption( + '--provider-smoke', + action='store_true', + help='Run one explicitly selected ML model against its real provider.', + ) + + +def pytest_collection_modifyitems(config, items) -> None: + selected_slugs = set(config.getoption('--model-slug')) + provider_smoke = config.getoption('--provider-smoke') + + available_slugs = { + marker.args[0] + for item in items + if (marker := item.get_closest_marker('ml_model')) and marker.args + } + if unknown_slugs := selected_slugs - available_slugs: + available = ', '.join(sorted(available_slugs)) + unknown = ', '.join(sorted(unknown_slugs)) + + raise pytest.UsageError(f'Unknown ML model slug: {unknown}. Available slugs: {available}') + + if provider_smoke and len(selected_slugs) != 1: + raise pytest.UsageError('--provider-smoke requires exactly one --model-slug.') + + deselected_items = [] + selected_items = [] + for item in items: + is_provider_smoke = item.get_closest_marker('provider_smoke') is not None + marker = item.get_closest_marker('ml_model') + slug = marker.args[0] if marker and marker.args else None + + if provider_smoke: + is_selected = is_provider_smoke and slug in selected_slugs + else: + is_selected = not is_provider_smoke + + if is_selected: + selected_items.append(item) + else: + deselected_items.append(item) + + if deselected_items: + config.hook.pytest_deselected(items=deselected_items) + items[:] = selected_items + + if provider_smoke: + return + + for item in items: + marker = item.get_closest_marker('ml_model') + slug = marker.args[0] if marker and marker.args else None + if selected_slugs and slug not in selected_slugs: + item.add_marker(pytest.mark.skip(reason='ML model slug was not selected')) + + +@pytest.fixture(autouse=True) +def resource_profile(request) -> Iterator[None]: + if not request.config.getoption('--profile-resources'): + yield + + return + + usage_before = resource.getrusage(resource.RUSAGE_SELF) + wall_started = time.perf_counter() + cpu_started = time.process_time() + + yield + + usage_after = resource.getrusage(resource.RUSAGE_SELF) + peak_rss_mb = usage_after.ru_maxrss / 1024 + RESOURCE_PROFILES.append( + { + 'test': request.node.nodeid, + 'wall_seconds': round(time.perf_counter() - wall_started, 6), + 'cpu_seconds': round(time.process_time() - cpu_started, 6), + 'peak_rss_mb': round(peak_rss_mb, 3), + 'minor_page_faults': usage_after.ru_minflt - usage_before.ru_minflt, + } + ) + + +def pytest_terminal_summary(terminalreporter, config) -> None: + if not config.getoption('--profile-resources') or not RESOURCE_PROFILES: + return + + terminalreporter.section('resource profile') + for profile in RESOURCE_PROFILES: + terminalreporter.write_line( + '{test}: wall={wall_seconds:.3f}s cpu={cpu_seconds:.3f}s ' + 'peak_rss={peak_rss_mb:.1f}MB minor_faults={minor_page_faults}'.format(**profile) + ) @@ -0,0 +1,84 @@ +from decimal import Decimal + +import factory +from factory.django import DjangoModelFactory + +from authentication.models import CustomUserModel +from ml_model.models import ModelCategory, ModelInput, ModelSettings, NeuronModel +from payments.models import PaymentPlan +from tools.chats.models import Chat + + +class PaymentPlanFactory(DjangoModelFactory): + class Meta: + model = PaymentPlan + django_get_or_create = ('price',) + + price = Decimal('1') + tokens_per_plan = Decimal('10000') + + +class FreePaymentPlanFactory(DjangoModelFactory): + class Meta: + model = PaymentPlan + django_get_or_create = ('price', 'is_corporate') + + price = Decimal('0') + tokens_per_plan = Decimal('10000') + is_corporate = False + + +class UserFactory(DjangoModelFactory): + class Meta: + model = CustomUserModel + skip_postgeneration_save = True + + email = factory.Sequence(lambda number: f'ml-model-{number}@test.local') + password = factory.PostGenerationMethodCall('set_password', 'test-password') + + @classmethod + def _create(cls, model_class, *args, **kwargs): + _ = FreePaymentPlanFactory() + paid_plan = PaymentPlanFactory() + + user = model_class.objects.create_user(*args, **kwargs) + user.payment_plan.plan = paid_plan + user.payment_plan.current_token_balance = paid_plan.tokens_per_plan + user.payment_plan.save(update_fields=('plan', 'current_token_balance')) + + return user + + +class ModelCategoryFactory(DjangoModelFactory): + class Meta: + model = ModelCategory + django_get_or_create = ('slug',) + + title = 'Chat-bots' + slug = 'chat-bots' + + +class NeuronModelFactory(DjangoModelFactory): + class Meta: + model = NeuronModel + skip_postgeneration_save = True + + title = factory.LazyAttribute(lambda model: model.slug.replace('_', ' ').title()) + slug = factory.Sequence(lambda number: f'test_model_{number}') + category = factory.SubFactory(ModelCategoryFactory) + + @factory.post_generation + def configure(model, create, extracted, **kwargs): + if not create: + return + ModelSettings.objects.create(model=model, is_active=True) + ModelInput.objects.create(model=model, type=ModelInput.TypeChoices.TEXT, required=True) + + +class ChatFactory(DjangoModelFactory): + class Meta: + model = Chat + + title = 'ML model API contract' + user = factory.SubFactory(UserFactory) + model = factory.SubFactory(NeuronModelFactory) @@ -8,7 +8,7 @@ from dataclasses import dataclass from unittest.mock import patch import orjson -from cacheops import invalidate_all +from cachalot.api import invalidate from django.db import close_old_connections, connections from django.test import Client, TransactionTestCase from rest_framework_simplejwt.tokens import RefreshToken @@ -75,7 +75,7 @@ class SSEStreamLoadTest(TransactionTestCase): scenario: LoadScenario def setUp(self) -> None: - invalidate_all() + invalidate() self.user = CustomUserModel.objects.create_user(email='sse-load@test.test', password='test') self.access_token = str(RefreshToken.for_user(self.user).access_token) category = ModelCategory.objects.create(title='Chat-bots', slug='chat-bots') @@ -3,7 +3,7 @@ import time from unittest.mock import patch import orjson -from cacheops import invalidate_all +from cachalot.api import invalidate from django.test import Client, TestCase from rest_framework_simplejwt.tokens import RefreshToken @@ -37,7 +37,7 @@ class SSEStreamAPITest(TestCase): ) def setUp(self) -> None: - invalidate_all() + invalidate() self.client = Client() self.user = CustomUserModel.objects.create_user(email='sse-test@test.test', password='test') self.access_token = str(RefreshToken.for_user(self.user).access_token) @@ -2,7 +2,6 @@ from datetime import date from decimal import Decimal from django.contrib.admin.models import ADDITION, DELETION, LogEntry -from django.contrib.contenttypes.models import ContentType from django.db import IntegrityError from django.utils.translation import gettext as _ @@ -23,13 +22,11 @@ class APIKeyService(BaseService): if 'unique' in str(exc): raise DuplicateError(model=APIKey, attrs=(_('Name'), _('Owner'))) raise UnknownError from exc - LogEntry.objects.log_action( - self.user.pk, - ContentType.objects.get_for_model(api_key).pk, - api_key.pk, - str(api_key), - ADDITION, - [ + LogEntry.objects.log_actions( + user_id=self.user.pk, + queryset=[api_key], + action_flag=ADDITION, + change_message=[ { 'added': { 'name': 'API-ключ', @@ -37,6 +34,7 @@ class APIKeyService(BaseService): } } ], + single_object=True, ) if serialize: return APIKeyResultSerializer(api_key) @@ -65,13 +63,11 @@ class APIKeyService(BaseService): def delete(self, key_name): api_key = APIKeySelector(self.user).get_by_name(name=key_name) - LogEntry.objects.log_action( - self.user.pk, - ContentType.objects.get_for_model(api_key).pk, - api_key.pk, - str(api_key), - DELETION, - [ + LogEntry.objects.log_actions( + user_id=self.user.pk, + queryset=[api_key], + action_flag=DELETION, + change_message=[ { 'deleted': { 'name': 'API-ключ', @@ -79,6 +75,7 @@ class APIKeyService(BaseService): } } ], + single_object=True, ) api_key.is_deleted = True api_key.save() @@ -0,0 +1,24 @@ +SECRET_KEY=ci-only-secret-key-that-is-long-enough +TELEGRAM_BOT_TOKEN=ci-only-token +DOMAIN=localhost +DJANGO_SUPERUSER_USERNAME=ci-admin +DJANGO_SUPERUSER_EMAIL=ci-admin@example.test +DJANGO_SUPERUSER_PASSWORD=ci-only-password + +POSTGRES_DB=backend_ci +POSTGRES_USER=backend_ci +POSTGRES_PASSWORD=backend_ci +POSTGRES_HOST=db +POSTGRES_PORT=5432 + +MINIO_ENDPOINT=s3:9000 +MINIO_ACCESS_KEY=ci-only-access-key +MINIO_SECRET_KEY=ci-only-secret-key + +REDIS_HOST=cache-mdb +CACHE_BROKER_URL=redis://cache-mdb:6379/0 +CELERY_BROKER_URL=redis://cache-mdb:6379/1 +CELERY_RESULT_BACKEND=redis://cache-mdb:6379/2 + +RELEASE=ci +ENVIRONMENT=test @@ -119,4 +119,8 @@ PYTHONWARNINGS=ignore::UserWarning:polymorphic # temporarily # SSE STREAMING FF__STREAMING_ENABLED=True -DATA_UPLOAD_MAX_MEMORY_SIZE=5 # MB \ No newline at end of file +DATA_UPLOAD_MAX_MEMORY_SIZE=5 # MB + +# MANUAL PROVIDER SMOKE TEST +PROVIDER_SMOKE_PROXY_ADDRESS= +PROVIDER_SMOKE_PROXY_PROTOCOL=http @@ -10,6 +10,7 @@ venv/ .venv/ virtualenv/ air_reports/ +test-results/ .python-version **/locales/**/*.mo @@ -1,5 +1,6 @@ stages: - Build + - Test - Deploy default: @@ -19,6 +20,33 @@ build: - staging when: on_success +test: + stage: Test + variables: + IMAGE_TAG: $CI_REGISTRY_IMAGE:$CI_COMMIT_SHA + COMPOSE_FILE: docker-compose.yml:docker-compose.local.yml + COMPOSE_PROJECT_NAME: $CI_PROJECT_NAME-test-$CI_PIPELINE_ID + script: + - cp .env.ci .env + - docker pull $IMAGE_TAG + - docker compose up -d --wait db cache-mdb + - until docker compose exec -T db pg_isready --username=backend_ci --dbname=backend_ci; do sleep 1; done + - mkdir -p test-results + - docker compose run --rm app python manage.py runtests --profile-resources --pytest-arg=--junitxml=/app/test-results/junit.xml + after_script: + - docker compose down --volumes --remove-orphans + artifacts: + when: always + expire_in: 1 week + reports: + junit: test-results/junit.xml + paths: + - test-results/junit.xml + only: + - main + - staging + when: on_success + .deploy_template: &default_deploy_job stage: Deploy services: @@ -66,4 +94,4 @@ deploy_production: deployment_tier: production url: $DOMAIN only: - - main \ No newline at end of file + - main @@ -21,14 +21,12 @@ RUN --mount=target=/var/lib/apt/lists,type=cache,sharing=locked \ && apt-get update \ && apt-get -y --no-install-recommends install gettext antiword ffmpeg git +ARG UV_SYNC_ARGS= + RUN --mount=from=ghcr.io/astral-sh/uv,source=/uv,target=/bin/uv \ --mount=type=cache,target=/root/.cache/uv \ --mount=type=bind,source=uv.lock,target=uv.lock \ --mount=type=bind,source=pyproject.toml,target=pyproject.toml \ - uv sync --no-install-project + uv sync --no-install-project ${UV_SYNC_ARGS} COPY . . - -RUN --mount=from=ghcr.io/astral-sh/uv,source=/uv,target=/bin/uv \ - --mount=type=cache,target=/root/.cache/uv \ - uv sync @@ -15,6 +15,11 @@ services: python manage.py runserver 0.0.0.0:8000 ports: - "8000:8000" + build: + context: . + dockerfile: Dockerfile + args: + UV_SYNC_ARGS: --all-groups celery: <<: *dev-app-config @@ -8,14 +8,14 @@ dependencies = [ "channels[daphne]==4.2.0", "deepl==1.21.1", "dj-rest-auth==4.0.1", - "django==5.0.*", - "django-cacheops==7.0.2", + "django==6.0.*", + "django-cachalot==2.9.0", "django-celery-beat>=2.9.0", "django-cors-headers==4.2.0", "django-filter==23.2", "django-import-export==4.0.9", - "django-minio-backend", - "django-ninja==1.3.0", + "django-minio-backend==4.5.0", + "django-ninja==1.6.2", "django-oauth-toolkit==2.3.0", "django-ordered-model==3.7.4", "django-polymorphic==3.1.0", @@ -38,6 +38,7 @@ dependencies = [ "langchain-openai==0.3.6", "langchainhub==0.1.15", "langserve[client]==0.0.46", + "markupsafe>=3.0.2", "minio>=7.0,<=8.0", "mutagen==1.47.0", "openpyxl==3.1.2", @@ -55,6 +56,7 @@ dependencies = [ "sentry-sdk[django]==2.39.0", "setuptools<81", "social-auth-app-django==5.3.0", + "social-auth-core==4.9.1", "tiktoken==0.9.0", "unleashclient==6.4.0", "yookassa==3.10.1", @@ -65,6 +67,9 @@ debug = [ "debugpy>=1.8.20", ] dev = [ + "factory-boy>=3.3.3", + "pytest>=9.0.2", + "pytest-django>=4.11.1", "ruff>=0.15.15", ] @@ -125,8 +130,15 @@ line-ending = "lf" docstring-code-format = false docstring-code-line-length = "dynamic" -[tool.uv.sources] -django-minio-backend = { git = "https://github.com/theriverman/django-minio-backend", tag = "3.7.0" } +[tool.pytest.ini_options] +DJANGO_SETTINGS_MODULE = "backend.settings" +python_files = ["test_*.py"] +testpaths = ["tests"] +addopts = "-ra" +markers = [ + "ml_model(slug): API contract for a concrete ML model service", + "provider_smoke: opt-in API test that calls a real external provider", +] [tool.ruff.lint.per-file-ignores] "__init__.py" = ["E402", "F401"]