from __future__ import annotations import logging from pathlib import Path import subprocess import sys from PIL import Image, ImageChops import pytest from rich.console import Console import kneiff.utils as ut import kneiff.utils.config as config_utils import kneiff.utils.env as env_utils import kneiff.utils.image.caption.io as caption_io import kneiff.utils.image.caption.tags as caption_tags import kneiff.utils.image.caption.text as caption_text import kneiff.utils.image.discovery as image_discovery import kneiff.utils.image.encoding as image_encoding import kneiff.utils.image.grid as image_grid import kneiff.utils.image.sidecars as image_sidecars import kneiff.utils.logging as logging_utils import kneiff.utils.tablefmt as tablefmt_utils import kneiff.utils.text as text_utils import kneiff.utils.yaml as yaml_utils def test_utils_facade_reexports_shared_helpers() -> None: """Keep the public utility facade aligned with the shared helper modules.""" assert ut.read_dotenv is env_utils.read_dotenv assert ut.read_json_mapping is config_utils.read_json_mapping assert ut.read_yaml_document is config_utils.read_yaml_document assert ut.dump_yaml is yaml_utils.dump_yaml assert ut.export_yaml is yaml_utils.export_yaml assert ut.load_yaml is yaml_utils.load_yaml assert ut.YamlError is yaml_utils.YamlError assert ut.optional_text is config_utils.optional_text assert ut.text_tuple is config_utils.text_tuple assert ut.format_text_table is tablefmt_utils.format_text_table assert ut.split_comma_separated_values is text_utils.split_comma_separated_values def test_utils_facade_keeps_rich_lazy_for_non_cli_imports() -> None: """Protect Space imports that only need config helpers from CLI-only Rich.""" script = ( "import sys\n" "import kneiff.utils\n" "raise SystemExit(1 if 'rich' in sys.modules else 0)\n" ) result = subprocess.run( [sys.executable, "-c", script], capture_output=True, text=True, check=False, ) assert result.returncode == 0, result.stdout + result.stderr def test_stdlib_deep_mapping_helpers() -> None: """Cover nested mapping helpers exposed through the utility facade.""" data: dict[object, object] = {"a": {"b": 1}} assert ut.deep_get(data, ("a", "b")) == 1 assert ut.deep_get(data, ("a", "missing"), default="fallback") == "fallback" assert ut.deep_get(data, (), default="fallback") == "fallback" ut.deep_set(data, ("a", "c"), 2) assert data == {"a": {"b": 1, "c": 2}} assert ut.deep_right_merge( {"a": {"b": 1}, "keep": True}, {"a": {"c": 2}}, ) == {"a": {"b": 1, "c": 2}, "keep": True} def test_read_dotenv_parses_exports_quotes_and_comments(tmp_path: Path) -> None: env_path = tmp_path / ".env.shared" env_path.write_text( """ # ignored export KNF_WORKERS=' 4 ' EMPTY= MODEL_NAME = "local model" invalid line """, encoding="utf-8", ) values = env_utils.read_dotenv(env_path) assert values == { "KNF_WORKERS": "4", "EMPTY": "", "MODEL_NAME": "local model", } def test_read_dotenv_value_returns_non_empty_value(tmp_path: Path) -> None: env_path = tmp_path / ".env.shared" env_path.write_text("KNF_WORKERS= 4 \nEMPTY=\n", encoding="utf-8") assert env_utils.read_dotenv_value(env_path, "KNF_WORKERS") == "4" assert env_utils.read_dotenv_value(env_path, "EMPTY") is None assert env_utils.read_dotenv_value(env_path, "MISSING") is None assert env_utils.read_dotenv_value(tmp_path / "missing.env", "KNF_WORKERS") is None with pytest.raises(RuntimeError, match=r"\.env file not found"): env_utils.read_dotenv_value( tmp_path / "missing.env", "KNF_WORKERS", require_file=True, ) def test_first_env_value_prefers_process_env( monkeypatch: pytest.MonkeyPatch, ) -> None: file_values = {"KNF_WORKERS": "2", "ALT_WORKERS": "3"} monkeypatch.delenv("KNF_WORKERS", raising=False) monkeypatch.delenv("ALT_WORKERS", raising=False) assert ( env_utils.first_env_value( ["KNF_WORKERS", "ALT_WORKERS"], dotenv_values=file_values ) == "2" ) monkeypatch.setenv("KNF_WORKERS", "5") assert ( env_utils.first_env_value( ["KNF_WORKERS", "ALT_WORKERS"], dotenv_values=file_values ) == "5" ) def test_parse_positive_int_defaults_and_rejects_invalid_values() -> None: assert env_utils.parse_positive_int(None, name="KNF_WORKERS", source="env") == 1 assert env_utils.parse_positive_int(" 3 ", name="KNF_WORKERS", source="env") == 3 with pytest.raises(ValueError, match="KNF_WORKERS"): env_utils.parse_positive_int("0", name="KNF_WORKERS", source="env") def test_config_helpers_load_yaml_and_validate_scalar_nodes(tmp_path: Path) -> None: config_path = tmp_path / "config.yaml" config_path.write_text("enabled: true\ncount: 3\n", encoding="utf-8") document = config_utils.read_yaml_document( config_path, dependency_error="py-yaml12 is required.", ) mapping = config_utils.require_mapping( document, message="config must be a mapping." ) assert config_utils.require_bool(mapping["enabled"], name="enabled") is True assert config_utils.require_positive_int(mapping["count"], name="count") == 3 assert config_utils.require_exact_text("subset", name="subset") == "subset" assert config_utils.optional_text(" caption ", name="caption") == "caption" assert config_utils.text_tuple( [" tag ", "", "tag", "other"], name="tags", ) == ("tag", "other") assert config_utils.require_int(5, name="seed") == 5 assert ( config_utils.optional_mapping(None, message="optional must be a mapping.") is None ) with pytest.raises(ValueError, match="enabled must be true or false"): config_utils.require_bool("yes", name="enabled") with pytest.raises(ValueError, match="count must be a positive integer"): config_utils.require_positive_int(True, name="count") with pytest.raises(ValueError, match="seed must be an integer"): config_utils.require_int("5", name="seed") with pytest.raises(ValueError, match="subset must be a non-empty exact string"): config_utils.require_exact_text(" subset ", name="subset") def test_yaml_helpers_export_readable_multiline_strings(tmp_path: Path) -> None: output_path = tmp_path / "nested" / "manifest.yaml" yaml_utils.export_yaml( { "caption": "line 1\nline 2\n", "path": Path("SOURCE/scene.png"), "tags": {"beta", "alpha"}, }, output_path, ) text = output_path.read_text(encoding="utf-8") assert "caption: |2" in text assert "line 1" in text assert "line 2" in text assert yaml_utils.load_yaml(output_path) == { "caption": "line 1\nline 2\n", "path": "SOURCE/scene.png", "tags": ["alpha", "beta"], } def test_yaml_helpers_load_empty_text_as_none() -> None: assert yaml_utils.load_yaml_text("") is None def test_read_yaml_document_rejects_duplicate_top_level_keys(tmp_path: Path) -> None: config_path = tmp_path / "config.yaml" config_path.write_text("enabled: true\nenabled: false\n", encoding="utf-8") with pytest.raises(ValueError, match="Duplicate YAML key 'enabled'"): config_utils.read_yaml_document( config_path, dependency_error="py-yaml12 is required.", ) def test_read_yaml_document_rejects_duplicate_nested_keys(tmp_path: Path) -> None: config_path = tmp_path / "config.yaml" config_path.write_text( """ training: simpletuner: trainer: caption_dropout_probability: 0.0 caption_dropout_probability: 0.05 """, encoding="utf-8", ) with pytest.raises( ValueError, match="Duplicate YAML key 'caption_dropout_probability'", ): config_utils.read_yaml_document( config_path, dependency_error="py-yaml12 is required.", ) def test_config_helpers_load_json_file_shapes(tmp_path: Path) -> None: mapping_path = tmp_path / "mapping.json" sequence_path = tmp_path / "sequence.json" mapping_path.write_text('{"alpha": 1}', encoding="utf-8") sequence_path.write_text('[{"id": "a"}, {"id": "b"}]', encoding="utf-8") assert config_utils.read_json_mapping(mapping_path) == {"alpha": 1} assert config_utils.read_json_sequence(sequence_path) == [ {"id": "a"}, {"id": "b"}, ] assert config_utils.read_json_mapping_sequence(sequence_path) == [ {"id": "a"}, {"id": "b"}, ] def test_config_helpers_reject_unexpected_json_shapes(tmp_path: Path) -> None: mapping_path = tmp_path / "mapping.json" sequence_path = tmp_path / "sequence.json" invalid_path = tmp_path / "invalid.json" mapping_path.write_text('{"alpha": 1}', encoding="utf-8") sequence_path.write_text('["not-an-object"]', encoding="utf-8") invalid_path.write_text("{", encoding="utf-8") with pytest.raises(ValueError, match="Expected JSON array"): config_utils.read_json_sequence(mapping_path) with pytest.raises(ValueError, match="Expected JSON object entries"): config_utils.read_json_mapping_sequence(sequence_path) with pytest.raises(ValueError, match="Broken JSON"): config_utils.read_json_document( invalid_path, decode_error_message="Broken JSON", ) def test_tablefmt_styles_plain_text_table_for_console() -> None: rendered = tablefmt_utils.format_text_table( headers=("Subset", "Images", "Train %"), rows=[("fullbody", 30, "100.0%")], right_align_headers={"Images", "Train %"}, ) styled = tablefmt_utils.style_text_table_for_console( rendered, column_styles={"Subset": "bold", "Train %": "green"}, ) assert styled.plain == rendered assert styled.spans plain_console = Console( force_terminal=True, color_system="truecolor", width=120, highlight=False, ) with plain_console.capture() as capture: plain_console.print(rendered) plain_output = capture.get() styled_console = Console( force_terminal=True, color_system="truecolor", width=120, highlight=False, ) with styled_console.capture() as capture: styled_console.print(styled) styled_output = capture.get() assert "\x1b[" not in plain_output assert "\x1b[" in styled_output assert "fullbody" in styled_output def test_tablefmt_styles_exact_cell_values_for_console() -> None: rendered = tablefmt_utils.format_text_table( headers=("Status", "Images"), rows=[("KEPT", 2), ("KICKED", 1)], right_align_headers={"Images"}, ) styled = tablefmt_utils.style_text_table_for_console( rendered, column_styles={"Status": "green"}, cell_styles={"KICKED": "bold bright_red"}, ) assert styled.plain == rendered kicked_start = rendered.index("KICKED") kicked_spans = [ span for span in styled.spans if span.start <= kicked_start and span.end >= kicked_start + len("KICKED") ] assert any("bright_red" in str(span.style) for span in kicked_spans) def test_image_grid_writes_grouped_contact_sheet_with_labels(tmp_path: Path) -> None: red_path = tmp_path / "red.png" blue_path = tmp_path / "blue.png" output_path = tmp_path / "grid.png" Image.new("RGB", (12, 8), color="red").save(red_path) Image.new("RGB", (8, 12), color="blue").save(blue_path) written_path = image_grid.write_grouped_contact_sheet( [ image_grid.ContactSheetGroup( label="fullbody", images=( image_grid.GridImage(red_path), image_grid.GridImage(blue_path), ), ) ], output_path, thumbnail_size=16, columns=2, ) assert written_path == output_path with Image.open(output_path) as rendered: assert rendered.size[0] >= 16 * 2 assert rendered.size[1] > 16 header = rendered.crop((0, 0, rendered.width, 16)) blank_header = Image.new("RGB", header.size, color="white") assert ImageChops.difference(header, blank_header).getbbox() is not None def test_image_grid_wrap_text_block_caps_lines_with_ellipsis() -> None: font = image_grid._load_font(24) wrapped = image_grid._wrap_text_block( ( "alpha\n" "This validation prompt text is long enough to wrap across many lines " "and should end with an ellipsis once the limit is reached." ), font=font, max_width=120, max_lines=5, ) assert len(wrapped) == 5 assert wrapped[0] == "alpha" assert wrapped[-1].endswith("...") def test_image_grid_prompt_label_limit_allows_ten_wrapped_lines() -> None: font = image_grid._load_font(24) wrapped = image_grid._wrap_text_block( " ".join("word" for _ in range(100)), font=font, max_width=100, max_lines=image_grid.DEFAULT_PROMPT_LABEL_MAX_LINES, ) assert image_grid.DEFAULT_PROMPT_LABEL_MAX_LINES == 10 assert len(wrapped) == 10 assert wrapped[-1].endswith("...") def test_image_grid_writes_matrix_contact_sheet_with_footer_labels( tmp_path: Path, ) -> None: image_path = tmp_path / "white.png" output_path = tmp_path / "matrix-grid.png" Image.new("RGB", (16, 16), color="white").save(image_path) written_path = image_grid.write_matrix_contact_sheet( row_labels=("step 1",), column_labels=("alpha_prompt",), cells={("step 1", "alpha_prompt"): image_grid.GridImage(image_path)}, output_path=output_path, thumbnail_size=32, column_display_text={ "alpha_prompt": ( "alpha_prompt\n" "A long validation prompt that wraps across several lines so the " "footer label region is visible in the rendered sheet." ) }, repeat_column_labels_at_bottom=True, column_label_max_lines=5, ) assert written_path == output_path with Image.open(output_path) as rendered: top = rendered.crop((120, 0, rendered.width, min(100, rendered.height))) bottom = rendered.crop( (120, max(0, rendered.height - 100), rendered.width, rendered.height) ) blank_top = Image.new("RGB", top.size, color="white") blank_bottom = Image.new("RGB", bottom.size, color="white") assert ImageChops.difference(top, blank_top).getbbox() is not None assert ImageChops.difference(bottom, blank_bottom).getbbox() is not None def test_split_comma_separated_values_ignores_empty_tokens() -> None: assert text_utils.split_comma_separated_values(" alpha, beta ,, gamma ") == [ "alpha", "beta", "gamma", ] def test_iter_image_files_returns_sorted_supported_images(tmp_path: Path) -> None: (tmp_path / "b.JPG").write_bytes(b"image") (tmp_path / "a.png").write_bytes(b"image") (tmp_path / "notes.txt").write_text("skip", encoding="utf-8") nested = tmp_path / "nested" nested.mkdir() (nested / "c.webp").write_bytes(b"image") assert image_discovery.iter_image_files(tmp_path, recursive=False) == [ tmp_path / "a.png", tmp_path / "b.JPG", ] assert image_discovery.iter_image_files(tmp_path) == [ tmp_path / "a.png", tmp_path / "b.JPG", nested / "c.webp", ] def test_comma_sidecars_read_and_write_without_overwrite(tmp_path: Path) -> None: image_path = tmp_path / "scene.png" image_path.write_bytes(b"image") sidecar = image_path.with_suffix(".txt") assert image_sidecars.read_comma_sidecar(image_path, ".txt") == [] sidecar.write_text(" alpha, beta , , gamma\n", encoding="utf-8") assert image_sidecars.read_comma_sidecar(image_path, ".txt") == [ "alpha", "beta", "gamma", ] image_sidecars.write_text_sidecar(sidecar, "replacement", overwrite=False) assert sidecar.read_text(encoding="utf-8") == " alpha, beta , , gamma\n" image_sidecars.write_text_sidecar(sidecar, "replacement", overwrite=True) assert sidecar.read_text(encoding="utf-8") == "replacement\n" def test_caption_sidecar_io_builds_output_paths() -> None: io = caption_io.CaptionSidecarIO(out_suffix=".cap.txt") assert io.output_path(Path("frame.png")) == Path("frame.cap.txt") def test_caption_tag_filtering_can_match_blip_aliases() -> None: tags = [" alpha_tag ", "Beta", "alpha tag", "gamma"] assert caption_tags.filter_caption_tags( tags, ignore_tags=["beta"], replace_underscores=True, dedupe=True, ) == ["alpha_tag", "gamma"] assert caption_tags.filter_caption_tags( tags, ignore_tags=["beta"], ) == ["alpha_tag", "alpha tag", "gamma"] def test_caption_text_rules_preserve_backend_specific_ban_behavior() -> None: raw_caption = '"masterpiece Rook_Kaefer standing in best quality light"' blip_rules = caption_text.build_caption_text_rules( must_include=["felkin"], ban_phrases=None, max_words=6, strip_banned=True, ) server_rules = caption_text.build_caption_text_rules( must_include=["felkin"], ban_phrases=None, max_words=6, strip_banned=False, ) assert blip_rules.apply(raw_caption) == "felkin Rook_Kaefer standing in light" assert ( server_rules.apply(raw_caption) == "felkin masterpiece Rook_Kaefer standing in best" ) def test_image_to_jpeg_data_url_downscales_and_encodes(tmp_path: Path) -> None: image_path = tmp_path / "source.png" Image.new("RGB", (12, 6), color=(20, 30, 40)).save(image_path) assert ( image_encoding.image_to_jpeg_data_url(image_path, max_side=4).split(",", 1)[0] == "data:image/jpeg;base64" ) def test_get_logger_builds_named_logger() -> None: assert logging_utils.get_logger("kneiff.tests").name == "kneiff.tests" def test_configure_logging_sets_root_level() -> None: root_logger = logging.getLogger() original_level = root_logger.level try: logging_utils.configure_logging(level="DEBUG") assert root_logger.level == logging.DEBUG finally: root_logger.setLevel(original_level) def test_configure_logging_rejects_unknown_level() -> None: with pytest.raises(ValueError, match="Unknown logging level"): logging_utils.configure_logging(level="verbose-ish")