Files
ComfyUI-Lora-Manager/tests/services/test_post_processor.py
T
Will Miao 5ab0e88abc feat(links): support ModelScope and TensorArt as model sources
A model file could only ever be linked to huggingface.co: `set_hf_url`
validated the URL with a huggingface-only regex, the agent fetched the card
from a hardcoded HF URL, and the readme processor built every relative image
path off `https://huggingface.co/{repo}/resolve/main`. ModelScope publishes the
same model-card convention (README.md + YAML frontmatter, often carrying
`base_model:` and `trigger_words:`) behind a public, key-less API, so the
enrichment pipeline could already serve it - it was the plumbing that was
HF-shaped, not the idea.

Make the external source a first-class, provider-driven concept:

- New `py/services/model_sources/` registry. A `ModelSource` owns URL
  recognition (lenient for stored values, strict for user input), the
  canonical page URL, model-card fetching, the asset base URL and the
  capability flags. `HuggingFaceSource` is the previous logic relocated;
  `ModelScopeSource` reads `/models/{o}/{n}/resolve/{master|main}/README.md`
  and falls back to `/api/v1/models/{o}/{n}/repo`. `TensorArtSource` is
  link-only on purpose: tensor.art answers plain HTTP clients with a
  Cloudflare challenge and its internal API (ap-east-1.tensorart.cloud /
  cn.tensorart.net) rejects every /v1/model/* route with "invalid
  authorization header", so it declares supports_enrichment=False rather than
  failing silently later.
- Metadata gains `source_platform` + `source_url`; `hf_url` stays as a
  read/write alias, written only for Hugging Face, so existing sidecars,
  cached rows and third-party consumers keep working. Normalisation runs at
  the scanner, the persistent cache (both directions, plus two new columns
  behind an ALTER migration) and the linking handler - which is what stops a
  user who switches sources from leaving a stale `hf_url` on a ModelScope
  model.
- The agent pipeline keys off the provider instead of `hf_url`: the fast-fail
  gate now explains *why* a model is skipped (no source / unknown source /
  source without a reachable card), the prompt context exposes
  source_url/source_id/source_label/asset_base_url while still filling the
  legacy hf_url/repo aliases, and the four README image extractors take a
  base_url (defaulting to HF) so relative paths resolve against the right
  site. Version grouping generalises to hf: / ms: / ta: keys.
- `POST /api/lm/set-hf-url` keeps its path and its legacy payload keys but
  accepts `source_url`, validates against every provider and returns the
  platform. `GET /api/lm/model-sources` lets the UI render the supported-site
  list from the server.
- Frontend: a `modelSourceHelpers` mirror of the registry drives the link
  dialog, the card/modal globe (branded "View on ModelScope/TensorArt"), the
  version-group key and the enrichment gate; the versions tab no longer sends
  ms:/ta: keys to the CivitAI API.

TensorArt stays in the list because provenance is worth keeping even when the
card is unreadable - the dialog says so plainly ("Sites that don't expose one
(currently TensorArt) can only be linked") and the context menu disables
enrichment with a matching tooltip, instead of the user getting
"Unsupported URL".

Verified against the real ModelScope API: jj3550945163/Krea-2-LORA returns a
1882-byte card whose frontmatter carries base_model/tags/trigger_words, and
relative images resolve to .../resolve/master/....

Tests: backend 2815 passed; frontend 1130 JS + 91 Vue passed; pytest
tests/i18n and a Jinja compile pass over templates/. The nine locales carry
[TODO: Translate] for the new strings, completed in the next commit.
2026-09-14 07:24:08 +08:00

526 lines
21 KiB
Python

"""Tests for the PostProcessor (py/services/agent/post_processor.py).
PostProcessor delegates all I/O to AgentCLI — these tests mock AgentCLI
functions and verify the business logic (conditions, merges, dispatch).
"""
from __future__ import annotations
from datetime import datetime, timezone
from unittest import mock
import pytest
from py.services.agent.post_processor import PostProcessor
@pytest.fixture
def processor():
return PostProcessor()
# ======================================================================
# process() — routing
# ======================================================================
class TestProcessDispatch:
@pytest.mark.asyncio
async def test_unknown_skill_returns_error(self, processor):
result = await processor.process(
skill_name="nonexistent",
model_path="/p.safetensors",
llm_output={},
metadata={},
)
assert result["success"] is False
assert "nonexistent" in result["errors"][0]
@pytest.mark.asyncio
async def test_enrich_hf_metadata_routes_correctly(self, processor):
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview") as mock_dl,
mock.patch("py.metadata_ops.refresh_cache") as mock_ref,
):
mock_apply.return_value = ["metadata_source"]
mock_dl.return_value = None
result = await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output={},
metadata={"from_civitai": True},
)
assert result["success"] is True
# ======================================================================
# enrich_hf_metadata — field-level logic
# ======================================================================
class TestEnrichHfMetadata:
"""Business logic tests for the enrich_hf_metadata post-processor."""
MIN_LLM_OUTPUT = {
"base_model": "",
"trigger_words": [],
"short_description": "",
"tags": [],
"recommended_width": 0,
"recommended_height": 0,
"preview_url": "",
"confidence": "low",
}
# -- base_model ------------------------------------------------------
@pytest.mark.asyncio
async def test_base_model_overwrites_empty(self, processor):
"""Empty current base_model → new value is applied."""
llm = {**self.MIN_LLM_OUTPUT, "base_model": "Flux.1 D"}
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=False),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=llm,
metadata={"base_model": ""},
)
applied = mock_apply.call_args[0][1]
assert applied["base_model"] == "Flux.1 D"
@pytest.mark.asyncio
async def test_base_model_does_not_overwrite_existing_civitai(self, processor):
"""Existing base_model from CivitAI → not overwritten."""
llm = {**self.MIN_LLM_OUTPUT, "base_model": "Flux.1 D"}
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=False),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=llm,
metadata={"base_model": "SDXL 1.0", "from_civitai": True},
)
# apply IS called (metadata_source, llm_enriched_at) but base_model not in it
applied = mock_apply.call_args[0][1]
assert "base_model" not in applied
@pytest.mark.asyncio
async def test_base_model_overwrites_existing_hf_model(self, processor):
"""Existing base_model from HF → overwritten (LLM is more reliable)."""
llm = {**self.MIN_LLM_OUTPUT, "base_model": "Flux.1 D"}
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=False),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=llm,
metadata={
"base_model": "SD 1.5",
"from_civitai": False,
"hf_url": "https://huggingface.co/user/repo",
},
)
applied = mock_apply.call_args[0][1]
assert applied["base_model"] == "Flux.1 D"
@pytest.mark.asyncio
async def test_base_model_skipped_when_llm_empty(self, processor):
"""LLM returns empty base_model → nothing written."""
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=False),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=self.MIN_LLM_OUTPUT,
metadata={"base_model": ""},
)
applied = mock_apply.call_args[0][1]
assert "base_model" not in applied
# -- trigger_words ---------------------------------------------------
@pytest.mark.asyncio
async def test_trigger_words_merged(self, processor):
"""New trigger words written when current list is empty."""
llm = {**self.MIN_LLM_OUTPUT, "trigger_words": ["trigger1", "trigger2"]}
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=None),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=llm,
metadata={},
)
applied = mock_apply.call_args[0][1]
assert applied["civitai"]["trainedWords"] == ["trigger1", "trigger2"]
# -- short_description → civitai.description -------------------------
@pytest.mark.asyncio
async def test_short_description_written_to_civitai(self, processor):
"""short_description written to civitai.description for HF models."""
llm = {**self.MIN_LLM_OUTPUT, "short_description": "A short summary"}
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=None),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=llm,
metadata={
"from_civitai": False,
"hf_url": "https://huggingface.co/user/repo",
},
)
applied = mock_apply.call_args[0][1]
assert applied["civitai"]["description"] == "A short summary"
@pytest.mark.asyncio
async def test_short_description_skipped_without_hf_url(self, processor):
"""short_description NOT written when the model has no HF source."""
llm = {**self.MIN_LLM_OUTPUT, "short_description": "A short summary"}
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=None),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=llm,
metadata={"from_civitai": True},
)
applied = mock_apply.call_args[0][1]
assert "civitai" not in applied or "description" not in applied.get("civitai", {})
# -- readme_content → modelDescription -------------------------------
@pytest.mark.asyncio
async def test_readme_content_converted_to_model_description(self, processor):
"""Raw README converted to HTML and stored as modelDescription."""
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=None),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=self.MIN_LLM_OUTPUT,
metadata={
"from_civitai": False,
"hf_url": "https://huggingface.co/user/repo",
},
readme_content="# Hello\n\nThis is **bold**.",
)
applied = mock_apply.call_args[0][1]
assert "<h1>Hello</h1>" in applied.get("modelDescription", "")
assert "<strong>bold</strong>" in applied.get("modelDescription", "")
@pytest.mark.asyncio
async def test_readme_content_skipped_without_hf_url(self, processor):
"""README content NOT converted when the model has no HF source."""
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=None),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=self.MIN_LLM_OUTPUT,
metadata={"from_civitai": True},
readme_content="# Hello",
)
applied = mock_apply.call_args[0][1]
assert "modelDescription" not in applied
# -- gallery images → civitai.images ---------------------------------
@pytest.mark.asyncio
async def test_gallery_images_extracted_from_readme(self, processor):
"""Widget entries in README → civitai.images."""
readme = """---
widget:
- text: "a cat"
output:
url: images/cat.png
---
Content
"""
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=None),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=self.MIN_LLM_OUTPUT,
metadata={
"from_civitai": False,
"hf_url": "https://huggingface.co/user/repo",
},
readme_content=readme,
)
applied = mock_apply.call_args[0][1]
images = applied.get("civitai", {}).get("images", [])
assert len(images) == 1
assert images[0]["url"] == (
"https://huggingface.co/user/repo/resolve/main/images/cat.png"
)
assert images[0]["meta"]["prompt"] == "a cat"
@pytest.mark.asyncio
async def test_gallery_images_use_modelscope_asset_base_url(self, processor):
"""A ModelScope-linked model resolves relative images against ModelScope."""
readme = """---
widget:
- text: "a cat"
output:
url: images/cat.png
---
Content
"""
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=None),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=self.MIN_LLM_OUTPUT,
metadata={
"from_civitai": False,
"source_platform": "modelscope",
"source_url": "https://modelscope.cn/models/user/repo",
},
readme_content=readme,
)
applied = mock_apply.call_args[0][1]
images = applied.get("civitai", {}).get("images", [])
assert len(images) == 1
assert images[0]["url"] == (
"https://modelscope.cn/models/user/repo/resolve/master/images/cat.png"
)
@pytest.mark.asyncio
async def test_base_model_overwrites_existing_modelscope_model(self, processor):
"""ModelScope is an external source, so the LLM may overwrite base_model."""
llm = {**self.MIN_LLM_OUTPUT, "base_model": "Flux.1 D"}
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=False),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=llm,
metadata={
"base_model": "SD 1.5",
"source_platform": "modelscope",
"source_url": "https://modelscope.cn/models/user/repo",
},
)
assert mock_apply.call_args[0][1]["base_model"] == "Flux.1 D"
@pytest.mark.asyncio
async def test_gallery_images_skipped_without_hf_url(self, processor):
"""Gallery images NOT extracted when the model has no HF source."""
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=None),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=self.MIN_LLM_OUTPUT,
metadata={"from_civitai": True},
readme_content="---\nwidget:\n- text: a\n output:\n url: x.png\n---\n",
)
applied = mock_apply.call_args[0][1]
civitai = applied.get("civitai", {})
assert "images" not in civitai
@pytest.mark.asyncio
async def test_gallery_images_extracted_for_civitai_linked_model(self, processor):
"""A model may be on CivitAI and HuggingFace at once (#1094).
HF enrichment is gated on ``hf_url``, not on ``from_civitai``, so the
README gallery is still applied when both sources are present.
"""
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=None),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=self.MIN_LLM_OUTPUT,
metadata={
"from_civitai": True,
"hf_url": "https://huggingface.co/user/repo",
},
readme_content="---\nwidget:\n- text: a\n output:\n url: x.png\n---\n",
)
applied = mock_apply.call_args[0][1]
images = applied.get("civitai", {}).get("images", [])
assert len(images) == 1
assert images[0]["url"] == (
"https://huggingface.co/user/repo/resolve/main/x.png"
)
# -- tags ------------------------------------------------------------
@pytest.mark.asyncio
async def test_tags_merged_and_deduplicated(self, processor):
llm = {**self.MIN_LLM_OUTPUT, "tags": ["flux", "lora", "STYLE"]}
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=False),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=llm,
metadata={"tags": ["anime"], "from_civitai": False},
)
merged = mock_apply.call_args[0][1]["tags"]
assert "anime" in merged
assert "flux" in merged
assert "style" in merged # lowercased
# "lora" and "STYLE" → "lora" and "style"
assert len(merged) == 4 # anime, flux, lora, style
# -- metadata_source & llm_enriched_at --------------------------------
@pytest.mark.asyncio
async def test_audit_fields_always_set(self, processor):
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview", return_value=False),
mock.patch("py.metadata_ops.refresh_cache"),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=self.MIN_LLM_OUTPUT,
metadata={},
)
applied = mock_apply.call_args[0][1]
assert applied["metadata_source"] == "agent:enrich_hf_metadata"
assert "llm_enriched_at" in applied
# -- preview download ------------------------------------------------
@pytest.mark.asyncio
async def test_preview_downloaded_when_url_provided(self, processor):
llm = {**self.MIN_LLM_OUTPUT, "preview_url": "https://ex.com/img.png"}
with (
mock.patch("py.metadata_ops.apply_metadata_updates") as mock_apply,
mock.patch("py.metadata_ops.download_preview") as mock_dl,
mock.patch("py.metadata_ops.refresh_cache"),
):
mock_dl.return_value = "/p.webp"
result = await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=llm,
metadata={},
)
assert result["preview_downloaded"] is True
mock_dl.assert_awaited_once_with("/p.safetensors", "https://ex.com/img.png")
applied = mock_apply.call_args[0][1]
assert applied["preview_url"] == "/p.webp"
@pytest.mark.asyncio
async def test_preview_skipped_when_exists(self, processor):
"""If current_preview file exists on disk, skip download."""
llm = {**self.MIN_LLM_OUTPUT, "preview_url": "https://ex.com/img.png"}
with (
mock.patch("py.metadata_ops.apply_metadata_updates"),
mock.patch("py.metadata_ops.download_preview") as mock_dl,
mock.patch("py.metadata_ops.refresh_cache"),
mock.patch("os.path.exists", return_value=True),
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=llm,
metadata={"preview_url": "/existing/preview.webp"},
)
mock_dl.assert_not_called()
# -- cache refresh ---------------------------------------------------
@pytest.mark.asyncio
async def test_cache_refreshed_when_updates_applied(self, processor):
llm = {**self.MIN_LLM_OUTPUT, "base_model": "Flux.1 D"}
with (
mock.patch("py.metadata_ops.apply_metadata_updates", return_value=["base_model"]),
mock.patch("py.metadata_ops.download_preview", return_value=False),
mock.patch("py.metadata_ops.refresh_cache") as mock_ref,
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=llm,
metadata={"base_model": ""},
)
mock_ref.assert_awaited_once_with("/p.safetensors")
@pytest.mark.asyncio
async def test_cache_not_refreshed_when_nothing_changed(self, processor):
with (
mock.patch("py.metadata_ops.apply_metadata_updates", return_value=[]),
mock.patch("py.metadata_ops.download_preview", return_value=False),
mock.patch("py.metadata_ops.refresh_cache") as mock_ref,
):
await processor.process(
skill_name="enrich_hf_metadata",
model_path="/p.safetensors",
llm_output=self.MIN_LLM_OUTPUT,
metadata={"base_model": ""},
)
mock_ref.assert_not_called()
# ======================================================================
# Unit: _merge_tags
# ======================================================================
class TestMergeTags:
def test_deduplicates_case_insensitive(self):
existing = ["anime", "Flux"]
new = ["flux", "LORA", "anime"]
result = PostProcessor._merge_tags(existing, new)
# All tags are lowercased (matching TagUpdateService behaviour)
assert result == ["anime", "flux", "lora"]