Shaman
74 lines
2.3 KiB
Python
74 lines
2.3 KiB
Python
import pytest
|
|
|
|
from app.config import DEFAULT_MAPPING
|
|
from app.errors import InvalidRequest
|
|
from app.injector import build_workflow, new_seed, parse_size
|
|
from app.templates import load_all
|
|
from app.config import Settings
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def templates():
|
|
return load_all(Settings(_env_file=None))
|
|
|
|
|
|
def test_parse_size():
|
|
assert parse_size("1024x1024") == (1024, 1024)
|
|
assert parse_size("832X1216") == (832, 1216)
|
|
|
|
|
|
@pytest.mark.parametrize("bad", ["1024", "ax1024", "1024x", "0x100", "1024x-2", "wide"])
|
|
def test_parse_size_invalid(bad):
|
|
with pytest.raises(InvalidRequest):
|
|
parse_size(bad)
|
|
|
|
|
|
def test_seed_random_differs():
|
|
assert new_seed(None) != new_seed(None)
|
|
|
|
|
|
def test_seed_fixed_and_range():
|
|
assert new_seed(42) == 42
|
|
with pytest.raises(InvalidRequest):
|
|
new_seed(-1)
|
|
|
|
|
|
def test_create_injection(templates):
|
|
create, _ = templates
|
|
mapping = DEFAULT_MAPPING["create"]
|
|
graph = build_workflow(create, mapping, prompt="a cat", seed=77, size="832x1216")
|
|
assert graph["83:27"]["inputs"]["text"] == "a cat"
|
|
assert graph["83:13"]["inputs"]["width"] == 832
|
|
assert graph["83:13"]["inputs"]["height"] == 1216
|
|
assert graph["83:3"]["inputs"]["seed"] == 77
|
|
# template untouched
|
|
assert create["83:27"]["inputs"]["text"] == "\n"
|
|
assert create["83:13"]["inputs"]["width"] == 1280
|
|
|
|
|
|
def test_edit_injection(templates):
|
|
_, edit = templates
|
|
mapping = DEFAULT_MAPPING["edit"]
|
|
graph = build_workflow(edit, mapping, prompt="make it night", seed=5, image_name="adapter_x.png")
|
|
assert graph["115:111"]["inputs"]["prompt"] == "make it night"
|
|
assert graph["115:110"]["inputs"]["prompt"] == ""
|
|
assert graph["78"]["inputs"]["image"] == "adapter_x.png"
|
|
assert graph["115:3"]["inputs"]["seed"] == 5
|
|
|
|
|
|
def test_edit_seed_varies_per_call(templates):
|
|
_, edit = templates
|
|
mapping = DEFAULT_MAPPING["edit"]
|
|
seeds = set()
|
|
for _ in range(3):
|
|
g = build_workflow(edit, mapping, prompt="p", seed=new_seed(None))
|
|
seeds.add(g["115:3"]["inputs"]["seed"])
|
|
assert len(seeds) == 3
|
|
|
|
|
|
def test_unknown_mapping_node_raises(templates):
|
|
create, _ = templates
|
|
bad_mapping = {"prompt": {"node": "999", "field": "text"}}
|
|
with pytest.raises(InvalidRequest, match="999"):
|
|
build_workflow(create, bad_mapping, prompt="p", seed=1)
|