Will Miao
6d7cf74405
CivitAI Workflow archives can bundle only a PNG/WebP screenshot whose metadata carries the graph, with no workflow JSON member. Zip candidate discovery now falls back to image members (bounded reads, full-size images first) and parses their metadata in memory, while the archive keeps its original bytes and SHA256 identity for CivitAI matching.
285 lines
9.5 KiB
Python
285 lines
9.5 KiB
Python
"""Tests for workflow document extraction from library artifacts."""
|
|
|
|
import json
|
|
import os
|
|
import zipfile
|
|
|
|
import pytest
|
|
|
|
from py.utils.workflow_extraction import (
|
|
WORKFLOW_DOCUMENT_SUFFIX,
|
|
extract_workflow_document,
|
|
get_workflow_document_path,
|
|
load_artifact_workflow_document,
|
|
materialize_workflow_document,
|
|
)
|
|
|
|
API_WORKFLOW = {
|
|
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "base.safetensors"}},
|
|
"2": {"class_type": "LoraLoader", "inputs": {"lora_name": "lora.safetensors"}},
|
|
}
|
|
|
|
UI_WORKFLOW = {
|
|
"nodes": [{"id": 1, "type": "CheckpointLoaderSimple", "widgets_values": ["base.safetensors"]}],
|
|
"links": [],
|
|
}
|
|
|
|
|
|
def _write_json(path, payload):
|
|
path.write_text(json.dumps(payload), encoding="utf-8")
|
|
|
|
|
|
def _write_zip(path, members):
|
|
with zipfile.ZipFile(path, "w") as archive:
|
|
for name, payload in members.items():
|
|
archive.writestr(name, payload)
|
|
|
|
|
|
def _write_png_with_workflow(path, workflow):
|
|
PIL = pytest.importorskip("PIL.Image")
|
|
from py.utils.exif_utils import ExifUtils
|
|
|
|
img = PIL.new("RGB", (8, 8), color=(0, 0, 0))
|
|
img.save(path, format="PNG")
|
|
ExifUtils.embed_workflow(str(path), json.dumps(workflow))
|
|
|
|
|
|
def _write_webp_with_workflow(path, workflow):
|
|
PIL = pytest.importorskip("PIL.Image")
|
|
from py.utils.exif_utils import ExifUtils
|
|
|
|
img = PIL.new("RGB", (8, 8), color=(0, 0, 0))
|
|
img.save(path, format="WEBP")
|
|
ExifUtils.embed_workflow(str(path), json.dumps(workflow))
|
|
|
|
|
|
def _zip_with_image_member(tmp_path, workflow, extension=".png"):
|
|
"""Build an image carrying *workflow* and bundle it into a zip."""
|
|
|
|
image_path = tmp_path / f"member{extension}"
|
|
if extension == ".png":
|
|
_write_png_with_workflow(image_path, workflow)
|
|
else:
|
|
_write_webp_with_workflow(image_path, workflow)
|
|
artifact = tmp_path / "pack.zip"
|
|
_write_zip(
|
|
artifact,
|
|
{"images/member" + extension: image_path.read_bytes()},
|
|
)
|
|
return artifact
|
|
|
|
|
|
def test_document_path_identity_for_json(tmp_path):
|
|
artifact = tmp_path / "flow.json"
|
|
assert get_workflow_document_path(str(artifact)) == str(artifact)
|
|
|
|
|
|
def test_document_path_derived_for_containers(tmp_path):
|
|
artifact = tmp_path / "flow.zip"
|
|
derived = get_workflow_document_path(str(artifact))
|
|
assert derived == str(tmp_path / f"flow{WORKFLOW_DOCUMENT_SUFFIX}")
|
|
|
|
image = tmp_path / "flow.png"
|
|
assert get_workflow_document_path(str(image)) == str(
|
|
tmp_path / f"flow{WORKFLOW_DOCUMENT_SUFFIX}"
|
|
)
|
|
|
|
|
|
def test_extract_json_artifact(tmp_path):
|
|
artifact = tmp_path / "flow.json"
|
|
_write_json(artifact, API_WORKFLOW)
|
|
document = extract_workflow_document(str(artifact))
|
|
assert document == API_WORKFLOW
|
|
|
|
|
|
def test_extract_rejects_non_workflow_json(tmp_path):
|
|
artifact = tmp_path / "settings.json"
|
|
_write_json(artifact, {"some": "config"})
|
|
assert extract_workflow_document(str(artifact)) is None
|
|
assert materialize_workflow_document(str(artifact)) is None
|
|
|
|
|
|
def test_extract_from_zip_top_level(tmp_path):
|
|
artifact = tmp_path / "pack.zip"
|
|
_write_zip(
|
|
artifact,
|
|
{
|
|
"workflow.json": json.dumps(UI_WORKFLOW),
|
|
"readme.txt": "hello",
|
|
},
|
|
)
|
|
document = extract_workflow_document(str(artifact))
|
|
assert document == UI_WORKFLOW
|
|
|
|
|
|
def test_extract_from_zip_prefers_shallow_candidates(tmp_path):
|
|
artifact = tmp_path / "pack.zip"
|
|
_write_zip(
|
|
artifact,
|
|
{
|
|
"nested/deep/config.json": json.dumps(API_WORKFLOW),
|
|
"top.json": json.dumps(UI_WORKFLOW),
|
|
},
|
|
)
|
|
document = extract_workflow_document(str(artifact))
|
|
assert document == UI_WORKFLOW
|
|
|
|
|
|
def test_extract_from_zip_skips_macosx_and_hidden(tmp_path):
|
|
artifact = tmp_path / "pack.zip"
|
|
_write_zip(
|
|
artifact,
|
|
{
|
|
"__MACOSX/workflow.json": json.dumps(UI_WORKFLOW),
|
|
".hidden/workflow.json": json.dumps(UI_WORKFLOW),
|
|
"real.json": json.dumps(API_WORKFLOW),
|
|
},
|
|
)
|
|
document = extract_workflow_document(str(artifact))
|
|
assert document == API_WORKFLOW
|
|
|
|
|
|
def test_extract_from_zip_without_workflow_returns_none(tmp_path):
|
|
artifact = tmp_path / "images.zip"
|
|
_write_zip(artifact, {"note.txt": "no graphs here", "config.json": "{}"})
|
|
assert extract_workflow_document(str(artifact)) is None
|
|
|
|
|
|
def test_extract_from_zip_png_with_embedded_workflow(tmp_path):
|
|
artifact = _zip_with_image_member(tmp_path, UI_WORKFLOW)
|
|
assert extract_workflow_document(str(artifact)) == UI_WORKFLOW
|
|
|
|
|
|
def test_extract_from_zip_webp_with_embedded_workflow(tmp_path):
|
|
artifact = _zip_with_image_member(tmp_path, UI_WORKFLOW, extension=".webp")
|
|
assert extract_workflow_document(str(artifact)) == UI_WORKFLOW
|
|
|
|
|
|
def test_extract_from_zip_prefers_json_over_image(tmp_path):
|
|
image_path = tmp_path / "shot.png"
|
|
_write_png_with_workflow(image_path, UI_WORKFLOW)
|
|
artifact = tmp_path / "pack.zip"
|
|
_write_zip(
|
|
artifact,
|
|
{
|
|
"workflow.json": json.dumps(API_WORKFLOW),
|
|
"shot.png": image_path.read_bytes(),
|
|
},
|
|
)
|
|
assert extract_workflow_document(str(artifact)) == API_WORKFLOW
|
|
|
|
|
|
def test_extract_from_zip_plain_image_returns_none(tmp_path):
|
|
PIL = pytest.importorskip("PIL.Image")
|
|
image_path = tmp_path / "plain.png"
|
|
PIL.new("RGB", (8, 8)).save(image_path, format="PNG")
|
|
artifact = tmp_path / "images.zip"
|
|
_write_zip(
|
|
artifact,
|
|
{"plain.png": image_path.read_bytes(), "note.txt": "no graphs here"},
|
|
)
|
|
assert extract_workflow_document(str(artifact)) is None
|
|
|
|
|
|
def test_extract_from_zip_continues_after_bad_image_member(tmp_path):
|
|
image_path = tmp_path / "good.png"
|
|
_write_png_with_workflow(image_path, UI_WORKFLOW)
|
|
artifact = tmp_path / "pack.zip"
|
|
# The corrupt member is larger, so the "full-size first" ordering tries
|
|
# it before the valid screenshot; the loop must move on to the next one.
|
|
_write_zip(
|
|
artifact,
|
|
{
|
|
"a_broken.png": b"not an image" * 512,
|
|
"b_good.png": image_path.read_bytes(),
|
|
},
|
|
)
|
|
assert extract_workflow_document(str(artifact)) == UI_WORKFLOW
|
|
|
|
|
|
def test_extract_from_zip_skips_oversized_image_member(tmp_path, monkeypatch):
|
|
from py.utils import workflow_extraction
|
|
|
|
artifact = _zip_with_image_member(tmp_path, UI_WORKFLOW)
|
|
monkeypatch.setattr(workflow_extraction, "_MAX_ZIP_IMAGE_BYTES", 1)
|
|
assert extract_workflow_document(str(artifact)) is None
|
|
|
|
|
|
def test_materialize_from_zip_image_member(tmp_path):
|
|
artifact = _zip_with_image_member(tmp_path, UI_WORKFLOW)
|
|
derived = materialize_workflow_document(str(artifact))
|
|
assert derived == str(tmp_path / f"pack{WORKFLOW_DOCUMENT_SUFFIX}")
|
|
with open(derived, "r", encoding="utf-8") as handle:
|
|
assert json.load(handle) == UI_WORKFLOW
|
|
|
|
|
|
def test_extract_from_corrupt_zip_returns_none(tmp_path):
|
|
artifact = tmp_path / "broken.zip"
|
|
artifact.write_bytes(b"not a zip at all")
|
|
assert extract_workflow_document(str(artifact)) is None
|
|
|
|
|
|
def test_extract_from_png_with_embedded_workflow(tmp_path):
|
|
artifact = tmp_path / "flow.png"
|
|
_write_png_with_workflow(artifact, UI_WORKFLOW)
|
|
document = extract_workflow_document(str(artifact))
|
|
assert document == UI_WORKFLOW
|
|
|
|
|
|
def test_extract_from_webp_with_embedded_workflow(tmp_path):
|
|
artifact = tmp_path / "flow.webp"
|
|
_write_webp_with_workflow(artifact, UI_WORKFLOW)
|
|
document = extract_workflow_document(str(artifact))
|
|
assert document == UI_WORKFLOW
|
|
|
|
|
|
def test_extract_from_plain_image_returns_none(tmp_path):
|
|
PIL = pytest.importorskip("PIL.Image")
|
|
artifact = tmp_path / "plain.png"
|
|
PIL.new("RGB", (8, 8)).save(artifact, format="PNG")
|
|
assert extract_workflow_document(str(artifact)) is None
|
|
assert materialize_workflow_document(str(artifact)) is None
|
|
|
|
|
|
def test_materialize_writes_derived_document_once(tmp_path):
|
|
artifact = tmp_path / "pack.zip"
|
|
_write_zip(artifact, {"workflow.json": json.dumps(API_WORKFLOW)})
|
|
|
|
derived = materialize_workflow_document(str(artifact))
|
|
assert derived == str(tmp_path / f"pack{WORKFLOW_DOCUMENT_SUFFIX}")
|
|
with open(derived, "r", encoding="utf-8") as handle:
|
|
assert json.load(handle) == API_WORKFLOW
|
|
|
|
# A fresh derived file is reused without re-extracting: corrupt the zip
|
|
# and confirm the cached document is still returned.
|
|
artifact_mtime = os.path.getmtime(artifact)
|
|
os.utime(derived, (artifact_mtime + 1, artifact_mtime + 1))
|
|
artifact.write_bytes(b"corrupted")
|
|
assert materialize_workflow_document(str(artifact)) == derived
|
|
|
|
|
|
def test_materialize_reextracts_when_artifact_is_newer(tmp_path):
|
|
artifact = tmp_path / "pack.zip"
|
|
_write_zip(artifact, {"workflow.json": json.dumps(API_WORKFLOW)})
|
|
|
|
derived = materialize_workflow_document(str(artifact))
|
|
|
|
# Replace the artifact with a newer mtime than the derived document.
|
|
_write_zip(artifact, {"workflow.json": json.dumps(UI_WORKFLOW)})
|
|
artifact_mtime = os.path.getmtime(artifact)
|
|
os.utime(artifact, (artifact_mtime + 10, artifact_mtime + 10))
|
|
|
|
assert materialize_workflow_document(str(artifact)) == derived
|
|
with open(derived, "r", encoding="utf-8") as handle:
|
|
assert json.load(handle) == UI_WORKFLOW
|
|
|
|
|
|
def test_load_artifact_workflow_document_round_trip(tmp_path):
|
|
artifact = tmp_path / "pack.zip"
|
|
_write_zip(artifact, {"workflow.json": json.dumps(API_WORKFLOW)})
|
|
assert load_artifact_workflow_document(str(artifact)) == API_WORKFLOW
|
|
|
|
bare = tmp_path / "bare.json"
|
|
_write_json(bare, UI_WORKFLOW)
|
|
assert load_artifact_workflow_document(str(bare)) == UI_WORKFLOW
|