Доработка управления моделями

tests/test_registry_integration.py

@@ -0,0 +1,101 @@
+"""A third model added purely by editing models.json (no code change): it shows up
+in /v1/models, is dry-inject validated at startup, and is selectable by id."""
+import contextlib
+import json
+
+import httpx
+import respx
+
+from app.config import REPO_ROOT, Settings
+from app.main import create_app
+from tests.conftest import COMFY, FAKE_KEY
+
+PROMPT_ID = "p-int"
+
+
+def _extra_generation_workflow() -> dict:
+    # declared = {prompt, negative, seed, width, height} -> generations-compatible, no {image}
+    return {
+        "1": {
+            "class_type": "X",
+            "inputs": {
+                "text": "{prompt}",
+                "neg": "{negative}",
+                "sd": "{seed}",
+                "w": "{width}",
+                "h": "{height}",
+            },
+        }
+    }
+
+
+def _models_file(tmp_path) -> Settings:
+    (tmp_path / "extra.json").write_text(json.dumps(_extra_generation_workflow()), encoding="utf-8")
+    registry = {
+        "z-image-turbo": {"workflow": str(REPO_ROOT / "workflows" / "z_image_turbo.json")},
+        "extra-gen": {"workflow": "extra.json"},  # relative to models.json dir
+        "qwen-image-edit": {"workflow": str(REPO_ROOT / "workflows" / "qwen_image_edit.json")},
+    }
+    models = tmp_path / "models.json"
+    models.write_text(json.dumps(registry), encoding="utf-8")
+    return Settings(
+        _env_file=None,
+        comfy_base_url=COMFY,
+        adapter_api_key=FAKE_KEY,
+        models_file=models,
+        validate_templates_on_server=True,
+        poll_interval=0.0,
+        execution_timeout=5.0,
+    )
+
+
+async def _with_lifespan(app):
+    stack = contextlib.AsyncExitStack()
+    await stack.enter_async_context(app.router.lifespan_context(app))
+    transport = httpx.ASGITransport(app=app)
+    ac = await stack.enter_async_context(httpx.AsyncClient(transport=transport, base_url="http://adapter.test"))
+    ac.headers["Authorization"] = f"Bearer {FAKE_KEY}"
+    return stack, ac
+
+
+@respx.mock
+async def test_third_model_listed_and_validated_without_code_change(tmp_path):
+    validate_route = respx.post(f"{COMFY}/prompt").mock(return_value=httpx.Response(200, json={"prompt_id": PROMPT_ID}))
+    app = create_app(_models_file(tmp_path))
+    stack, ac = await _with_lifespan(app)
+    try:
+        ids = [m["id"] for m in (await ac.get("/v1/models")).json()["data"]]
+        assert ids == ["z-image-turbo", "extra-gen", "qwen-image-edit"]
+        # every model dry-injected and validated once at startup
+        assert validate_route.call_count == 3
+    finally:
+        await stack.aclose()
+
+
+@respx.mock
+async def test_extra_model_selectable_and_injects(tmp_path):
+    prompt_route = respx.post(f"{COMFY}/prompt").mock(return_value=httpx.Response(200, json={"prompt_id": PROMPT_ID}))
+    respx.get(f"{COMFY}/history/{PROMPT_ID}").mock(
+        return_value=httpx.Response(
+            200,
+            json={PROMPT_ID: {"outputs": {"1": {"images": [{"filename": "o.png", "subfolder": "", "type": "output"}]}},
+                              "status": {"status_str": "success", "completed": True}}},
+        )
+    )
+    respx.get(f"{COMFY}/view").mock(return_value=httpx.Response(200, content=b"\x89PNG\r\n\x1a\n"))
+
+    settings = _models_file(tmp_path).model_copy(update={"validate_templates_on_server": False})
+    app = create_app(settings)
+    stack, ac = await _with_lifespan(app)
+    try:
+        r = await ac.post(
+            "http://adapter.test/v1/images/generations",
+            json={"model": "extra-gen", "prompt": "hello", "size": "512x768", "seed": 9},
+        )
+        assert r.status_code == 200
+        sent = json.loads(prompt_route.calls.last.request.content.decode())["prompt"]
+        assert sent["1"]["inputs"]["text"] == "hello"
+        assert sent["1"]["inputs"]["sd"] == 9
+        assert sent["1"]["inputs"]["w"] == 512 and sent["1"]["inputs"]["h"] == 768
+    finally:
+        await stack.aclose()