320 lines
8.9 KiB
Python
320 lines
8.9 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import random
|
|
|
|
|
|
def test_load_entries_supports_common_json_shapes(isolated_gallery_env):
|
|
extractor = isolated_gallery_env.extractor
|
|
library_path = isolated_gallery_env.data_dir / "prompts.json"
|
|
library_path.write_text(
|
|
json.dumps(
|
|
[
|
|
"plain text",
|
|
{"prompt": "prompt field"},
|
|
{"name": "name field"},
|
|
{"title": "title field"},
|
|
["first list item", "ignored"],
|
|
{"prompt": ""},
|
|
42,
|
|
],
|
|
ensure_ascii=False,
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
assert extractor.load_entries("prompts.json") == [
|
|
"plain text",
|
|
"prompt field",
|
|
"name field",
|
|
"title field",
|
|
"first list item",
|
|
]
|
|
|
|
|
|
def test_load_segments_supports_nested_paths_lists_and_filters(isolated_gallery_env):
|
|
extractor = isolated_gallery_env.extractor
|
|
library = [
|
|
{
|
|
"name": "artist_a",
|
|
"other_names": ["alias_a", "alias_b"],
|
|
"meta": {"style": ["sketch", "line art"], "group": "ink"},
|
|
"post_count": 120,
|
|
},
|
|
{
|
|
"name": "artist_b",
|
|
"other_names": ["alias_c"],
|
|
"meta": {"style": ["watercolor"], "group": "paint"},
|
|
"post_count": 20,
|
|
},
|
|
]
|
|
(isolated_gallery_env.data_dir / "segments.json").write_text(json.dumps(library), encoding="utf-8")
|
|
|
|
assert extractor.load_segments("segments.json", "other_names") == ["alias_a", "alias_b", "alias_c"]
|
|
assert extractor.load_segments("segments.json", "meta.style") == ["sketch", "line art", "watercolor"]
|
|
assert extractor.load_segments("segments.json", "name, meta.group", "meta.group", "ink", "equals") == ["artist_a", "ink"]
|
|
|
|
|
|
def test_load_segments_accepts_plain_string_lists_for_tag_libraries(isolated_gallery_env):
|
|
extractor = isolated_gallery_env.extractor
|
|
(isolated_gallery_env.data_dir / "tags.json").write_text(json.dumps(["solo", "duo"]), encoding="utf-8")
|
|
|
|
assert extractor.load_segments("tags.json", "name") == ["solo", "duo"]
|
|
|
|
|
|
def test_universal_json_segment_randomizer_is_seed_deterministic(isolated_gallery_env):
|
|
extractor = isolated_gallery_env.extractor
|
|
library = [
|
|
{"name": "artist_a", "other_names": ["alias_a", "alias_b"]},
|
|
{"name": "artist_b", "other_names": ["alias_c"]},
|
|
]
|
|
(isolated_gallery_env.data_dir / "segments.json").write_text(json.dumps(library), encoding="utf-8")
|
|
|
|
prompt, selected_json = extractor.UniversalJsonSegmentRandomizer().extract_segments(
|
|
"segments.json",
|
|
"other_names",
|
|
2,
|
|
"random",
|
|
"auto",
|
|
"tags",
|
|
"<",
|
|
">",
|
|
"|",
|
|
11,
|
|
"",
|
|
"",
|
|
"contains",
|
|
1.0,
|
|
1.0,
|
|
"{tag}",
|
|
)
|
|
|
|
expected_items = random.Random(11).sample(["alias_a", "alias_b", "alias_c"], 2)
|
|
assert prompt == "|".join(f"<{item}>" for item in expected_items)
|
|
assert json.loads(selected_json) == expected_items
|
|
|
|
|
|
def test_universal_json_segment_randomizer_keeps_legacy_widget_order(isolated_gallery_env):
|
|
extractor = isolated_gallery_env.extractor
|
|
required_keys = list(extractor.UniversalJsonSegmentRandomizer.INPUT_TYPES()["required"])
|
|
|
|
assert required_keys[:13] == [
|
|
"file_name",
|
|
"field_paths",
|
|
"extract_count",
|
|
"mode",
|
|
"duplicate_policy",
|
|
"output_format",
|
|
"prefix",
|
|
"suffix",
|
|
"separator",
|
|
"seed",
|
|
"filter_path",
|
|
"filter_value",
|
|
"filter_mode",
|
|
]
|
|
assert required_keys[-3:] == ["weight_min", "weight_max", "custom_template"]
|
|
|
|
|
|
def test_universal_json_segment_randomizer_unique_only_caps_count(isolated_gallery_env):
|
|
extractor = isolated_gallery_env.extractor
|
|
(isolated_gallery_env.data_dir / "segments.json").write_text(json.dumps([{"tags": ["a", "b"]}]), encoding="utf-8")
|
|
|
|
prompt, selected_json = extractor.UniversalJsonSegmentRandomizer().extract_segments(
|
|
"segments.json",
|
|
"tags",
|
|
5,
|
|
"sequential",
|
|
"unique_only",
|
|
"tags",
|
|
"",
|
|
"",
|
|
",",
|
|
1,
|
|
"",
|
|
"",
|
|
"contains",
|
|
1.0,
|
|
1.0,
|
|
"{tag}",
|
|
)
|
|
|
|
assert prompt == "b,a"
|
|
assert json.loads(selected_json) == ["b", "a"]
|
|
|
|
|
|
def test_universal_json_segment_randomizer_formats_artist_outputs(isolated_gallery_env):
|
|
extractor = isolated_gallery_env.extractor
|
|
library = [{"name": "artist:kiyo_(kiyo_mariari)"}, {"name": "foo_bar"}]
|
|
(isolated_gallery_env.data_dir / "artists.json").write_text(json.dumps(library), encoding="utf-8")
|
|
|
|
node = extractor.UniversalJsonSegmentRandomizer()
|
|
|
|
anima_prompt, _ = node.extract_segments(
|
|
"artists.json",
|
|
"name",
|
|
1,
|
|
"sequential",
|
|
"auto",
|
|
"anima",
|
|
"",
|
|
"",
|
|
", ",
|
|
0,
|
|
"",
|
|
"",
|
|
"contains",
|
|
1.0,
|
|
1.0,
|
|
"{tag}",
|
|
)
|
|
assert anima_prompt == "@kiyo \\(kiyo mariari\\)"
|
|
|
|
artist_prompt, _ = node.extract_segments(
|
|
"artists.json",
|
|
"name",
|
|
1,
|
|
"sequential",
|
|
"auto",
|
|
"artist",
|
|
"",
|
|
"",
|
|
", ",
|
|
1,
|
|
"",
|
|
"",
|
|
"contains",
|
|
1.0,
|
|
1.0,
|
|
"{tag}",
|
|
)
|
|
assert artist_prompt == "artist:foo bar"
|
|
|
|
weighted_prompt, _ = node.extract_segments(
|
|
"artists.json",
|
|
"name",
|
|
1,
|
|
"sequential",
|
|
"auto",
|
|
"weighted_artist",
|
|
"",
|
|
"",
|
|
", ",
|
|
1,
|
|
"",
|
|
"",
|
|
"contains",
|
|
1.2,
|
|
1.2,
|
|
"{tag}",
|
|
)
|
|
assert weighted_prompt == "(foo bar:1.2)"
|
|
|
|
|
|
def test_universal_json_segment_randomizer_custom_template(isolated_gallery_env):
|
|
extractor = isolated_gallery_env.extractor
|
|
(isolated_gallery_env.data_dir / "segments.json").write_text(json.dumps([{"name": "foo_bar"}]), encoding="utf-8")
|
|
|
|
prompt, selected_json = extractor.UniversalJsonSegmentRandomizer().extract_segments(
|
|
"segments.json",
|
|
"name",
|
|
1,
|
|
"sequential",
|
|
"auto",
|
|
"custom",
|
|
"",
|
|
"",
|
|
",",
|
|
0,
|
|
"",
|
|
"",
|
|
"contains",
|
|
1.0,
|
|
1.0,
|
|
"{index}:{clean}:{anima}",
|
|
)
|
|
|
|
assert prompt == "1:foo bar:@foo bar"
|
|
assert json.loads(selected_json) == ["foo_bar"]
|
|
|
|
|
|
def test_universal_json_segment_randomizer_polling_advances_per_node(isolated_gallery_env):
|
|
extractor = isolated_gallery_env.extractor
|
|
extractor.POLLING_STATE.clear()
|
|
(isolated_gallery_env.data_dir / "segments.json").write_text(json.dumps([{"tags": ["a", "b", "c"]}]), encoding="utf-8")
|
|
|
|
node = extractor.UniversalJsonSegmentRandomizer()
|
|
first_prompt, first_json = node.extract_segments(
|
|
"segments.json",
|
|
"tags",
|
|
2,
|
|
"polling",
|
|
"auto",
|
|
"tags",
|
|
"",
|
|
"",
|
|
",",
|
|
0,
|
|
"",
|
|
"",
|
|
"contains",
|
|
1.0,
|
|
1.0,
|
|
"{tag}",
|
|
unique_id="node-1",
|
|
)
|
|
second_prompt, second_json = node.extract_segments(
|
|
"segments.json",
|
|
"tags",
|
|
2,
|
|
"polling",
|
|
"auto",
|
|
"tags",
|
|
"",
|
|
"",
|
|
",",
|
|
0,
|
|
"",
|
|
"",
|
|
"contains",
|
|
1.0,
|
|
1.0,
|
|
"{tag}",
|
|
unique_id="node-1",
|
|
)
|
|
|
|
assert first_prompt == "a,b"
|
|
assert json.loads(first_json) == ["a", "b"]
|
|
assert second_prompt == "c,a"
|
|
assert json.loads(second_json) == ["c", "a"]
|
|
|
|
|
|
def test_old_text_extractor_node_is_not_registered():
|
|
from py import plugin
|
|
|
|
node_classes = plugin.load_node_classes()
|
|
|
|
assert set(node_classes) == {"UniversalJsonSegmentRandomizer", "UniversalPromptSnapshot"}
|
|
assert "UniversalTextExtractor" not in node_classes
|
|
|
|
|
|
def test_runtime_gallery_state_files_are_hidden_from_node_choices(isolated_gallery_env):
|
|
extractor = isolated_gallery_env.extractor
|
|
(isolated_gallery_env.data_dir / "gallery_state.json").write_text("{}", encoding="utf-8")
|
|
(isolated_gallery_env.data_dir / "trash_state.json").write_text("{}", encoding="utf-8")
|
|
(isolated_gallery_env.data_dir / "usable.json").write_text("[]", encoding="utf-8")
|
|
|
|
assert extractor.list_json_files() == ["usable.json"]
|
|
|
|
|
|
def test_node_library_loader_rejects_runtime_and_path_escape_inputs(isolated_gallery_env):
|
|
extractor = isolated_gallery_env.extractor
|
|
(isolated_gallery_env.data_dir / "usable.json").write_text(json.dumps(["ok"]), encoding="utf-8")
|
|
(isolated_gallery_env.data_dir / "gallery_state.json").write_text(json.dumps(["hidden"]), encoding="utf-8")
|
|
|
|
assert extractor.load_entries("usable.json") == ["ok"]
|
|
assert extractor.load_entries("gallery_state.json") == []
|
|
assert extractor.load_entries("../usable.json") == []
|
|
assert extractor.load_entries("../gallery_state.json") == []
|
|
assert extractor.load_entries("not-json.txt") == []
|