Shaman
114 lines
4.0 KiB
Python
114 lines
4.0 KiB
Python
import pytest
|
|
|
|
from app.config import Settings
|
|
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():
|
|
graph = {
|
|
"1": {"class_type": "X", "inputs": {"m": ["9", 0]}},
|
|
}
|
|
errors = find_broken_links(graph)
|
|
assert len(errors) == 1
|
|
assert "missing node '9'" in errors[0]
|
|
assert "node 1" in errors[0]
|
|
|
|
|
|
def test_find_broken_links_ok_on_valid_graph():
|
|
graph = {
|
|
"1": {"class_type": "X", "inputs": {"m": ["2", 0]}},
|
|
"2": {"class_type": "Y", "inputs": {}},
|
|
}
|
|
assert find_broken_links(graph) == []
|
|
|
|
|
|
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_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_generations_graph_incompatible_with_edits():
|
|
assert not is_compatible(frozenset({"prompt", "seed", "width", "height"}), "edits")
|
|
|
|
|
|
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")
|
|
models = tmp_path / "models.json"
|
|
models.write_text('{"m1": {"workflow": "bad.json"}}', encoding="utf-8")
|
|
with pytest.raises(TemplateError) as exc:
|
|
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))
|