"""Offline tests for the image and video AI-writing helpers.""" from __future__ import annotations import asyncio import base64 from io import BytesIO import json from pathlib import Path import sys import tempfile import unittest from PIL import Image ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) from utils import chat_support as CHAT # noqa: E402 class _FakeResponse: def __init__(self, status, payload): self.status = status self.payload = payload async def __aenter__(self): return self async def __aexit__(self, _exc_type, _exc, _tb): return False async def json(self, **_kwargs): return self.payload class _FakeSession: def __init__(self, response): self.response = response self.calls = [] def post(self, url, **kwargs): self.calls.append((url, kwargs)) return self.response class PromptOptimizerTests(unittest.TestCase): def test_payload_uses_fixed_model_high_reasoning_and_ordered_images(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) Image.new("RGB", (80, 120), "red").save(root / "first.png") Image.new("RGB", (120, 80), "blue").save(root / "second.png") payload = CHAT.build_prompt_optimization_payload( "保留人物身份，改成雨夜街景", [ {"name": "first.png", "subfolder": "", "type": "input"}, {"name": "second.png", "subfolder": "", "type": "input"}, ], directory, ) self.assertEqual(payload["model"], "gpt-5.6-sol") self.assertEqual(payload["reasoning_effort"], "high") self.assertFalse(payload["stream"]) self.assertIn("视觉元素绑定优先", payload["messages"][0]["content"]) self.assertIn("必须保持不变", payload["messages"][0]["content"]) content = payload["messages"][1]["content"] self.assertEqual(content[1]["text"], "参考图 1（上传顺序第 1 张）") self.assertEqual(content[3]["text"], "参考图 2（上传顺序第 2 张）") image_items = [item for item in content if item["type"] == "image_url"] self.assertEqual(len(image_items), 2) for item in image_items: url = item["image_url"]["url"] self.assertTrue(url.startswith("data:image/jpeg;base64,")) with Image.open(BytesIO(base64.b64decode(url.split(",", 1)[1]))) as image: self.assertEqual(image.format, "JPEG") body_size = len( json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8") ) self.assertLessEqual(body_size, CHAT.PROMPT_OPTIMIZER_BODY_LIMIT_BYTES) def test_rejects_reference_path_outside_input_directory(self): with tempfile.TemporaryDirectory() as directory: with self.assertRaisesRegex(ValueError, "路径不安全"): CHAT.build_prompt_optimization_payload( "测试提示词", [{"name": "outside.png", "subfolder": "..", "type": "input"}], directory, ) def test_non_streaming_request_extracts_only_the_final_prompt(self): session = _FakeSession(_FakeResponse(200, { "choices": [{"message": {"content": "```text\n最终视觉提示词\n```"}}], })) with tempfile.TemporaryDirectory() as directory: result = asyncio.run(CHAT.optimize_image_prompt( session, "https://example.invalid", "offline-test-token", "原始提示词", [], directory, )) self.assertEqual(result, "最终视觉提示词") self.assertEqual(len(session.calls), 1) url, kwargs = session.calls[0] self.assertEqual(url, "https://example.invalid/v1/chat/completions") self.assertEqual(kwargs["json"]["model"], "gpt-5.6-sol") self.assertEqual(kwargs["json"]["reasoning_effort"], "high") self.assertFalse(kwargs["json"]["stream"]) def test_video_payload_uses_dedicated_preset_and_ordered_frame_roles(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) Image.new("RGB", (96, 128), "green").save(root / "start.png") Image.new("RGB", (128, 96), "yellow").save(root / "end.png") payload = CHAT.build_video_prompt_writing_payload( "人物从门口走到窗前", [ {"name": "start.png", "subfolder": "", "type": "input"}, {"name": "end.png", "subfolder": "", "type": "input"}, ], directory, { "generation_mode": "first_last_frame", "duration": "8", "aspect_ratio": "16:9", "generate_audio": True, "reference_video_count": 1, "reference_audio_count": 2, }, ) self.assertEqual(payload["model"], "gpt-5.6-sol") self.assertEqual(payload["reasoning_effort"], "high") self.assertFalse(payload["stream"]) system_prompt = payload["messages"][0]["content"] self.assertIn("连续的时间过程", system_prompt) self.assertIn("首尾帧生视频", system_prompt) content = payload["messages"][1]["content"] self.assertIn("生成模式：首尾帧生视频", content[0]["text"]) self.assertIn("目标时长：8", content[0]["text"]) self.assertIn("生成音频：开启", content[0]["text"]) self.assertIn("参考视频 1 个、参考音频 2 个", content[0]["text"]) self.assertEqual(content[1]["text"], "首帧（第 1 张分析图）") self.assertEqual(content[3]["text"], "尾帧（第 2 张分析图）") self.assertEqual(len([item for item in content if item["type"] == "image_url"]), 2) def test_video_request_uses_non_streaming_completion(self): session = _FakeSession(_FakeResponse(200, { "choices": [{"message": {"content": "连续镜头中的最终视频提示词"}}], })) with tempfile.TemporaryDirectory() as directory: result = asyncio.run(CHAT.write_video_prompt( session, "https://example.invalid", "offline-test-token", "原始视频创意", [], directory, {"generation_mode": "text", "duration": "5"}, )) self.assertEqual(result, "连续镜头中的最终视频提示词") self.assertEqual(session.calls[0][0], "https://example.invalid/v1/chat/completions") request = session.calls[0][1]["json"] self.assertIn("AI 视频生成提示词编写助手", request["messages"][0]["content"]) self.assertFalse(request["stream"]) if __name__ == "__main__": unittest.main(verbosity=2)