cmhub/apps/ai/tests.py

import base64

from django.test import SimpleTestCase

from apps.ai.providers import AiCapabilityError, ResolvedModel, get_provider, resolve_api_type
from apps.ai.providers.openai_compatible import ChatCompletionsProvider, ImagesEditsProvider


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)