import base64 from cryptography.fernet import Fernet from django.contrib.admin.sites import AdminSite from django.test import RequestFactory, SimpleTestCase, TestCase, override_settings from apps.ai.admin import AiModelAdmin from apps.ai.aliases import AliasNotFoundError, ModelCapabilityError, resolve_alias from apps.ai.importers import import_ai_models_config from apps.ai.models import AiModel, ModelAlias from apps.ai.providers import AiCapabilityError, ResolvedModel, get_provider, resolve_api_type from apps.ai.providers.openai_compatible import ChatCompletionsProvider, ImagesEditsProvider TEST_ENCRYPTION_KEY = Fernet.generate_key().decode("ascii") class FakeResponse: def __init__(self, payload, content=b""): self.payload = payload self.content = content def json(self): return self.payload def raise_for_status(self): return None class FakeSession: def __init__(self, *responses): self.responses = list(responses) self.posts = [] self.gets = [] self.trust_env = True def post(self, url, **kwargs): self.posts.append({"url": url, **kwargs}) return self.responses.pop(0) def get(self, url, **kwargs): self.gets.append({"url": url, **kwargs}) return self.responses.pop(0) class ProviderRegistryTests(SimpleTestCase): def test_auto_api_type_resolves_chat_provider_from_url(self): url = "https://api.vectorengine.ai/v1/chat/completions" self.assertEqual(resolve_api_type("auto", url), "chat") self.assertIsInstance(get_provider("auto", url), ChatCompletionsProvider) class ChatCompletionsProviderTests(SimpleTestCase): def test_generate_text_builds_chat_payload_and_cleans_titles(self): session = FakeSession( FakeResponse( { "choices": [ {"message": {"content": "1. Red Dress\n2. Blue Coat"}} ] } ) ) provider = ChatCompletionsProvider(session=session) model = ResolvedModel( name="GPT-5.5 text", url="https://api.vectorengine.ai/v1", model="gpt-5.5", api_key="test-key", api_type="chat", ) result = provider.generate_text( "Generate titles", model, parameters={"temperature": 0.2}, ) self.assertEqual(result.text, "Red Dress") self.assertEqual(result.titles, ("Red Dress", "Blue Coat")) request = session.posts[0] self.assertEqual( request["url"], "https://api.vectorengine.ai/v1/chat/completions", ) self.assertEqual(request["headers"]["Authorization"], "Bearer test-key") self.assertEqual(request["json"]["model"], "gpt-5.5") self.assertFalse(request["json"]["stream"]) self.assertEqual(request["json"]["temperature"], 0.2) def test_generate_image_parses_chat_multimodal_data_url(self): generated = b"generated-image" encoded = base64.b64encode(generated).decode("ascii") session = FakeSession( FakeResponse( { "choices": [ { "message": { "content": [ {"type": "text", "text": "done"}, { "type": "image_url", "image_url": { "url": f"data:image/png;base64,{encoded}" }, }, ] } } ] } ) ) provider = ChatCompletionsProvider(session=session) model = ResolvedModel( name="Nano Banana 2", url="https://api.vectorengine.ai/v1/chat/completions", model="gemini-3.1-flash-image-preview", api_key="test-key", api_type="auto", ) result = provider.generate_image( "Generate product image", model, image=b"input-image", image_mime_type="image/jpeg", ) self.assertEqual(result.image, generated) content = session.posts[0]["json"]["messages"][0]["content"] self.assertEqual(content[0], {"type": "text", "text": "Generate product image"}) self.assertTrue(content[1]["image_url"]["url"].startswith("data:image/jpeg;base64,")) class ImagesEditsProviderTests(SimpleTestCase): def test_generate_image_builds_multipart_request_and_parses_base64(self): generated = b"edited-image" encoded = base64.b64encode(generated).decode("ascii") session = FakeSession(FakeResponse({"data": [{"b64_json": encoded}]})) provider = ImagesEditsProvider(session=session) model = ResolvedModel( name="GPT Image 2", url="https://api.vectorengine.ai/v1/images/edits", model="gpt-image-2", api_key="test-key", api_type="images_edits", ) result = provider.generate_image( "Replace background", model, image=b"source-image", image_mime_type="image/png", image_filename="source.png", resolution="1K", ) self.assertEqual(result.image, generated) request = session.posts[0] self.assertEqual(request["url"], "https://api.vectorengine.ai/v1/images/edits") self.assertEqual(request["headers"]["Authorization"], "Bearer test-key") self.assertEqual(request["data"]["model"], "gpt-image-2") self.assertEqual(request["data"]["size"], "1024x1024") self.assertEqual( request["files"]["image"], ("source.png", b"source-image", "image/png"), ) def test_images_edits_requires_input_image(self): provider = ImagesEditsProvider(session=FakeSession()) model = ResolvedModel( name="GPT Image 2", url="https://api.vectorengine.ai/v1/images/edits", model="gpt-image-2", api_key="test-key", api_type="images_edits", ) with self.assertRaises(AiCapabilityError): provider.generate_image("Replace background", model) with self.assertRaises(AiCapabilityError): provider.generate_text("Generate title", model) @override_settings(AI_KEY_ENCRYPTION_KEY=TEST_ENCRYPTION_KEY) class AiModelEncryptionTests(TestCase): def test_api_key_is_encrypted_and_resolved_model_decrypts_it(self): model = AiModel( name="GPT-5.5 text", url="https://api.vectorengine.ai/v1", model="gpt-5.5", api_type=AiModel.ApiType.CHAT, capabilities=["text"], ) model.set_api_key("sk-test-secret") model.save() self.assertNotIn("sk-test-secret", model.api_key_encrypted) self.assertTrue(model.api_key_encrypted.startswith("fernet:")) self.assertEqual(model.get_api_key(), "sk-test-secret") resolved = model.to_resolved_model() self.assertEqual(resolved.api_key, "sk-test-secret") self.assertEqual(resolved.capabilities, frozenset({"text"})) def test_capabilities_accept_dict_shape(self): model = AiModel( name="Vision image", url="https://api.vectorengine.ai/v1/images/edits", model="gpt-image-2", api_type=AiModel.ApiType.IMAGES_EDITS, capabilities={"image": True, "vision": True, "text": False}, ) self.assertEqual(model.capabilities_set(), frozenset({"image", "vision"})) @override_settings(AI_KEY_ENCRYPTION_KEY=TEST_ENCRYPTION_KEY) class AliasResolutionTests(TestCase): def setUp(self): self.text_model = self.create_model( name="GPT-5.5 text", model="gpt-5.5", api_type=AiModel.ApiType.CHAT, capabilities=["text", "vision"], ) self.image_model = self.create_model( name="GPT Image 2", url="https://api.vectorengine.ai/v1/images/edits", model="gpt-image-2", api_type=AiModel.ApiType.IMAGES_EDITS, capabilities=["image", "vision"], ) def create_model( self, *, name, model, capabilities, url="https://api.vectorengine.ai/v1", api_type=AiModel.ApiType.CHAT, ): ai_model = AiModel( name=name, url=url, model=model, api_type=api_type, capabilities=capabilities, ) ai_model.set_api_key("sk-test-secret") ai_model.save() return ai_model def test_resolve_default_title_alias(self): ModelAlias.objects.create( operation_type=ModelAlias.OperationType.TITLE, alias="title-standard", ai_model=self.text_model, is_default=True, ) resolved = resolve_alias(ModelAlias.OperationType.TITLE) self.assertEqual(resolved.model, "gpt-5.5") self.assertEqual(resolved.api_key, "sk-test-secret") def test_resolve_named_image_alias(self): ModelAlias.objects.create( operation_type=ModelAlias.OperationType.IMAGE, alias="image-edit", ai_model=self.image_model, ) resolved = resolve_alias(ModelAlias.OperationType.IMAGE, "image-edit") self.assertEqual(resolved.model, "gpt-image-2") self.assertIn("image", resolved.capabilities) def test_resolve_alias_rejects_capability_mismatch(self): ModelAlias.objects.create( operation_type=ModelAlias.OperationType.TITLE, alias="bad-title", ai_model=self.image_model, ) with self.assertRaises(ModelCapabilityError): resolve_alias(ModelAlias.OperationType.TITLE, "bad-title") def test_resolve_alias_ignores_inactive_aliases(self): ModelAlias.objects.create( operation_type=ModelAlias.OperationType.TITLE, alias="inactive-title", ai_model=self.text_model, is_active=False, ) with self.assertRaises(AliasNotFoundError): resolve_alias(ModelAlias.OperationType.TITLE, "inactive-title") def test_model_alias_save_rejects_second_default_for_operation(self): ModelAlias.objects.create( operation_type=ModelAlias.OperationType.TITLE, alias="title-standard", ai_model=self.text_model, is_default=True, ) with self.assertRaisesMessage(Exception, "Only one default alias"): ModelAlias.objects.create( operation_type=ModelAlias.OperationType.TITLE, alias="title-backup", ai_model=self.text_model, is_default=True, ) @override_settings(AI_KEY_ENCRYPTION_KEY=TEST_ENCRYPTION_KEY) class AiModelsImportTests(TestCase): def test_import_cmbot_config_encrypts_keys_and_creates_default_aliases(self): result = import_ai_models_config( { "models": [ { "name": "Nano Banana 2", "url": "https://api.vectorengine.ai/v1/chat/completions", "model": "gemini-3.1-flash-image-preview", "api_key": "sk-image", "api_type": "auto", "timeout_seconds": 0, "connect_timeout_seconds": 30, "extra_body": {}, }, { "name": "GPT-5.5 text", "url": "https://api.vectorengine.ai/v1", "model": "gpt-5.5", "api_key": "sk-text", "api_type": "chat", "timeout_seconds": 0, "connect_timeout_seconds": 30, "extra_body": {}, }, ] }, create_default_aliases=True, ) self.assertEqual(result, {"created": 2, "updated": 0, "aliases": 2}) text_model = AiModel.objects.get(name="GPT-5.5 text") image_model = AiModel.objects.get(name="Nano Banana 2") self.assertEqual(text_model.capabilities_set(), frozenset({"text", "vision"})) self.assertEqual(image_model.capabilities_set(), frozenset({"image", "vision"})) self.assertNotIn("sk-text", text_model.api_key_encrypted) self.assertEqual(text_model.get_api_key(), "sk-text") title_alias = ModelAlias.objects.get( operation_type=ModelAlias.OperationType.TITLE, alias="title-standard", ) image_alias = ModelAlias.objects.get( operation_type=ModelAlias.OperationType.IMAGE, alias="image-standard", ) self.assertEqual(title_alias.ai_model, text_model) self.assertEqual(image_alias.ai_model, image_model) @override_settings(AI_KEY_ENCRYPTION_KEY=TEST_ENCRYPTION_KEY) class AiModelAdminTests(TestCase): def test_admin_form_uses_write_only_api_key_field(self): request = RequestFactory().get("/admin/apps/ai/aimodel/add/") model_admin = AiModelAdmin(AiModel, AdminSite()) form_class = model_admin.get_form(request) self.assertIn("api_key", form_class.base_fields) self.assertNotIn("api_key_encrypted", form_class.base_fields) self.assertFalse(form_class.base_fields["api_key"].widget.render_value) def test_admin_form_encrypts_api_key_on_save(self): form_class = AiModelAdmin(AiModel, AdminSite()).form form = form_class( data={ "name": "GPT-5.5 text", "url": "https://api.vectorengine.ai/v1", "model": "gpt-5.5", "api_type": AiModel.ApiType.CHAT, "api_key": "sk-admin-secret", "capabilities": '["text"]', "timeout_seconds": 0, "connect_timeout_seconds": 30, "extra_body": "{}", "is_active": "on", } ) self.assertTrue(form.is_valid(), form.errors) ai_model = form.save() self.assertNotIn("sk-admin-secret", ai_model.api_key_encrypted) self.assertEqual(ai_model.get_api_key(), "sk-admin-secret") @override_settings(AI_KEY_ENCRYPTION_KEY="") def test_admin_form_reports_missing_encryption_key(self): form_class = AiModelAdmin(AiModel, AdminSite()).form form = form_class( data={ "name": "GPT-5.5 text", "url": "https://api.vectorengine.ai/v1", "model": "gpt-5.5", "api_type": AiModel.ApiType.CHAT, "api_key": "sk-admin-secret", "capabilities": '["text"]', "timeout_seconds": 0, "connect_timeout_seconds": 30, "extra_body": "{}", "is_active": "on", } ) self.assertFalse(form.is_valid()) self.assertIn("AI_KEY_ENCRYPTION_KEY is not configured", str(form.errors))