from __future__ import annotations import json import tempfile import unittest from pathlib import Path from unittest import mock from fastapi import FastAPI from fastapi.testclient import TestClient from api import ai as ai_module from services.chat_model_service import resolve_chat_models from services.config import ConfigStore, _normalize_image_model_overrides from services.image_model_service import resolve_image_models def _upstream(model_id: str, **extra: object) -> dict[str, object]: return {"id": model_id, "object": "model", "owned_by": "chatgpt", **extra} def _builtin_image(model_id: str = "gpt-image-2") -> dict[str, object]: return {"id": model_id, "object": "model", "owned_by": "chatgpt2api"} def _third_party(model_id: str, capabilities: list[str] | None = None) -> dict[str, object]: return { "id": model_id, "object": "model", "owned_by": "aichimo", "group_id": "g1", "capabilities": ["txt2img", "img2img"] if capabilities is None else capabilities, } class ImageModelOverrideNormalizationTests(unittest.TestCase): def test_trims_dedupes_and_include_wins(self) -> None: normalized = _normalize_image_model_overrides({ "include": [" gpt-image-2 ", "", None, "gpt-image-2"], "exclude": ["Nano Banana 2-4k", "Nano Banana 2-4k", "gpt-image-2"], }) # include 与 exclude 冲突时以 include 为准 self.assertEqual(normalized, {"include": ["gpt-image-2"], "exclude": ["Nano Banana 2-4k"]}) def test_invalid_payloads_fall_back_to_empty(self) -> None: for value in (None, "gpt-image-2", 42, ["gpt-image-2"]): self.assertEqual(_normalize_image_model_overrides(value), {"include": [], "exclude": []}) def test_overrides_round_trip_through_config_store(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: path = Path(tmp_dir) / "config.json" path.write_text(json.dumps({"auth-key": "test-auth"}), encoding="utf-8") store = ConfigStore(path) self.assertEqual(store.image_models, {"include": [], "exclude": []}) saved = store.update({"image_models": {"include": [], "exclude": ["Nano Banana 2-4k"]}}) self.assertEqual(saved["image_models"], {"include": [], "exclude": ["Nano Banana 2-4k"]}) self.assertEqual(store.get()["image_models"], {"include": [], "exclude": ["Nano Banana 2-4k"]}) # 未提及 image_models 的局部更新不应把它清掉。 store.update({"proxy": "http://proxy.example:8080"}) self.assertEqual(store.get()["image_models"], {"include": [], "exclude": ["Nano Banana 2-4k"]}) # 与对话侧互不干扰，且都能落盘。 store.update({"chat_models": {"include": ["research"], "exclude": []}}) reloaded = ConfigStore(path) self.assertEqual(reloaded.image_models, {"include": [], "exclude": ["Nano Banana 2-4k"]}) self.assertEqual(reloaded.chat_models, {"include": ["research"], "exclude": []}) class ResolveImageModelsTests(unittest.TestCase): def setUp(self) -> None: self.items = [ _upstream("auto"), _upstream("gpt-5-6"), _upstream("research"), _builtin_image("gpt-image-2"), _builtin_image("codex-gpt-image-2"), _third_party("Nano Banana 2-4k"), _third_party("gemini-3-pro-image"), _third_party("qwen3-max", capabilities=[]), # 多模态：既能聊也能生图 _third_party("omni-chat-image", capabilities=["chat", "txt2img"]), ] def resolved(self, overrides: dict[str, object] | None = None) -> dict[str, dict[str, object]]: return {item["id"]: item for item in resolve_image_models(self.items, overrides)} def test_only_image_capable_models_are_candidates(self) -> None: resolved = self.resolved() self.assertEqual( set(resolved), {"gpt-image-2", "codex-gpt-image-2", "Nano Banana 2-4k", "gemini-3-pro-image", "omni-chat-image"}, ) for model_id in ("auto", "gpt-5-6", "research", "qwen3-max"): self.assertNotIn(model_id, resolved, f"{model_id} 不能生图，不该进候选") def test_sources_are_labelled(self) -> None: resolved = self.resolved() self.assertEqual(resolved["gpt-image-2"]["image_source"], "builtin") self.assertEqual(resolved["codex-gpt-image-2"]["image_source"], "builtin") self.assertEqual(resolved["Nano Banana 2-4k"]["image_source"], "third_party") self.assertEqual(resolved["gpt-image-2"]["image_capabilities"], []) self.assertEqual(resolved["Nano Banana 2-4k"]["image_capabilities"], ["img2img", "txt2img"]) def test_all_models_default_to_enabled(self) -> None: """默认全开：装完这个功能不应该让现有生图模型凭空消失。""" for item in resolve_image_models(self.items): self.assertTrue(item["image_default_enabled"]) self.assertTrue(item["image_enabled"]) def test_exclude_hides_a_model(self) -> None: resolved = self.resolved({"include": [], "exclude": ["Nano Banana 2-4k", "gpt-image-2"]}) self.assertFalse(resolved["Nano Banana 2-4k"]["image_enabled"]) self.assertFalse(resolved["gpt-image-2"]["image_enabled"]) self.assertTrue(resolved["gemini-3-pro-image"]["image_enabled"]) def test_include_wins_over_exclude(self) -> None: resolved = self.resolved({"include": ["Nano Banana 2-4k"], "exclude": ["Nano Banana 2-4k"]}) self.assertTrue(resolved["Nano Banana 2-4k"]["image_enabled"]) def test_enabled_models_sort_first_and_builtin_before_third_party(self) -> None: resolved = resolve_image_models(self.items, {"include": [], "exclude": ["gpt-image-2"]}) keys = [(0 if item["image_enabled"] else 1, item["image_source"], item["id"]) for item in resolved] self.assertEqual(keys, sorted(keys)) self.assertEqual(resolved[-1]["id"], "gpt-image-2") def test_multimodal_model_appears_in_both_lists(self) -> None: image_ids = {item["id"] for item in resolve_image_models(self.items)} chat_flags = {item["id"]: item["chat_enabled"] for item in resolve_chat_models(self.items)} self.assertIn("omni-chat-image", image_ids) self.assertTrue(chat_flags["omni-chat-image"], "多模态模型应同时出现在对话下拉里") # 纯生图模型仍然只在生图列表、不会混进对话 self.assertIn("Nano Banana 2-4k", image_ids) self.assertFalse(chat_flags["Nano Banana 2-4k"]) def test_chat_only_third_party_model_stays_out_of_the_image_list(self) -> None: """第三方接入里的文本模型（capabilities 只有 chat）不是生图候选，也不该被当成生图模型。""" items = [_third_party("gpt-5.5", capabilities=["chat"])] self.assertEqual(resolve_image_models(items), []) chat = {item["id"]: item for item in resolve_chat_models(items)}["gpt-5.5"] self.assertFalse(chat["chat_is_image"], "文本模型不该带「生图」角标") self.assertTrue(chat["chat_default_enabled"], "文本模型应默认出现在对话下拉里") self.assertEqual(chat["chat_default_reason"], "") class ImageModelsEndpointTests(unittest.TestCase): def _get( self, payload: list[dict[str, object]], overrides: dict[str, list[str]] | None = None, third_party: list[dict[str, object]] | None = None, ): with ( mock.patch.object(ai_module, "require_identity", return_value={"role": "admin"}), mock.patch.object( ai_module.openai_v1_models, "list_models", return_value={"object": "list", "data": payload}, ), mock.patch.object( ai_module.third_party_image_service, "public_models_for_identity", return_value=third_party or [], ), mock.patch.object( type(ai_module.config), "image_models", new_callable=mock.PropertyMock, return_value=overrides or {"include": [], "exclude": []}, ), ): app = FastAPI() app.include_router(ai_module.create_router()) return TestClient(app).get("/api/image/models") def test_returns_only_image_models_with_flags(self) -> None: response = self._get([ _upstream("auto"), _upstream("gpt-5-6"), _builtin_image("gpt-image-2"), ]) self.assertEqual(response.status_code, 200) body = response.json() self.assertEqual(body["object"], "list") self.assertEqual(body["include"], []) self.assertEqual(body["exclude"], []) ids = [item["id"] for item in body["data"]] self.assertEqual(ids, ["gpt-image-2"]) self.assertTrue(body["data"][0]["image_enabled"]) self.assertEqual(body["data"][0]["image_source"], "builtin") def test_merges_third_party_image_models(self) -> None: response = self._get( [_upstream("auto"), _builtin_image("gpt-image-2")], third_party=[_third_party("Nano Banana 2-4k"), _third_party("qwen3-max", capabilities=[])], ) flags = {item["id"]: item["image_enabled"] for item in response.json()["data"]} self.assertEqual(flags, {"gpt-image-2": True, "Nano Banana 2-4k": True}) sources = {item["id"]: item["image_source"] for item in response.json()["data"]} self.assertEqual(sources["Nano Banana 2-4k"], "third_party") def test_echoes_overrides_and_applies_them(self) -> None: response = self._get( [_builtin_image("gpt-image-2")], overrides={"include": [], "exclude": ["gpt-image-2"]}, third_party=[_third_party("Nano Banana 2-4k")], ) body = response.json() self.assertEqual(body["exclude"], ["gpt-image-2"]) flags = {item["id"]: item["image_enabled"] for item in body["data"]} self.assertFalse(flags["gpt-image-2"]) self.assertTrue(flags["Nano Banana 2-4k"]) if __name__ == "__main__": unittest.main()