Ewald
e3afb1a0d3
- New POST /mjr/am/collect-files route (CSRF, write token, rate limit, path allowlist)
- collect_service: extracts workflow media/model refs, resolves inputs within allowed roots, builds {stem}_collected.zip next to the asset with writable-dir fallback to output/_mjr_collected
- Zip contains asset, workflow.json, prompt_graph.json, prompt.json (positive/negative traced via geninfo parser), inputs/ and collected_files.txt manifest
- comfy_core adapter: get_model_full_path helper
- Frontend: 'Collect files' grid context menu item + sidebar (D) MButton, collectFiles() API client op
- Tests: tests/features/test_collect_service.py (10 tests)
211 lines
8.5 KiB
Python
211 lines
8.5 KiB
Python
"""Tests for the collect-files service (workflow inputs bundling)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import zipfile
|
|
from pathlib import Path
|
|
|
|
from mjr_am_backend.features.assets import collect_service as cs
|
|
from mjr_am_backend.shared import Result
|
|
|
|
|
|
def _patch_roots(monkeypatch, *, output_root: Path, input_root: Path) -> None:
|
|
monkeypatch.setattr(cs, "get_runtime_output_root", lambda: str(output_root))
|
|
monkeypatch.setattr(cs, "get_input_directory", lambda: str(input_root))
|
|
monkeypatch.setattr(cs, "get_temp_directory", lambda: None)
|
|
monkeypatch.setattr(cs, "list_custom_roots", lambda: Result.Ok([]))
|
|
monkeypatch.setattr(cs, "get_model_full_path", lambda name: None)
|
|
|
|
|
|
def _sd_prompt_graph() -> dict:
|
|
"""Minimal valid API prompt graph: text encoders -> KSampler -> SaveImage."""
|
|
return {
|
|
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "sd.safetensors"}},
|
|
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "a beautiful castle", "clip": ["1", 1]}},
|
|
"3": {"class_type": "CLIPTextEncode", "inputs": {"text": "blurry, ugly", "clip": ["1", 1]}},
|
|
"4": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"model": ["1", 0],
|
|
"positive": ["2", 0],
|
|
"negative": ["3", 0],
|
|
"latent_image": ["5", 0],
|
|
"seed": 1,
|
|
"steps": 20,
|
|
"cfg": 7.5,
|
|
"sampler_name": "euler",
|
|
"scheduler": "normal",
|
|
"denoise": 1.0,
|
|
},
|
|
},
|
|
"5": {"class_type": "EmptyLatentImage", "inputs": {"width": 512, "height": 512, "batch_size": 1}},
|
|
"6": {"class_type": "VAEDecode", "inputs": {"samples": ["4", 0], "vae": ["1", 2]}},
|
|
"7": {"class_type": "SaveImage", "inputs": {"images": ["6", 0], "filename_prefix": "out"}},
|
|
}
|
|
|
|
|
|
def test_extract_prompt_texts_traces_text_encoders():
|
|
payload = cs.extract_prompt_texts(None, _sd_prompt_graph())
|
|
assert payload is not None
|
|
assert payload["positive"] == "a beautiful castle"
|
|
assert payload["negative"] == "blurry, ugly"
|
|
|
|
|
|
def test_extract_prompt_texts_returns_none_without_prompts():
|
|
prompt = {"1": {"class_type": "LoadImage", "inputs": {"image": "photo.png"}}}
|
|
assert cs.extract_prompt_texts(None, prompt) is None
|
|
|
|
|
|
def test_extract_workflow_refs_media_and_models():
|
|
prompt = {
|
|
"1": {"class_type": "LoadImage", "inputs": {"image": "photo.png [input]"}},
|
|
"2": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "sdxl.safetensors"}},
|
|
"3": {"class_type": "KSampler", "inputs": {"seed": 42, "model": ["2", 0]}},
|
|
}
|
|
workflow = {
|
|
"nodes": [
|
|
{"type": "LoadImage", "widgets_values": ["photo.png [input]", "image"]},
|
|
{"type": "VHS_LoadVideo", "widgets_values": {"video": "clips/intro.mp4"}},
|
|
{"type": "LoraLoader", "widgets_values": ["style.safetensors", 1.0]},
|
|
]
|
|
}
|
|
media, models = cs.extract_workflow_refs(workflow, prompt)
|
|
assert "photo.png [input]" in media
|
|
assert "clips/intro.mp4" in media
|
|
# Deduped: photo.png appears in prompt and workflow but only once.
|
|
assert len([m for m in media if "photo.png" in m]) == 1
|
|
assert "sdxl.safetensors" in models
|
|
assert "style.safetensors" in models
|
|
|
|
|
|
def test_split_annotation():
|
|
assert cs._split_annotation("img.png [input]") == ("img.png", "input")
|
|
assert cs._split_annotation("sub/img.png [output]") == ("sub/img.png", "output")
|
|
assert cs._split_annotation("plain.png") == ("plain.png", "")
|
|
assert cs._split_annotation("weird [brackets].png") == ("weird [brackets].png", "")
|
|
|
|
|
|
def test_resolve_media_refs_relative_and_missing(monkeypatch, tmp_path: Path):
|
|
input_root = tmp_path / "input"
|
|
output_root = tmp_path / "output"
|
|
(input_root / "clips").mkdir(parents=True)
|
|
output_root.mkdir()
|
|
(input_root / "photo.png").write_bytes(b"png")
|
|
(input_root / "clips" / "intro.mp4").write_bytes(b"mp4")
|
|
_patch_roots(monkeypatch, output_root=output_root, input_root=input_root)
|
|
|
|
entries = cs.resolve_media_refs(["photo.png [input]", "clips/intro.mp4", "gone.png"])
|
|
by_name = {e["name"]: e for e in entries}
|
|
assert by_name["photo.png"]["status"] == "ok"
|
|
assert by_name["intro.mp4"]["status"] == "ok"
|
|
assert by_name["gone.png"]["status"] == "missing"
|
|
|
|
|
|
def test_resolve_media_refs_absolute_outside_roots_not_copied(monkeypatch, tmp_path: Path):
|
|
input_root = tmp_path / "input"
|
|
output_root = tmp_path / "output"
|
|
outside = tmp_path / "outside"
|
|
for d in (input_root, output_root, outside):
|
|
d.mkdir()
|
|
secret = outside / "secret.png"
|
|
secret.write_bytes(b"png")
|
|
_patch_roots(monkeypatch, output_root=output_root, input_root=input_root)
|
|
|
|
entries = cs.resolve_media_refs([str(secret)])
|
|
assert entries[0]["status"] == "skipped_outside_roots"
|
|
|
|
|
|
def test_build_collect_zip_end_to_end(monkeypatch, tmp_path: Path):
|
|
input_root = tmp_path / "input"
|
|
output_root = tmp_path / "output"
|
|
input_root.mkdir()
|
|
output_root.mkdir()
|
|
asset = output_root / "render.png"
|
|
asset.write_bytes(b"fake png")
|
|
(input_root / "photo.png").write_bytes(b"input bytes")
|
|
_patch_roots(monkeypatch, output_root=output_root, input_root=input_root)
|
|
|
|
workflow = {"nodes": [{"type": "LoadImage", "widgets_values": ["photo.png [input]"]}]}
|
|
prompt = _sd_prompt_graph()
|
|
prompt["0"] = {"class_type": "LoadImage", "inputs": {"image": "photo.png [input]"}}
|
|
|
|
res = cs.build_collect_zip(asset, workflow=workflow, prompt=prompt)
|
|
assert res.ok, res.error
|
|
data = res.data or {}
|
|
zip_path = Path(data["zip_path"])
|
|
assert zip_path.exists()
|
|
assert zip_path.parent == output_root # next to the asset
|
|
assert data["fallback_used"] is False
|
|
assert data["missing"] == []
|
|
assert data["has_prompt_text"] is True
|
|
|
|
with zipfile.ZipFile(zip_path) as zf:
|
|
names = set(zf.namelist())
|
|
assert "render.png" in names
|
|
assert "workflow.json" in names
|
|
assert "prompt_graph.json" in names
|
|
assert "prompt.json" in names
|
|
assert "inputs/photo.png" in names
|
|
assert "collected_files.txt" in names
|
|
manifest = zf.read("collected_files.txt").decode("utf-8")
|
|
assert "photo.png" in manifest
|
|
assert "copied" in manifest
|
|
parsed = json.loads(zf.read("workflow.json").decode("utf-8"))
|
|
assert parsed == workflow
|
|
graph = json.loads(zf.read("prompt_graph.json").decode("utf-8"))
|
|
assert graph == prompt
|
|
texts = json.loads(zf.read("prompt.json").decode("utf-8"))
|
|
assert texts["positive"] == "a beautiful castle"
|
|
assert texts["negative"] == "blurry, ugly"
|
|
|
|
|
|
def test_build_collect_zip_missing_inputs_listed(monkeypatch, tmp_path: Path):
|
|
input_root = tmp_path / "input"
|
|
output_root = tmp_path / "output"
|
|
input_root.mkdir()
|
|
output_root.mkdir()
|
|
asset = output_root / "render.png"
|
|
asset.write_bytes(b"fake png")
|
|
_patch_roots(monkeypatch, output_root=output_root, input_root=input_root)
|
|
|
|
workflow = {"nodes": [{"type": "LoadImage", "widgets_values": ["gone.png [input]"]}]}
|
|
res = cs.build_collect_zip(asset, workflow=workflow, prompt=None)
|
|
assert res.ok, res.error
|
|
data = res.data or {}
|
|
assert data["missing"] == ["gone.png"]
|
|
with zipfile.ZipFile(Path(data["zip_path"])) as zf:
|
|
manifest = zf.read("collected_files.txt").decode("utf-8")
|
|
assert "MISSING" in manifest
|
|
|
|
|
|
def test_build_collect_zip_unique_name(monkeypatch, tmp_path: Path):
|
|
input_root = tmp_path / "input"
|
|
output_root = tmp_path / "output"
|
|
input_root.mkdir()
|
|
output_root.mkdir()
|
|
asset = output_root / "render.png"
|
|
asset.write_bytes(b"fake png")
|
|
(output_root / "render_collected.zip").write_bytes(b"existing")
|
|
_patch_roots(monkeypatch, output_root=output_root, input_root=input_root)
|
|
|
|
res = cs.build_collect_zip(asset, workflow=None, prompt=None)
|
|
assert res.ok, res.error
|
|
assert Path((res.data or {})["zip_path"]).name == "render_collected (2).zip"
|
|
|
|
|
|
def test_writable_dest_dir_fallback(monkeypatch, tmp_path: Path):
|
|
output_root = tmp_path / "output"
|
|
output_root.mkdir()
|
|
monkeypatch.setattr(cs, "get_runtime_output_root", lambda: str(output_root))
|
|
|
|
missing_dir = tmp_path / "does-not-exist"
|
|
dest, fallback_used = cs._writable_dest_dir(missing_dir)
|
|
assert fallback_used is True
|
|
assert dest == output_root / cs._FALLBACK_SUBDIR
|
|
assert dest.is_dir()
|
|
|
|
dest2, fallback2 = cs._writable_dest_dir(output_root)
|
|
assert fallback2 is False
|
|
assert dest2 == output_root
|