"""Tests for ``readme_processor.py`` — HF README processing for enrich_hf_metadata. Import via ``importlib`` to avoid the ``folder_paths`` dependency in ``py.services.agent.__init__``. """ from __future__ import annotations import importlib.util import re from pathlib import Path import pytest _MODULE_PATH = Path(__file__).parents[2] / "py" / "services" / "agent" / "skills" / "enrich_hf_metadata" / "readme_processor.py" @pytest.fixture(scope="session") def R(): """Load the ``readme_processor`` module once per session.""" spec = importlib.util.spec_from_file_location("readme_processor", str(_MODULE_PATH)) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) return mod # ====================================================================== # extract_gallery_images # ====================================================================== class TestExtractGalleryImages: def test_empty(self, R): assert R.extract_gallery_images("", "repo") == [] assert R.extract_gallery_images("no frontmatter", "repo") == [] def test_no_widget(self, R): readme = "---\ntags: [test]\n---\nbody" assert R.extract_gallery_images(readme, "repo") == [] def test_widget_simple_text(self, R): """YAML ``text: 'plain'`` → extracted as-is.""" readme = """--- widget: - text: 'a cute cat' output: url: images/cat.png ---""" imgs = R.extract_gallery_images(readme, "user/repo") assert len(imgs) == 1 assert imgs[0]["meta"]["prompt"] == "a cute cat" assert "images/cat.png" in imgs[0]["url"] def test_widget_unquoted_text(self, R): """YAML ``text: plain value`` without quotes.""" readme = """--- widget: - text: simple text output: url: img.png ---""" imgs = R.extract_gallery_images(readme, "user/repo") assert len(imgs) == 1 assert imgs[0]["meta"]["prompt"] == "simple text" def test_widget_block_scalar(self, R): """YAML ``text: >-`` folded block scalar — extract actual content.""" readme = """--- widget: - text: >- Long toons, a close-up of a cartoon characters face is featured in a vibrant red backdrop. output: url: images/LT4.png ---""" imgs = R.extract_gallery_images(readme, "user/repo") assert len(imgs) == 1 prompt = imgs[0]["meta"]["prompt"] assert "Long toons" in prompt assert "vibrant red backdrop" in prompt assert prompt != ">-" def test_widget_dash_prefix_output(self, R): """YAML ``- output:`` (dash prefix) — regression for widget parsing.""" readme = """--- widget: - output: url: images/test.png text: dash test ---""" imgs = R.extract_gallery_images(readme, "user/repo") assert len(imgs) == 1 assert imgs[0]["meta"]["prompt"] == "dash test" assert "images/test.png" in imgs[0]["url"] def test_widget_mixed_entries(self, R): """Multiple widget entries with different text styles.""" readme = """--- widget: - text: >- First entry description. output: url: img1.png - text: second entry output: url: img2.png - text: 'third entry' output: url: img3.png ---""" imgs = R.extract_gallery_images(readme, "user/repo") assert len(imgs) == 3 assert imgs[0]["meta"]["prompt"] == "First entry description." assert imgs[1]["meta"]["prompt"] == "second entry" assert imgs[2]["meta"]["prompt"] == "third entry" # ====================================================================== # extract_simple_markdown_images # ====================================================================== class TestExtractSimpleMarkdownImages: def test_empty(self, R): assert R.extract_simple_markdown_images("", "repo") == [] def test_basic_markdown_image(self, R): """``![alt](./img.png)`` → absolute URL.""" imgs = R.extract_simple_markdown_images("![test](./image_0.png)", "u/r") assert len(imgs) == 1 assert "image_0.png" in imgs[0]["url"] assert imgs[0]["meta"]["prompt"] == "test" def test_absolute_url(self, R): """``![alt](https://...)`` → keep as-is.""" imgs = R.extract_simple_markdown_images( "![img](https://example.com/img.png)", "u/r" ) assert len(imgs) == 1 assert imgs[0]["url"] == "https://example.com/img.png" def test_skips_code_fences(self, R): """Inside ``` blocks should be ignored.""" text = """outside ``` ![inside](./img.png) ``` outside again ![valid](./valid.png)""" imgs = R.extract_simple_markdown_images(text, "u/r") assert len(imgs) == 1 assert "valid.png" in imgs[0]["url"] def test_deduplicates(self, R): text = "![a](./img.png)\n![b](./img.png)" imgs = R.extract_simple_markdown_images(text, "u/r") assert len(imgs) == 1 # deduplicated # ====================================================================== # extract_html_img_tags # ====================================================================== class TestExtractHtmlImgTags: def test_double_quoted_src(self, R): imgs = R.extract_html_img_tags('', "u/r") assert len(imgs) == 1 assert "img.png" in imgs[0]["url"] def test_single_quoted_src(self, R): imgs = R.extract_html_img_tags("", "u/r") assert len(imgs) == 1 assert "img.png" in imgs[0]["url"] def test_absolute_url(self, R): imgs = R.extract_html_img_tags( '', "u/r" ) assert len(imgs) == 1 assert imgs[0]["url"] == "https://cdn.example.com/img.png" def test_deduplicates_across_formats(self, R): text = '\n' imgs = R.extract_html_img_tags(text, "u/r") assert len(imgs) == 1 # ====================================================================== # extract_gallery_table_images # ====================================================================== class TestExtractGalleryTableImages: def test_gallery_table(self, R): text = """| Preview | Prompt | |--------|--------| | ![img](./a.png) | a cat | | ![img](./b.png) | a dog |""" imgs = R.extract_gallery_table_images(text, "u/r") assert len(imgs) == 2 assert imgs[0]["meta"]["prompt"] == "a cat" assert "a.png" in imgs[0]["url"] assert imgs[1]["meta"]["prompt"] == "a dog" def test_skips_non_gallery_table(self, R): text = """| Parameter | Value | |----------|-------| | Steps | 4 |""" imgs = R.extract_gallery_table_images(text, "u/r") assert len(imgs) == 0 # ====================================================================== # clean_readme_for_llm + strip helpers # ====================================================================== class TestCleanReadmeForLlm: def test_preserves_plain_code_block(self, R): """`` ``` `` without language tag → preserved (trigger words).""" text = """Before ``` pixel art sprite, game asset ``` After""" cleaned = R.clean_readme_for_llm(text) assert "pixel art sprite" in cleaned assert "game asset" in cleaned def test_strips_fenced_code_with_lang(self, R): """`` ```python `` → stripped.""" text = "before\n```python\nimport torch\n```\nafter" cleaned = R.clean_readme_for_llm(text) assert "import torch" not in cleaned assert "before" in cleaned assert "after" in cleaned def test_preserves_markdown_image_url(self, R): """``![alt](url)`` → URL kept for LLM preview extraction.""" text = "![sample](./preview.png)" cleaned = R.clean_readme_for_llm(text) assert "./preview.png" in cleaned def test_converts_html_img_tag_to_markdown_image(self, R): """```` → ``![](src)`` preserving URL for LLM.""" text = 'before\n\nafter' cleaned = R.clean_readme_for_llm(text) assert "![](logo.png)" in cleaned assert "logo.png" in cleaned # URL preserved for LLM extraction def test_widget_stripped_frontmatter_preserved(self, R): """Widget YAML stripped but ``base_model:`` kept.""" text = """--- tags: [test] widget: - text: >- long description here output: url: img.png base_model: black-forest-labs/FLUX.1-dev instance_prompt: test ---""" cleaned = R.clean_readme_for_llm(text) assert "widget:" not in cleaned assert "black-forest-labs/FLUX.1-dev" in cleaned assert "instance_prompt: test" in cleaned def test_training_table_stripped(self, R): """Training-parameter table → stripped.""" text = """before | LR Scheduler | constant | |--------------|---------| | Optimizer | AdamW | after""" cleaned = R.clean_readme_for_llm(text) assert "LR Scheduler" not in cleaned assert "Optimizer" not in cleaned assert "before" in cleaned assert "after" in cleaned def test_best_dimensions_table_kept(self, R): """Non-training table (Best Dimensions) → kept.""" text = """## Best Dimensions - 768 x 1024 (Best) - 1024 x 1024 (Default)""" cleaned = R.clean_readme_for_llm(text) assert "768 x 1024" in cleaned def test_boilerplate_section_stripped(self, R): text = """stuff ## Download model [link](url) ## Next section content""" cleaned = R.clean_readme_for_llm(text) assert "Download model" not in cleaned assert "Next section" in cleaned assert "content" in cleaned def test_returns_empty_for_none(self, R): assert R.clean_readme_for_llm(None) == "" def test_returns_empty_for_empty(self, R): assert R.clean_readme_for_llm("") == "" # ====================================================================== # _is_heading / _heading_level # ====================================================================== class TestHeadingDetection: @pytest.mark.parametrize( "line,expected", [ ("# Title", 1), ("## Sub", 2), ("### Subsub", 3), ("#### Subsubsub", 4), ("

