106 lines
3.1 KiB
Python
106 lines
3.1 KiB
Python
"""Tests for core.metadata module."""
|
|
|
|
import json
|
|
import pytest
|
|
from PIL.PngImagePlugin import PngInfo
|
|
from core.metadata import (
|
|
build_iteration_metadata,
|
|
embed_in_pnginfo,
|
|
format_a1111_params,
|
|
extract_iteration_metadata,
|
|
)
|
|
|
|
|
|
class TestBuildIterationMetadata:
|
|
def test_basic(self):
|
|
meta = build_iteration_metadata(
|
|
batch_name="test_batch",
|
|
batch_id="test_batch_abc123",
|
|
index=0,
|
|
total=12,
|
|
params={"steps": 20, "cfg": 7.5},
|
|
)
|
|
assert meta["batch_name"] == "test_batch"
|
|
assert meta["batch_id"] == "test_batch_abc123"
|
|
assert meta["index"] == 0
|
|
assert meta["total"] == 12
|
|
assert meta["params"]["steps"] == 20
|
|
assert meta["params"]["cfg"] == 7.5
|
|
assert meta["mode"] == "matrix"
|
|
|
|
def test_linear_mode(self):
|
|
meta = build_iteration_metadata(
|
|
batch_name="test",
|
|
batch_id="test_123",
|
|
index=0,
|
|
total=1,
|
|
params={},
|
|
mode="linear",
|
|
)
|
|
assert meta["mode"] == "linear"
|
|
|
|
|
|
class TestEmbedInPnginfo:
|
|
def test_embeds_iteration_metadata(self):
|
|
pnginfo = PngInfo()
|
|
meta = build_iteration_metadata(
|
|
"batch", "batch_123", 0, 1, {"steps": 20}
|
|
)
|
|
embed_in_pnginfo(pnginfo, meta)
|
|
|
|
# PngInfo doesn't have a direct read API for testing,
|
|
# so we check it doesn't error
|
|
assert pnginfo is not None
|
|
|
|
def test_embeds_workflow_and_prompt(self):
|
|
pnginfo = PngInfo()
|
|
meta = build_iteration_metadata(
|
|
"batch", "batch_123", 0, 1, {"steps": 20}
|
|
)
|
|
workflow = {"nodes": []}
|
|
prompt = {"1": {"inputs": {}}}
|
|
embed_in_pnginfo(pnginfo, meta, workflow=workflow, prompt=prompt)
|
|
assert pnginfo is not None
|
|
|
|
|
|
class TestFormatA1111Params:
|
|
def test_basic(self):
|
|
result = format_a1111_params({"steps": 20, "cfg": 7.5})
|
|
assert "Steps: 20" in result
|
|
assert "CFG scale: 7.5" in result
|
|
|
|
def test_sampler(self):
|
|
result = format_a1111_params({"sampler": "euler"})
|
|
assert "Sampler: euler" in result
|
|
|
|
def test_unknown_param(self):
|
|
result = format_a1111_params({"my_custom_param": 42})
|
|
assert "My Custom Param: 42" in result
|
|
|
|
def test_empty(self):
|
|
assert format_a1111_params({}) == ""
|
|
|
|
def test_seed(self):
|
|
result = format_a1111_params({"seed": 12345})
|
|
assert "Seed: 12345" in result
|
|
|
|
|
|
class TestExtractIterationMetadata:
|
|
def test_from_dict(self):
|
|
data = {"wi_iteration": {"batch_name": "test", "index": 0}}
|
|
result = extract_iteration_metadata(data)
|
|
assert result["batch_name"] == "test"
|
|
|
|
def test_from_json_string(self):
|
|
data = {
|
|
"wi_iteration": json.dumps({"batch_name": "test", "index": 0})
|
|
}
|
|
result = extract_iteration_metadata(data)
|
|
assert result["batch_name"] == "test"
|
|
|
|
def test_missing(self):
|
|
assert extract_iteration_metadata({}) is None
|
|
|
|
def test_invalid_json(self):
|
|
assert extract_iteration_metadata({"wi_iteration": "not json"}) is None
|