309 lines
11 KiB
Python
309 lines
11 KiB
Python
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from mjr_am_backend.adapters.db.migrations import MigrationRunner
|
|
from mjr_am_backend.adapters.db.migrations.registry import MIGRATIONS
|
|
from mjr_am_backend.adapters.db.schema import migrate_schema
|
|
from mjr_am_backend.adapters.db.sqlite import Sqlite
|
|
from mjr_am_backend.features.index import metadata_helpers as mh
|
|
from mjr_am_backend.shared import Result
|
|
|
|
|
|
class _DummyLock:
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
|
|
class _DB:
|
|
def __init__(self):
|
|
self.exec_calls = []
|
|
self.query_calls = []
|
|
self.query_result = Result.Ok([])
|
|
|
|
async def aexecute(self, sql, params=()):
|
|
self.exec_calls.append((sql, params))
|
|
return Result.Ok({"ok": True})
|
|
|
|
async def aquery(self, sql, params=()):
|
|
self.query_calls.append((sql, params))
|
|
return self.query_result
|
|
|
|
def lock_for_asset(self, _asset_id):
|
|
return _DummyLock()
|
|
|
|
|
|
async def _make_db(tmp_path: Path) -> Sqlite:
|
|
db_path = tmp_path / "metadata-helpers.db"
|
|
db = Sqlite(str(db_path), attach={"vec": str(tmp_path / "vectors.sqlite")})
|
|
mig = await migrate_schema(db)
|
|
assert mig.ok
|
|
runner_res = await MigrationRunner(MIGRATIONS).run(db)
|
|
assert runner_res.ok, runner_res.error
|
|
return db
|
|
|
|
|
|
def test_json_guard_and_tag_helpers(monkeypatch):
|
|
monkeypatch.setattr(mh, "MAX_METADATA_JSON_BYTES", 20)
|
|
raw = json.dumps({"x": "y" * 100})
|
|
truncated = mh._apply_metadata_json_size_guard(1, Result.Ok({"a": 1}), raw, filepath="x.png")
|
|
assert '"_truncated"' in truncated
|
|
|
|
assert mh._sanitize_metadata_tags(["a", "a", " ", "b"]) == ["a", "b"]
|
|
r, tags_json, tags_txt = mh._extract_rating_and_tags(Result.Ok({"rating": 9, "tags": ["x", "y"]}))
|
|
assert r == 5 and tags_json == '["x", "y"]' and tags_txt == "x y"
|
|
|
|
|
|
def test_json_guard_preserves_priority_fields_when_truncated(monkeypatch):
|
|
monkeypatch.setattr(mh, "MAX_METADATA_JSON_BYTES", 220)
|
|
raw = json.dumps(
|
|
{
|
|
"workflow": {"nodes": [{"id": i, "type": "KSampler"} for i in range(50)]},
|
|
"prompt": {"1": {"inputs": {"text": "cinematic portrait lighting"}}},
|
|
"parameters": "cinematic portrait lighting\nNegative prompt: blurry\nSampler: Euler a\nModel: dreamshaperXL",
|
|
"geninfo": {
|
|
"engine": {"type": "TXT2IMG"},
|
|
"sampler": {"name": "Euler a"},
|
|
"models": {"checkpoint": {"name": "dreamshaperXL"}},
|
|
},
|
|
"workflow_type": "txt2img",
|
|
}
|
|
)
|
|
|
|
truncated = mh._apply_metadata_json_size_guard(1, Result.Ok({"a": 1}), raw, filepath="x.png")
|
|
payload = json.loads(truncated)
|
|
|
|
assert payload["_truncated"] is True
|
|
assert payload["workflow_type"] == "TXT2IMG"
|
|
assert payload["model"] == "dreamshaperXL"
|
|
assert payload["sampler"] == "Euler a"
|
|
assert "cinematic portrait lighting" in payload["prompt"]
|
|
|
|
|
|
def test_json_guard_uses_minimal_fallback_when_budget_is_tiny(monkeypatch):
|
|
monkeypatch.setattr(mh, "MAX_METADATA_JSON_BYTES", 40)
|
|
raw = json.dumps({"parameters": "x" * 400, "model": "model-a", "sampler": "Euler"})
|
|
|
|
truncated = mh._apply_metadata_json_size_guard(1, Result.Ok({"a": 1}), raw, filepath="x.png")
|
|
payload = json.loads(truncated)
|
|
|
|
assert payload["_truncated"] is True
|
|
assert payload["original_bytes"] > 40
|
|
assert "model" not in payload
|
|
|
|
|
|
def test_geninfo_extras_and_presence_flags():
|
|
meta = {
|
|
"parameters": "x",
|
|
"geninfo": {
|
|
"models": {"a": {"name": "m1"}},
|
|
"loras": [{"name": "l1"}],
|
|
"positive": {"value": "prompt"},
|
|
"engine": {"type": "tts"},
|
|
"inputs": [{"filename": "i.wav"}],
|
|
},
|
|
}
|
|
extras = mh._collect_geninfo_extras(meta)
|
|
assert "m1" in extras and "l1" in extras and "i.wav" in extras
|
|
enriched = mh._enrich_tags_text_with_metadata(Result.Ok(meta), "tag1")
|
|
assert "tag1" in enriched and "prompt" in enriched
|
|
|
|
has_wf, has_gen = mh._metadata_presence_flags({"parameters": "Steps:20"})
|
|
assert has_wf is True and has_gen is True
|
|
assert mh._graph_has_sampler({"1": {"class_type": "KSampler", "inputs": {"steps": 20}}}) is True
|
|
|
|
|
|
def test_generation_time_parser_accepts_units_and_nested_variants():
|
|
assert mh._parse_generation_time_ms("8421ms") == 8421
|
|
assert mh._parse_generation_time_ms("8.421s") == 8421
|
|
assert mh._parse_generation_time_ms("1.5m") == 90_000
|
|
assert mh._parse_generation_time_ms("1.5h") == 5_400_000
|
|
assert mh._parse_generation_time_ms(True) is None
|
|
assert mh._parse_generation_time_ms("2026:04:10 12:00:00") is None
|
|
|
|
workflow_type, generation_time_ms, positive_prompt = mh._denormalized_metadata_fields(
|
|
Result.Ok(
|
|
{
|
|
"quality": "partial",
|
|
"geninfo": {"generation_time_ms": "8.421s"},
|
|
"positive_prompt": "portrait",
|
|
"workflow_type": "txt2img",
|
|
}
|
|
)
|
|
)
|
|
assert workflow_type == "TXT2IMG"
|
|
assert generation_time_ms == 8421
|
|
assert positive_prompt == "portrait"
|
|
|
|
|
|
def test_denormalized_prompt_prefers_geninfo_positive_over_file_path_fallbacks():
|
|
_, _, positive_prompt = mh._denormalized_metadata_fields(
|
|
Result.Ok(
|
|
{
|
|
"prompt": r"C:\ComfyUI\output\image_00001.png",
|
|
"parameters": r"C:\ComfyUI\output\image_00001.png",
|
|
"positive_prompt": "stale prompt",
|
|
"geninfo": {"positive": {"value": "cinematic portrait"}},
|
|
}
|
|
)
|
|
)
|
|
|
|
assert positive_prompt == "cinematic portrait"
|
|
|
|
|
|
def test_denormalized_ltx_director_prompt_uses_timeline_global_prompt():
|
|
prompt_graph = {
|
|
"131": {
|
|
"class_type": "LTXDirector",
|
|
"inputs": {
|
|
"timeline_data": json.dumps(
|
|
{
|
|
"global_prompt": "golden portrait, cinematic sunset light",
|
|
"segments": [{"prompt": "camera turns"}],
|
|
}
|
|
)
|
|
},
|
|
},
|
|
"132": {"class_type": "SaveVideo", "inputs": {"video": ["131", 0]}},
|
|
}
|
|
_, _, positive_prompt = mh._denormalized_metadata_fields(
|
|
Result.Ok(
|
|
{
|
|
"prompt": r"d:\____comfy_outputs\video\ltx_director_00009_.mp4",
|
|
"raw_ffprobe": {"format": {"tags": {"prompt": json.dumps(prompt_graph)}}},
|
|
}
|
|
)
|
|
)
|
|
|
|
assert positive_prompt == "golden portrait, cinematic sunset light"
|
|
|
|
|
|
def test_prepare_and_error_payload_helpers():
|
|
has_wf, has_gen, quality, raw = mh.MetadataHelpers.prepare_metadata_fields(Result.Ok({"quality": "partial", "parameters": "x"}))
|
|
assert has_wf is True and has_gen is True and quality == "partial" and raw
|
|
assert mh.MetadataHelpers._bool_to_db(True) == 1
|
|
assert mh.MetadataHelpers._bool_to_db(False) == 0
|
|
assert mh.MetadataHelpers._bool_to_db(None) is None
|
|
|
|
err = Result.Err("E", "bad", quality="degraded", foo="x")
|
|
payload = mh.MetadataHelpers.metadata_error_payload(err, "a.png")
|
|
assert payload.ok and payload.data["quality"] == "degraded" and payload.data["foo"] == "x"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_write_row_retrieve_cache_and_store(monkeypatch):
|
|
db = _DB()
|
|
monkeypatch.setattr(mh, "MAX_METADATA_JSON_BYTES", 0)
|
|
meta = Result.Ok({"quality": "full", "rating": 4, "tags": ["a"], "parameters": "x"})
|
|
out = await mh.MetadataHelpers.write_asset_metadata_row(db, 10, meta, filepath="a.png")
|
|
assert out.ok and db.exec_calls
|
|
|
|
db.query_result = Result.Ok([{"metadata_raw": json.dumps({"x": 1})}])
|
|
cached = await mh.MetadataHelpers.retrieve_cached_metadata(db, "a.png", "h")
|
|
assert cached and cached.ok and cached.data["x"] == 1
|
|
|
|
store = await mh.MetadataHelpers.store_metadata_cache(db, "a.png", "h", Result.Ok({"x": 1}))
|
|
assert store.ok
|
|
|
|
monkeypatch.setattr(mh, "MAX_METADATA_JSON_BYTES", 1)
|
|
too_large = await mh.MetadataHelpers.store_metadata_cache(db, "a.png", "h", Result.Ok({"x": "long"}))
|
|
assert too_large.code == "CACHE_SKIPPED"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_refresh_if_needed_and_cleanup_paths(monkeypatch):
|
|
db = _DB()
|
|
db.query_result = Result.Ok([{"has_workflow": 0, "has_generation_data": 0, "metadata_raw": "{}"}])
|
|
called = {"journal": 0}
|
|
|
|
async def _journal(*args, **kwargs):
|
|
called["journal"] += 1
|
|
|
|
res = await mh.MetadataHelpers.refresh_metadata_if_needed(
|
|
db,
|
|
asset_id=1,
|
|
metadata_result=Result.Ok({"parameters": "x"}),
|
|
filepath="a.png",
|
|
base_dir=".",
|
|
state_hash="h",
|
|
mtime=1,
|
|
size=2,
|
|
write_journal_fn=_journal,
|
|
)
|
|
assert res is True and called["journal"] == 1
|
|
|
|
monkeypatch.setattr(mh, "METADATA_CACHE_MAX", 10)
|
|
monkeypatch.setattr(mh, "METADATA_CACHE_TTL_SECONDS", 10)
|
|
monkeypatch.setattr(mh, "METADATA_CACHE_CLEANUP_INTERVAL_SECONDS", 0)
|
|
db.query_result = Result.Ok([{"count": 20}])
|
|
await mh.MetadataHelpers._maybe_cleanup_metadata_cache(db)
|
|
assert db.query_calls and db.exec_calls
|
|
|
|
assert mh._cache_cleanup_config()[0] >= 0
|
|
await mh._cleanup_cache_by_ttl(db, 5)
|
|
await mh._cleanup_cache_by_max_entries(db, 10)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_write_row_preserves_generation_time_across_quality_upgrade(tmp_path: Path):
|
|
db = await _make_db(tmp_path)
|
|
try:
|
|
ins = await db.aexecute(
|
|
"INSERT INTO assets(filepath, filename, subfolder, source, root_id, kind, ext, size, mtime) "
|
|
"VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?)",
|
|
("C:\\x\\gentime.png", "gentime.png", "", "output", "output", "image", ".png", 123, 1700000000),
|
|
)
|
|
assert ins.ok
|
|
asset_rows = (await db.aquery("SELECT id FROM assets")).data
|
|
assert asset_rows
|
|
asset_id = int(asset_rows[0]["id"])
|
|
|
|
initial = await mh.MetadataHelpers.write_asset_metadata_row(
|
|
db,
|
|
asset_id,
|
|
Result.Ok({"quality": "partial", "parameters": "prompt", "generation_time_ms": 8421}),
|
|
filepath="gentime.png",
|
|
)
|
|
assert initial.ok
|
|
|
|
upgraded = await mh.MetadataHelpers.write_asset_metadata_row(
|
|
db,
|
|
asset_id,
|
|
Result.Ok(
|
|
{
|
|
"quality": "full",
|
|
"workflow": {"nodes": [{"id": 1, "type": "KSampler"}]},
|
|
"prompt": {"1": {"class_type": "KSampler", "inputs": {"steps": 20}}},
|
|
}
|
|
),
|
|
filepath="gentime.png",
|
|
)
|
|
assert upgraded.ok
|
|
|
|
rows = (
|
|
await db.aquery(
|
|
"SELECT metadata_quality, generation_time_ms FROM asset_metadata WHERE asset_id = ?",
|
|
(asset_id,),
|
|
)
|
|
).data
|
|
assert rows
|
|
row = rows[0]
|
|
assert row["metadata_quality"] == "full"
|
|
assert row["generation_time_ms"] == 8421
|
|
finally:
|
|
await db.aclose()
|
|
|
|
|
|
def test_metadata_payload_size_helpers(monkeypatch):
|
|
monkeypatch.setattr(mh, "MAX_METADATA_JSON_BYTES", 10)
|
|
raw = mh.MetadataHelpers._metadata_json_payload({"b": 1, "a": 2})
|
|
assert '"a":2' in raw and '"b":1' in raw
|
|
assert mh.MetadataHelpers._metadata_payload_size_bytes(raw) > 0
|
|
assert mh.MetadataHelpers.compute_metadata_hash(raw)
|
|
|
|
guard = mh.MetadataHelpers._metadata_payload_size_guard("x" * 100, "a.png")
|
|
assert guard and guard.code == "CACHE_SKIPPED"
|