cmshoppe/tests/test_ai.py

import base64
import io
import json
import os
import sys
import unittest
from unittest import mock

sys.path.insert(0, os.path.dirname(__file__))

from _helpers import TempDirMixin

from app import ai, appconfig


class _Response:
    def __init__(self, payload):
        self.payload = payload

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False

    def read(self, size=-1):
        return json.dumps(self.payload).encode("utf-8")


class AITests(TempDirMixin, unittest.TestCase):
    def _write_models(self, path, text=None, image=None):
        text = text or {
            "name": "Text",
            "category": "text",
            "enabled": True,
            "url": "https://example.invalid/v1/chat/completions",
            "model": "text-model",
            "api_key": "sk-text-secret",
            "api_type": "chat",
            "connect_timeout_seconds": 1,
            "timeout_seconds": 1,
            "extra_body": {"temperature": 0},
        }
        image = image or {
            "name": "Image",
            "category": "image",
            "enabled": True,
            "url": "https://example.invalid/v1/chat/completions",
            "model": "image-model",
            "api_key": "sk-image-secret",
            "api_type": "auto",
            "connect_timeout_seconds": 1,
            "timeout_seconds": 1,
            "extra_body": {},
        }
        appconfig.save_ai_models_config({"models": [text, image]}, path=path)

    def _config(self):
        cfg = appconfig.default_config()
        cfg["ai"]["default_text_model"] = "Text"
        cfg["ai"]["default_image_model"] = "Image"
        cfg["ai"]["retry"] = 1
        cfg["ai"]["resolution"] = "512"
        cfg["ai"]["jpg_quality"] = 80
        return cfg

    def test_gen_title_uses_configured_model_and_retries(self):
        with self.make_temp_dir() as temp_dir:
            models_path = os.path.join(temp_dir, "ai_models.json")
            self._write_models(models_path)
            calls = []

            def fake_urlopen(request, timeout=None):
                calls.append((request, timeout))
                if len(calls) == 1:
                    raise ai.urllib.error.URLError("temporary")
                return _Response({"choices": [{"message": {"content": " 新标题 "}}]})

            with mock.patch("app.ai.urllib.request.urlopen", side_effect=fake_urlopen):
                title = ai.gen_title(
                    "优化标题",
                    "旧标题",
                    config=self._config(),
                    models_path=models_path,
                )

            self.assertEqual("新标题", title)
            self.assertEqual(2, len(calls))
            body = json.loads(calls[-1][0].data.decode("utf-8"))
            self.assertEqual("text-model", body["model"])
            self.assertEqual(0, body["temperature"])
            self.assertNotIn("sk-text-secret", body["messages"][1]["content"])

        self.assert_removed(temp_dir)

    def test_missing_model_fields_raise_clear_error_without_secret(self):
        with self.make_temp_dir() as temp_dir:
            models_path = os.path.join(temp_dir, "ai_models.json")
            self._write_models(
                models_path,
                text={
                    "name": "Text",
                    "category": "text",
                    "enabled": True,
                    "url": "",
                    "model": "text-model",
                    "api_key": "sk-text-secret",
                    "api_type": "chat",
                    "connect_timeout_seconds": 1,
                    "timeout_seconds": 1,
                    "extra_body": {},
                },
            )

            with self.assertRaises(ai.AIError) as raised:
                ai.gen_title("prompt", "old", config=self._config(), models_path=models_path)

            message = str(raised.exception)
            self.assertIn("缺少字段: url", message)
            self.assertNotIn("sk-text-secret", message)

        self.assert_removed(temp_dir)

    def test_gen_cover_saves_jpeg_with_resolution_and_quality(self):
        try:
            from PIL import Image
        except ImportError:
            self.skipTest("Pillow not installed")

        with self.make_temp_dir() as temp_dir:
            models_path = os.path.join(temp_dir, "ai_models.json")
            self._write_models(models_path)
            old_cover = os.path.join(temp_dir, "old.jpg")
            output = os.path.join(temp_dir, "new.jpg")
            Image.new("RGB", (16, 16), (20, 30, 40)).save(old_cover, "JPEG")
            generated = io.BytesIO()
            Image.new("RGB", (8, 8), (200, 120, 80)).save(generated, "PNG")
            b64_image = base64.b64encode(generated.getvalue()).decode("ascii")

            def fake_urlopen(request, timeout=None):
                body = json.loads(request.data.decode("utf-8"))
                self.assertEqual("image-model", body["model"])
                self.assertIn("目标分辨率：512", body["messages"][0]["content"][0]["text"])
                return _Response({"data": [{"b64_json": b64_image}]})

            with mock.patch("app.ai.urllib.request.urlopen", side_effect=fake_urlopen):
                result = ai.gen_cover(
                    "生成封面",
                    old_cover,
                    output,
                    config=self._config(),
                    models_path=models_path,
                )

            self.assertEqual(os.path.abspath(output), result)
            with Image.open(output) as saved:
                self.assertEqual((512, 512), saved.size)
                self.assertEqual("JPEG", saved.format)

        self.assert_removed(temp_dir)


if __name__ == "__main__":
    unittest.main()