Shaman
102 lines
3.8 KiB
Python
102 lines
3.8 KiB
Python
"""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()
|