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()