Title

", 1), ("

Sub

", 2), ("

Sub

", 3), ("

Sub

", 4), ("not a heading", 0), ("###", 0), # no text after ### ("", 0), # closing tag, not a heading ("", 0), ], ) def test_heading_level(self, R, line, expected): assert R._heading_level(line) == expected @pytest.mark.parametrize( "line,expected", [ ("# Title", True), ("

Sub

", True), ("", False), # closing tag ("not heading", False), ], ) def test_is_heading(self, R, line, expected): assert R._is_heading(line) == expected # ====================================================================== # extract_relevant_section # ====================================================================== class TestExtractRelevantSection: def test_fallback_full_readme(self, R): """No match → full README returned.""" readme = "# Title\n\nsome content" assert R.extract_relevant_section(readme, "nonexistent") == readme def test_empty_basename_returns_full(self, R): readme = "# Title" assert R.extract_relevant_section(readme, "") == readme def test_match_heading_includes_yaml(self, R): """Matching heading should still include YAML frontmatter.""" readme = """--- base_model: foo --- # My-Model-Title content ## Subsection more""" section = R.extract_relevant_section(readme, "My-Model") assert "base_model: foo" in section assert "content" in section assert "Subsection" in section def test_match_heading_includes_subheadings(self, R): """``# Title`` match includes all ``##`` children.""" readme = """# Main Title ## Child A content A ## Child B content B ## Child C content C""" section = R.extract_relevant_section(readme, "Main Title") assert "Child A" in section assert "Child B" in section assert "Child C" in section def test_match_download_link(self, R): """Download link containing basename → section extracted.""" readme = """# Collection ## Model A [Download](./model_a.safetensors) ## MyModel [Download](./mymodel.safetensors) content here ## Model B other""" section = R.extract_relevant_section(readme, "mymodel") assert "content here" in section assert "Model A" not in section # should not include sibling def test_heading_closing_tag_not_boundary(self, R): """```` should NOT be treated as a section boundary.""" readme = """# Title

