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

tests/test_templates.py

@@ -1,7 +1,15 @@
 import pytest
 
 from app.config import Settings
-from app.templates import TemplateError, check_mapping, find_broken_links, load_all
+from app.errors import InvalidRequest
+from app.injector import extract_tokens
+from app.templates import (
+    TemplateError,
+    find_broken_links,
+    is_compatible,
+    load_registry,
+    select_model,
+)
 
 
 def test_find_broken_links_detects_missing_node():
@@ -22,36 +30,84 @@ def test_find_broken_links_ok_on_valid_graph():
     assert find_broken_links(graph) == []
 
 
-def test_check_mapping_missing_node():
-    graph = {"1": {"class_type": "X", "inputs": {"text": ""}}}
-    mapping = {"create": {"prompt": {"node": "404", "field": "text"}}}
-    errors = check_mapping(graph, mapping["create"], "create")
-    assert any("not found" in e for e in errors)
+def test_extract_tokens_reports_known_and_unknown():
+    graph = {
+        "1": {"class_type": "X", "inputs": {"text": "{prompt}", "seed": "{seed}"}},
+        "2": {"class_type": "Y", "inputs": {"note": "static {notatoken} text"}},
+    }
+    declared, unknown = extract_tokens(graph)
+    assert declared == {"prompt", "seed"}
+    assert len(unknown) == 1 and "notatoken" in unknown[0]
 
 
-def test_check_mapping_non_scalar_field():
-    graph = {"1": {"class_type": "X", "inputs": {"text": ["2", 0]}}}
-    mapping = {"create": {"prompt": {"node": "1", "field": "text"}}}
-    errors = check_mapping(graph, mapping["create"], "create")
-    assert any("no scalar input" in e for e in errors)
+def test_shipped_registry_loads_with_expected_tokens():
+    registry = load_registry(Settings(_env_file=None))
+    assert set(registry) == {"z-image-turbo", "qwen-image-edit"}
+    assert registry["z-image-turbo"].declared == frozenset({"prompt", "seed", "width", "height"})
+    assert registry["qwen-image-edit"].declared == frozenset({"prompt", "negative", "seed", "image"})
 
 
-def test_shipped_templates_load_cleanly():
-    s = Settings(_env_file=None)
-    create, edit = load_all(s)
-    assert create["83:28"]["class_type"] == "UNETLoader"
-    assert edit["115:89"]["class_type"] == "LoraLoaderModelOnly"
+def test_generations_graph_incompatible_with_edits():
+    assert not is_compatible(frozenset({"prompt", "seed", "width", "height"}), "edits")
 
 
-def test_broken_template_blocks_startup(tmp_path):
+def test_edits_graph_incompatible_with_generations():
+    assert not is_compatible(frozenset({"prompt", "negative", "seed", "image"}), "generations")
+
+
+def test_select_model_defaults(registry):
+    assert select_model(registry, None, "generations").id == "z-image-turbo"
+    assert select_model(registry, None, "edits").id == "qwen-image-edit"
+
+
+def test_select_model_explicit(registry):
+    assert select_model(registry, "qwen-image-edit", "edits").id == "qwen-image-edit"
+
+
+def test_select_model_unknown_400(registry):
+    with pytest.raises(InvalidRequest, match="unknown model"):
+        select_model(registry, "dall-e-3", "generations")
+
+
+def test_select_model_incompatible_400(registry):
+    with pytest.raises(InvalidRequest, match="not compatible"):
+        select_model(registry, "qwen-image-edit", "generations")
+    with pytest.raises(InvalidRequest, match="not compatible"):
+        select_model(registry, "z-image-turbo", "edits")
+
+
+def test_select_model_no_suitable_400(registry):
+    gen_only = {"z-image-turbo": registry["z-image-turbo"]}
+    with pytest.raises(InvalidRequest, match="no model"):
+        select_model(gen_only, None, "edits")
+
+
+def test_broken_workflow_blocks_startup(tmp_path):
     bad = tmp_path / "bad.json"
     bad.write_text('{"1": {"class_type": "X", "inputs": {"m": ["9", 0]}}}', encoding="utf-8")
-    good = tmp_path / "good.json"
-    good.write_text(
-        '{"1": {"class_type": "X", "inputs": {"m": ["2", 0]}}, "2": {"class_type": "Y", "inputs": {}}}',
-        encoding="utf-8",
-    )
-    s = Settings(_env_file=None, create_template_path=bad, edit_template_path=good)
+    models = tmp_path / "models.json"
+    models.write_text('{"m1": {"workflow": "bad.json"}}', encoding="utf-8")
     with pytest.raises(TemplateError) as exc:
-        load_all(s)
+        load_registry(Settings(_env_file=None, models_file=models))
     assert "missing node '9'" in str(exc.value)
+
+
+def test_unknown_token_blocks_startup(tmp_path):
+    wf = tmp_path / "wf.json"
+    wf.write_text('{"1": {"class_type": "X", "inputs": {"text": "{bogus}"}}}', encoding="utf-8")
+    models = tmp_path / "models.json"
+    models.write_text('{"m1": {"workflow": "wf.json"}}', encoding="utf-8")
+    with pytest.raises(TemplateError, match="unknown token"):
+        load_registry(Settings(_env_file=None, models_file=models))
+
+
+def test_missing_workflow_file_blocks_startup(tmp_path):
+    models = tmp_path / "models.json"
+    models.write_text('{"m1": {"workflow": "absent.json"}}', encoding="utf-8")
+    with pytest.raises(TemplateError, match="not found"):
+        load_registry(Settings(_env_file=None, models_file=models))
+
+
+@pytest.fixture(scope="module")
+def registry():
+    return load_registry(Settings(_env_file=None))