feat: update tests for workflow and prompt validation

This commit is contained in:
LuqP2
2026-01-09 14:34:46 -03:00
parent 12b7efc19e
commit da047e8ff0
2 changed files with 32 additions and 0 deletions
+9
View File
@@ -0,0 +1,9 @@
__pycache__/__init__.cpython-310.pyc
__pycache__/conftest.cpython-310-pytest-9.0.2.pyc
__pycache__/metadata_utils.cpython-310.pyc
__pycache__/metahub_save_node.cpython-310.pyc
__pycache__/timer_node.cpython-310.pyc
__pycache__/workflow_extractor.cpython-310.pyc
tests/__pycache__/conftest.cpython-310-pytest-9.0.2.pyc
tests/__pycache__/test_metadata_utils.cpython-310-pytest-9.0.2.pyc
tests/__pycache__/test_workflow_extractor.cpython-310-pytest-9.0.2.pyc
+23
View File
@@ -1,4 +1,5 @@
import json
from pathlib import Path
import copy
from metadata_utils import (
@@ -12,6 +13,24 @@ from metadata_utils import (
from PIL import Image
def _read_png_text_chunks(path: Path) -> dict:
data = path.read_bytes()
if data[:8] != b"\x89PNG\r\n\x1a\n":
raise AssertionError("Not a PNG file")
offset = 8
chunks = {}
while offset + 8 <= len(data):
length = int.from_bytes(data[offset:offset + 4], "big")
chunk_type = data[offset + 4:offset + 8]
chunk_data = data[offset + 8:offset + 8 + length]
offset += 12 + length
if chunk_type == b"tEXt":
if b"\x00" in chunk_data:
keyword, text = chunk_data.split(b"\x00", 1)
chunks[keyword.decode("latin-1")] = text.decode("latin-1")
return chunks
def test_ensure_prompt_in_workflow_adds_prompt():
workflow = {"workflow": {"nodes": []}}
prompt = {"1": {"class_type": "KSampler", "inputs": {}}}
@@ -162,3 +181,7 @@ def test_save_png_with_metadata_writes_workflow_prompt(tmp_path):
assert prompt_text
assert json.loads(workflow_text) == imh_metadata["workflow"]
assert json.loads(prompt_text) == imh_metadata["prompt_api"]
chunk_text = _read_png_text_chunks(file_path)
assert "workflow" in chunk_text
assert "prompt" in chunk_text