some text

## Real Section content""" section = R.extract_relevant_section(readme, "Title") assert "Real Section" in section # forward walk should not stop at assert "content" in section # ====================================================================== # _extract_frontmatter # ====================================================================== class TestExtractFrontmatter: def test_basic(self, R): assert R._extract_frontmatter("---\ntags: [a]\n---\nbody") == "\ntags: [a]\n" def test_no_frontmatter(self, R): assert R._extract_frontmatter("no dashes") == "" def test_empty_string(self, R): assert R._extract_frontmatter("") == "" # ====================================================================== # _strip_widget_section # ====================================================================== class TestStripWidgetSection: def test_strip_widget_keep_base_model(self, R): """Widget stripped but ``base_model:`` preserved.""" text = """--- tags: [test] widget: - text: >- long text output: url: img.png base_model: black-forest-labs/FLUX.1-dev ---""" result = R._strip_widget_section(text) assert "widget:" not in result assert "black-forest-labs/FLUX.1-dev" in result def test_no_widget_no_change(self, R): text = "---\ntags: [a]\n---" assert R._strip_widget_section(text) == text def test_widget_at_end_of_frontmatter(self, R): """Widget is the last YAML key before closing ---.""" text = """--- base_model: a widget: - text: x output: url: y.png ---""" result = R._strip_widget_section(text) assert "widget:" not in result assert "base_model: a" in result # ====================================================================== # _strip_fenced_code_blocks # ====================================================================== class TestStripFencedCodeBlocks: def test_strips_with_language(self, R): text = "a\n```python\ncode\n```\nb" assert R._strip_fenced_code_blocks(text) == "a\nb" def test_keeps_plain_fence(self, R): """`` ``` `` without language → preserved.""" text = "a\n```\ntrigger words\n```\nb" assert "trigger words" in R._strip_fenced_code_blocks(text) def test_pattern(self, R): text = "x\n```yaml\nkey: val\n```\ny" assert "key: val" not in R._strip_fenced_code_blocks(text)