| 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"], |
| }) |
| |
| 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"]}) |
|
|
| |
| 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() |
|
|
