Vito Sansevero
ed56c39745
Same bytes now give the server's record in five shapes where the local reader was stricter or kept different tags (PR #277 final review I1, I2, M1, M2, M4): - Invalid UTF-8 in a tag value is replaced as ffprobe's default string_validation=replace does (av_utf8_decode, one U+FFFD per failing decode step), instead of failing the read. Key names and TagNames that do not decode are no carrier and are skipped. - An ilst item whose keys index is 0 or out of range is skipped, as ffmpeg does, instead of failing the read. - Only the first keys box is read, as ffmpeg does. - A first per-key workflow over 4 MiB keeps its place in the carrier choice as an unread marker, so the server's carrier (per-key, workflow limited) is chosen; when the server's first probe would overflow, the re-probe leaves it out and its own bound applies. - An envelope integer of more than 4300 digits fails the Python json port, so that envelope is skipped as CPython's json.loads skips it. Each shape was checked against ffprobe 9.0.2 and read_inspection_file; the real-ffprobe parity test now builds them, and a 600-case random-byte comparison of the replacement against ffprobe matched.
1065 lines
40 KiB
Python
1065 lines
40 KiB
Python
"""The browser's local readers answer exactly as the server's on the same file.
|
||
|
||
``web/js/png-metadata.js`` ``readPngMetadataFile`` ports the bounded PNG reader in
|
||
``utils/prompt_inspector.py``. This test hands the inspection parity corpus (the
|
||
committed ``f*.png`` and the test-time shapes), extra compressed, oversized,
|
||
Latin-1 and batch-duplicate PNGs, and the committed JPEG/WebP header fixtures to a
|
||
Node oracle, and asserts per file that records, record formats, dimensions and
|
||
readability equal ``read_inspection_file``. The oracle reads each file through
|
||
``fs.openAsBlob`` slices, as the page reads a chosen ``File``.
|
||
|
||
``web/js/workspace/local-video.mjs`` reads the container tags the server asks
|
||
ffprobe for and ports ``utils/video_metadata.select_video_carrier`` with the
|
||
record budgets of ``read_inspection_file``. The video tests hand the oracle the
|
||
committed container fixtures (``tests/fixtures/local-inspect/*.{mp4,mov,m4v,mkv,
|
||
webm}`` with their recorded ffprobe tag maps in ``video-tags-v1.json``), recorded
|
||
and synthetic tag maps, and multi-megabyte per-key containers built at test time,
|
||
and assert that the local carrier equals the server's, which selects from the
|
||
same tag map through a mocked probe, or from the same file through real ffprobe.
|
||
|
||
``web/js/workspace/local-projection.mjs`` ports ``utils/inspection_response.py``:
|
||
the strict Decimal/duplicate-key parse, the prompt, workflow and carrier
|
||
projection, the depth and node ceilings and ``_metadata``. The projection tests
|
||
hand the oracle every file above plus the committed adversarial record sets in
|
||
``tests/fixtures/local-inspect/projection-cases.json`` and node-ceiling sets built
|
||
at test time, and assert that ``projectMetadata`` equals ``_metadata`` byte for
|
||
byte, as the server's ``_bytes`` writes it with sorted keys.
|
||
"""
|
||
|
||
import io
|
||
import json
|
||
import os
|
||
import shutil
|
||
import struct
|
||
import subprocess # nosec B404
|
||
import sys
|
||
import tempfile
|
||
import unittest
|
||
import zlib
|
||
from pathlib import Path
|
||
from unittest.mock import patch
|
||
|
||
from PIL import Image
|
||
from PIL.PngImagePlugin import PngInfo
|
||
|
||
from utils.inspection_response import GRAPH_NODE_CEILING, _bytes, _metadata
|
||
from utils.prompt_inspector import (
|
||
ROLE_METADATA_BYTES,
|
||
WORKFLOW_METADATA_BYTES,
|
||
InspectionRead,
|
||
MetadataReadError,
|
||
read_inspection_file,
|
||
)
|
||
|
||
try:
|
||
from tests import inspection_parity as parity
|
||
from tests.open_handles import assert_closed
|
||
from tests.test_utils_video_metadata import HAVE_FFMPEG, write_video
|
||
except ImportError: # pragma: no cover - discover -s tests puts tests/ first
|
||
import inspection_parity as parity
|
||
from open_handles import assert_closed
|
||
from test_utils_video_metadata import HAVE_FFMPEG, write_video
|
||
|
||
HERE = Path(__file__).resolve().parent
|
||
ORACLE = HERE / "js" / "support" / "local-inspect-oracle.mjs"
|
||
HEADER_FIXTURES = HERE / "fixtures" / "local-inspect"
|
||
VIDEO_FIXTURES = HEADER_FIXTURES
|
||
RECORDED_TAGS = json.loads(
|
||
(VIDEO_FIXTURES / "video-tags-v1.json").read_text(encoding="utf-8")
|
||
)["fixtures"]
|
||
PROJECTION_CASES = json.loads(
|
||
(HEADER_FIXTURES / "projection-cases.json").read_text(encoding="utf-8")
|
||
)["cases"]
|
||
NODE_TIMEOUT_SECONDS = 180
|
||
PNG_CLASS = {"media": "image", "kind": "png"}
|
||
_FOOTER = "Steps: 20, Sampler: Euler, Batch index: {}, Seed: 7"
|
||
MiB = 1024 * 1024
|
||
GRAPH = '{"3":{"class_type":"KSampler","inputs":{"seed":7}}}'
|
||
WORKFLOW = '{"nodes":[],"version":0.4}'
|
||
ENVELOPE = json.dumps({"prompt": '{"9":{}}', "workflow": {"v": 1}})
|
||
VIDEO_CONTAINERS = {
|
||
".mp4": "iso-bmff",
|
||
".mov": "iso-bmff",
|
||
".m4v": "iso-bmff",
|
||
".mkv": "matroska",
|
||
".webm": "matroska",
|
||
}
|
||
PER_KEY_NAMES = {"prompt", "workflow", "parameters", "comfyassets_prompt_pairs"}
|
||
|
||
# Tag maps as ffprobe would report them, covering every branch of the carrier rule:
|
||
# (name, tags, records the reader left out). The server answer is computed, not
|
||
# written down, so a change in either side fails here.
|
||
TAG_MAP_VECTORS = [
|
||
(
|
||
"per_key_wins_over_comment",
|
||
{"prompt": GRAPH, "workflow": WORKFLOW, "comment": ENVELOPE},
|
||
[],
|
||
),
|
||
(
|
||
"matroska_names_fold",
|
||
{"PROMPT": GRAPH, "Workflow": WORKFLOW, "ENCODER": "Lavf"},
|
||
[],
|
||
),
|
||
("identical_fold_duplicates", {"PROMPT": GRAPH, "prompt": GRAPH}, []),
|
||
("differing_fold_duplicates_fail", {"PROMPT": GRAPH, "prompt": '{"9":{}}'}, []),
|
||
(
|
||
"differing_workflow_duplicates_fail",
|
||
{"Workflow": WORKFLOW, "WORKFLOW": "{}", "prompt": GRAPH},
|
||
[],
|
||
),
|
||
(
|
||
"envelope_names_never_fold",
|
||
{"comment": ENVELOPE, "COMMENT": "a free-text note"},
|
||
[],
|
||
),
|
||
(
|
||
"non_json_comment_falls_to_description",
|
||
{"comment": "not json", "description": ENVELOPE},
|
||
[],
|
||
),
|
||
(
|
||
"array_comment_is_not_an_envelope",
|
||
{"comment": '["prompt"]', "description": '{"workflow": {}}'},
|
||
[],
|
||
),
|
||
(
|
||
"role_envelope_beats_workflow_only_per_key",
|
||
{"workflow": WORKFLOW, "comment": ENVELOPE},
|
||
[],
|
||
),
|
||
(
|
||
"role_description_beats_workflow_comment",
|
||
{"comment": '{"workflow": {"a": 1}}', "description": '{"parameters": "cat"}'},
|
||
[],
|
||
),
|
||
(
|
||
"workflow_only_carriers_never_merge",
|
||
{"comment": '{"workflow": {"a": 1}}', "description": '{"workflow": {"b": 2}}'},
|
||
[],
|
||
),
|
||
("description_only", {"description": ENVELOPE}, []),
|
||
("no_known_key", {"encoder": "Lavf", "major_brand": "isom"}, []),
|
||
("envelope_without_known_keys", {"comment": '{"title": "x"}'}, []),
|
||
("null_role_marker_refused", {"comment": '{"prompt": null, "workflow": {}}'}, []),
|
||
("numeric_role_refused", {"description": '{"parameters": 5}'}, []),
|
||
(
|
||
"invalid_envelope_workflow_keeps_roles",
|
||
{"comment": '{"prompt": "{}", "workflow": null}'},
|
||
[],
|
||
),
|
||
(
|
||
"object_prompt_reserialized",
|
||
{
|
||
"comment": '{"prompt": {"1": {"class_type": "A", "f": 8.0, "n": 1e16, '
|
||
'"u": "é😀"}}}'
|
||
},
|
||
[],
|
||
),
|
||
("workflow_only_invalid_envelope", {"comment": '{"workflow": 5}'}, []),
|
||
("non_text_per_key_role_refused", {"prompt": 5, "workflow": WORKFLOW}, []),
|
||
("non_text_per_key_workflow_invalid", {"prompt": GRAPH, "workflow": 5}, []),
|
||
(
|
||
"python_number_spellings",
|
||
{
|
||
"description": '{"workflow": {"a": [1.0, 0.1, 1E2, 1e-7, -0.0, -0, '
|
||
"123456789012345678, 1e400, NaN, 0.00001, 1e15], "
|
||
'"b": "\\u0001\\u007f/"}}'
|
||
},
|
||
[],
|
||
),
|
||
(
|
||
"integer_like_keys_keep_order",
|
||
{"description": '{"workflow": {"10": 1, "3": 2, "b": 3, "a": 4, "3": 5}}'},
|
||
[],
|
||
),
|
||
(
|
||
"workflow_over_budget_limited",
|
||
{"prompt": GRAPH, "workflow": "{" + "0" * WORKFLOW_METADATA_BYTES},
|
||
[],
|
||
),
|
||
("workflow_at_budget_kept", {"workflow": "0" * WORKFLOW_METADATA_BYTES}, []),
|
||
(
|
||
"roles_over_budget_fail",
|
||
{"prompt": "a" * ROLE_METADATA_BYTES, "parameters": "b"},
|
||
[],
|
||
),
|
||
("roles_at_budget_kept", {"prompt": "é" * (ROLE_METADATA_BYTES // 2)}, []),
|
||
(
|
||
"lone_surrogates_count_three_bytes",
|
||
{"prompt": "\ud800" * (ROLE_METADATA_BYTES // 3) + "aaa"},
|
||
[],
|
||
),
|
||
(
|
||
"envelope_workflow_over_budget_limited",
|
||
{
|
||
"comment": json.dumps(
|
||
{"prompt": "x", "workflow": {"k": "0" * WORKFLOW_METADATA_BYTES}}
|
||
)
|
||
},
|
||
[],
|
||
),
|
||
("reader_limited_workflow_with_per_key", {"prompt": GRAPH}, ["workflow"]),
|
||
("reader_limited_workflow_without_carrier", {"encoder": "x"}, ["workflow"]),
|
||
(
|
||
"reader_limited_workflow_cleared_by_envelope",
|
||
{"comment": ENVELOPE},
|
||
["workflow"],
|
||
),
|
||
# CPython's json.loads refuses an integer of more than 4300 digits, so the
|
||
# server skips that envelope for the next one.
|
||
(
|
||
"envelope_integer_over_digit_limit_skipped",
|
||
{
|
||
"comment": '{"workflow": {"a": 1' + "0" * 4400 + "}}",
|
||
"description": '{"workflow": {"b": 2}}',
|
||
},
|
||
[],
|
||
),
|
||
]
|
||
|
||
|
||
def _raw_chunk(kind, data):
|
||
return (
|
||
struct.pack(">I", len(data))
|
||
+ kind
|
||
+ data
|
||
+ struct.pack(">I", zlib.crc32(kind + data) & 0xFFFFFFFF)
|
||
)
|
||
|
||
|
||
def _text_chunk(key, text):
|
||
return _raw_chunk(b"tEXt", key.encode("latin-1") + b"\0" + text.encode("latin-1"))
|
||
|
||
|
||
def _write_png(path, info, extra_chunks=()):
|
||
"""A 16x16 PNG with Pillow-written text records and spliced raw chunks."""
|
||
buffer = io.BytesIO()
|
||
with Image.new("RGB", (16, 16), (10, 20, 30)) as image:
|
||
image.save(buffer, format="PNG", pnginfo=info, compress_level=4)
|
||
data = buffer.getvalue()
|
||
iend = data.rindex(b"IEND") - 4
|
||
with open(path, "wb") as stream:
|
||
stream.write(data[:iend] + b"".join(extra_chunks) + data[iend:])
|
||
|
||
|
||
def _info(**texts):
|
||
info = PngInfo()
|
||
for key, value in texts.items():
|
||
if isinstance(value, tuple):
|
||
text, options = value
|
||
if options.get("itxt"):
|
||
info.add_itxt(
|
||
key,
|
||
text,
|
||
lang=options.get("lang", ""),
|
||
tkey=options.get("tkey", ""),
|
||
zip=options.get("zip", False),
|
||
)
|
||
else:
|
||
info.add_text(key, text, zip=options.get("zip", False))
|
||
else:
|
||
info.add_text(key, value)
|
||
return info
|
||
|
||
|
||
PROMPT = json.dumps(parity.build("f02")["prompt"])
|
||
WORKFLOW = json.dumps(parity.build("f02")["workflow"])
|
||
|
||
# Shapes the committed corpus lacks: compressed chunks, the budgets, Latin-1 text,
|
||
# the batch-index rule and the pair execution record.
|
||
EXTRA_SHAPES = {
|
||
"ztxt_prompt_itxt_workflow": lambda path: _write_png(
|
||
path,
|
||
_info(
|
||
prompt=(PROMPT, {"zip": True}),
|
||
workflow=(WORKFLOW, {"itxt": True, "zip": True}),
|
||
),
|
||
),
|
||
"itxt_bom_lang": lambda path: _write_png(
|
||
path,
|
||
_info(
|
||
prompt=("" + PROMPT, {"itxt": True, "lang": "en", "tkey": "prompt"}),
|
||
workflow=(WORKFLOW, {"itxt": True}),
|
||
),
|
||
),
|
||
"latin1_prompt": lambda path: _write_png(
|
||
path, _info(prompt='{"1":{"class_type":"Note","inputs":{"text":"café ÿ"}}}')
|
||
),
|
||
"all_four_records": lambda path: _write_png(
|
||
path,
|
||
_info(
|
||
prompt=PROMPT,
|
||
workflow=WORKFLOW,
|
||
parameters="cat\nNegative prompt: dog\n" + _FOOTER.format(0),
|
||
comfyassets_prompt_pairs='{"pairs":[]}',
|
||
),
|
||
),
|
||
"workflow_over_budget_limited": lambda path: _write_png(
|
||
path,
|
||
_info(
|
||
prompt=PROMPT,
|
||
workflow="[" + "0," * ((WORKFLOW_METADATA_BYTES + 1) // 2) + "]",
|
||
),
|
||
),
|
||
"workflow_at_budget_kept": lambda path: _write_png(
|
||
path, _info(workflow=("{" * WORKFLOW_METADATA_BYTES, {"zip": True}))
|
||
),
|
||
"roles_over_budget_unreadable": lambda path: _write_png(
|
||
path, _info(prompt="A" * (ROLE_METADATA_BYTES + 1))
|
||
),
|
||
"roles_over_budget_decoded_unreadable": lambda path: _write_png(
|
||
path, _info(prompt=("A" * (ROLE_METADATA_BYTES + 1), {"zip": True}))
|
||
),
|
||
"roles_at_budget_kept": lambda path: _write_png(
|
||
path,
|
||
_info(
|
||
prompt=("A" * (ROLE_METADATA_BYTES // 2), {"zip": True}),
|
||
parameters=("B" * (ROLE_METADATA_BYTES // 2), {"zip": True}),
|
||
),
|
||
),
|
||
"batch_parameters_last_kept": lambda path: _write_png(
|
||
path,
|
||
_info(prompt=PROMPT),
|
||
[
|
||
_text_chunk("parameters", "cat\n" + _FOOTER.format(index))
|
||
for index in (0, 1, 2, 10)
|
||
],
|
||
),
|
||
"batch_parameters_conflict_unreadable": lambda path: _write_png(
|
||
path,
|
||
PngInfo(),
|
||
[
|
||
_text_chunk("parameters", "cat\n" + _FOOTER.format(0)),
|
||
_text_chunk("parameters", "dog\n" + _FOOTER.format(1)),
|
||
],
|
||
),
|
||
"duplicate_workflow_identical": lambda path: _write_png(
|
||
path,
|
||
_info(prompt=PROMPT),
|
||
[_text_chunk("workflow", WORKFLOW), _text_chunk("workflow", WORKFLOW)],
|
||
),
|
||
"unknown_keys_only": lambda path: _write_png(
|
||
path, _info(Software="ComfyUI", Comment=("café", {"zip": True}))
|
||
),
|
||
}
|
||
|
||
|
||
def _classes(results):
|
||
return {name: entry["classified"] for name, entry in results.items()}
|
||
|
||
|
||
def _run_oracle(testcase, entries):
|
||
node = shutil.which("node")
|
||
if node is None: # pragma: no cover - CI runners and dev hosts ship Node
|
||
if os.environ.get("CI"):
|
||
testcase.fail("Node.js is required in CI for the local reader parity test")
|
||
testcase.skipTest("Node.js is not installed")
|
||
# Fixed local Node harness, no shell.
|
||
completed = subprocess.run( # nosec B603
|
||
[node, str(ORACLE)],
|
||
input=json.dumps(entries).encode("utf-8"),
|
||
capture_output=True,
|
||
timeout=NODE_TIMEOUT_SECONDS,
|
||
check=False,
|
||
)
|
||
testcase.assertEqual(completed.returncode, 0, completed.stderr.decode("utf-8"))
|
||
return json.loads(completed.stdout)
|
||
|
||
|
||
def _known_tags(tags):
|
||
"""The tags the local reader keeps: per-key names folded, envelope names exact."""
|
||
return {
|
||
key: value
|
||
for key, value in tags.items()
|
||
if key.lower() in PER_KEY_NAMES or key in ("comment", "description")
|
||
}
|
||
|
||
|
||
def _fake_probe(tags, drop_workflow=False):
|
||
"""``_probe_video_tags`` answering from a tag map as ffprobe would.
|
||
|
||
ffprobe's ``format_tags=`` filter is case-insensitive and the stdout bound is
|
||
the only reason a probe answers ``None``; ``drop_workflow`` makes the first
|
||
probe (the one asking for ``workflow``) overflow, as a huge per-key workflow
|
||
tag does, so the server takes its role-only re-probe.
|
||
"""
|
||
|
||
def probe(binary, path, keys, limit):
|
||
if drop_workflow and "workflow" in keys:
|
||
return None
|
||
selected = {key: value for key, value in tags.items() if key.lower() in keys}
|
||
raw = json.dumps({"format": {"tags": selected}}, ensure_ascii=False).encode(
|
||
"utf-8", "surrogatepass"
|
||
)
|
||
return None if len(raw) > limit else raw
|
||
|
||
return probe
|
||
|
||
|
||
def _server_selection(tags, limited=()):
|
||
"""The server's carrier for a tag map, in the oracle's shape; None when it fails."""
|
||
with (
|
||
patch("utils.prompt_inspector.ffprobe_path", return_value="/tool/ffprobe"),
|
||
patch(
|
||
"utils.prompt_inspector._probe_video_tags",
|
||
side_effect=_fake_probe(tags, drop_workflow="workflow" in limited),
|
||
),
|
||
):
|
||
try:
|
||
read = read_inspection_file(Path("vector.mp4"))
|
||
except MetadataReadError:
|
||
return None
|
||
return _read_shape(read)
|
||
|
||
|
||
def _read_shape(read):
|
||
return {
|
||
"records": read.metadata,
|
||
"record_formats": read.record_formats,
|
||
"carrierText": read.video_carrier_text,
|
||
"limited": sorted(read.limited_records),
|
||
}
|
||
|
||
|
||
def _local_shape(selection):
|
||
if selection["status"] != "complete":
|
||
return None
|
||
return {
|
||
"records": selection["records"],
|
||
"record_formats": selection["record_formats"],
|
||
"carrierText": selection["carrierText"],
|
||
"limited": sorted(selection["limited"]),
|
||
}
|
||
|
||
|
||
def _box(kind, *parts):
|
||
payload = b"".join(parts)
|
||
return struct.pack(">I", 8 + len(payload)) + kind + payload
|
||
|
||
|
||
def _full_box(kind, *parts):
|
||
return _box(kind, b"\0\0\0\0", *parts)
|
||
|
||
|
||
def _mdta_udta(tags):
|
||
"""``moov/udta/meta`` with one mdta key per tag, laid out as ffmpeg writes it."""
|
||
names = [name.encode("utf-8") for name in tags]
|
||
values = [value.encode("utf-8") for value in tags.values()]
|
||
return _mdta_items(names, list(enumerate(values, 1)))
|
||
|
||
|
||
def _mdta_items(names, items, keys_boxes=(), legacy=()):
|
||
"""``moov/udta/meta`` from raw bytes: mdta key ``names``, ``(index, value)``
|
||
items, extra ``keys`` boxes after the first (each a list of names) and
|
||
``(fourcc, value)`` legacy atoms in the same ``ilst``."""
|
||
|
||
def keys(entries):
|
||
return _full_box(
|
||
b"keys",
|
||
struct.pack(">I", len(entries)),
|
||
*(struct.pack(">I", 8 + len(name)) + b"mdta" + name for name in entries),
|
||
)
|
||
|
||
def data(value):
|
||
return _box(b"data", struct.pack(">II", 1, 0), value)
|
||
|
||
ilst = _box(
|
||
b"ilst",
|
||
*(_box(struct.pack(">I", index), data(value)) for index, value in items),
|
||
*(_box(fourcc.encode("latin-1"), data(value)) for fourcc, value in legacy),
|
||
)
|
||
hdlr = _full_box(b"hdlr", b"\0\0\0\0", b"mdta", bytes(12), b"\0")
|
||
boxes = [keys(names), *(keys(entries) for entries in keys_boxes)]
|
||
return _box(b"udta", _full_box(b"meta", hdlr, *boxes, ilst))
|
||
|
||
|
||
def _with_udta(data, udta):
|
||
"""``data`` (an MP4 with moov after mdat) with its moov's udta replaced."""
|
||
out = bytearray()
|
||
offset = 0
|
||
while offset < len(data):
|
||
size, kind = struct.unpack(">I4s", data[offset : offset + 8])
|
||
if kind != b"moov":
|
||
out += data[offset : offset + size]
|
||
offset += size
|
||
continue
|
||
body = data[offset + 8 : offset + size]
|
||
children = bytearray()
|
||
inner = 0
|
||
while inner < len(body):
|
||
child_size, child_kind = struct.unpack(">I4s", body[inner : inner + 8])
|
||
if child_kind != b"udta":
|
||
children += body[inner : inner + child_size]
|
||
inner += child_size
|
||
children += udta
|
||
out += struct.pack(">I", 8 + len(children)) + b"moov" + bytes(children)
|
||
offset += size
|
||
return bytes(out)
|
||
|
||
|
||
def _text_of(size, prefix):
|
||
"""A valid JSON object of exactly ``size`` UTF-8 bytes with quotes throughout."""
|
||
unit = '"k":"v",'
|
||
body = (unit * (size // len(unit) + 1))[: size - len(prefix) - 2]
|
||
body = body[: body.rfind(",")] if body.endswith(",") else body
|
||
text = prefix + body + "}"
|
||
return text + " " * (size - len(text))
|
||
|
||
|
||
class LocalPngReaderParityTests(unittest.TestCase):
|
||
def setUp(self):
|
||
temporary = tempfile.TemporaryDirectory()
|
||
self.addCleanup(temporary.cleanup)
|
||
self.addCleanup(assert_closed, self, temporary.name)
|
||
self.directory = Path(temporary.name)
|
||
|
||
def _files(self):
|
||
files = {}
|
||
for name in parity.COMMITTED:
|
||
files[name] = parity.fixture_path(name)
|
||
for name in parity.TEST_TIME:
|
||
path = parity.fixture_path(name, self.directory)
|
||
parity.write_png(path, parity.build(name))
|
||
files[name] = path
|
||
for name, write in EXTRA_SHAPES.items():
|
||
path = self.directory / f"{name}.png"
|
||
write(path)
|
||
files[name] = path
|
||
for header in ("tiny.jpg", "tiny.webp"):
|
||
files[header] = HEADER_FIXTURES / header
|
||
return files
|
||
|
||
def _oracle(self, files):
|
||
return _run_oracle(
|
||
self, [{"name": name, "path": str(path)} for name, path in files.items()]
|
||
)
|
||
|
||
def test_local_png_reader_matches_read_inspection_file(self):
|
||
files = self._files()
|
||
observed = self._oracle(files)
|
||
self.assertEqual(set(observed), set(files))
|
||
outcomes = set()
|
||
for name, path in files.items():
|
||
if path.suffix != ".png":
|
||
continue
|
||
with self.subTest(fixture=name):
|
||
self.assertEqual(observed[name]["classified"], PNG_CLASS)
|
||
result = observed[name]["result"]
|
||
try:
|
||
read = read_inspection_file(path)
|
||
except MetadataReadError:
|
||
outcomes.add("unreadable")
|
||
self.assertEqual(result["status"], "unreadable", result)
|
||
continue
|
||
self.assertEqual(result["status"], "complete", result)
|
||
self.assertEqual(result["records"], read.metadata)
|
||
self.assertEqual(
|
||
sorted(result["limited"]), sorted(read.limited_records)
|
||
)
|
||
formats = {key: "original-text" for key in result["records"]}
|
||
formats.update({key: "limited" for key in result["limited"]})
|
||
self.assertEqual(formats, read.record_formats)
|
||
self.assertEqual(
|
||
(result["width"], result["height"]), (read.width, read.height)
|
||
)
|
||
self.assertEqual(read.format, "PNG")
|
||
outcomes.add("limited" if read.limited_records else "complete")
|
||
# The corpus must exercise every outcome, or agreement proves little.
|
||
self.assertEqual(outcomes, {"complete", "limited", "unreadable"})
|
||
|
||
def test_local_image_header_matches_pillow(self):
|
||
files = {name: HEADER_FIXTURES / name for name in ("tiny.jpg", "tiny.webp")}
|
||
observed = self._oracle(files)
|
||
self.assertEqual(
|
||
_classes(observed),
|
||
{
|
||
"tiny.jpg": {"media": "image", "kind": "jpeg"},
|
||
"tiny.webp": {"media": "image", "kind": "webp"},
|
||
},
|
||
)
|
||
for name, path in files.items():
|
||
with self.subTest(fixture=name):
|
||
read = read_inspection_file(path)
|
||
self.assertEqual(read.metadata, {})
|
||
self.assertEqual(
|
||
observed[name]["result"],
|
||
{"format": read.format, "width": read.width, "height": read.height},
|
||
)
|
||
self.assertLess(path.stat().st_size, 4096)
|
||
|
||
|
||
class LocalVideoCarrierParityTests(unittest.TestCase):
|
||
"""The local video readers and carrier rule answer as the server does."""
|
||
|
||
def setUp(self):
|
||
temporary = tempfile.TemporaryDirectory()
|
||
self.addCleanup(temporary.cleanup)
|
||
self.addCleanup(assert_closed, self, temporary.name)
|
||
self.directory = Path(temporary.name)
|
||
|
||
def test_local_carrier_matches_server_selection(self):
|
||
vectors = [(name, tags, []) for name, tags in RECORDED_TAGS.items()]
|
||
vectors += TAG_MAP_VECTORS
|
||
observed = _run_oracle(
|
||
self,
|
||
[
|
||
{"name": name, "tags": tags, "limited": limited}
|
||
for name, tags, limited in vectors
|
||
],
|
||
)
|
||
outcomes = set()
|
||
for name, tags, limited in vectors:
|
||
with self.subTest(vector=name):
|
||
selection = observed[name]["result"]["selection"]
|
||
server = _server_selection(tags, limited)
|
||
self.assertEqual(
|
||
_local_shape(selection), server, selection.get("reason")
|
||
)
|
||
if server is None:
|
||
outcomes.add(selection["status"] + ":" + selection["reason"])
|
||
elif not server["records"] and not server["record_formats"]:
|
||
outcomes.add("no-carrier")
|
||
else:
|
||
outcomes.update(server["record_formats"].values())
|
||
# Every outcome of the rule must occur, or agreement proves little.
|
||
self.assertEqual(
|
||
outcomes,
|
||
{
|
||
"original-text",
|
||
"reserialized-json",
|
||
"invalid",
|
||
"limited",
|
||
"no-carrier",
|
||
"unreadable:duplicate-metadata",
|
||
"unreadable:invalid-metadata",
|
||
"limited:metadata-limit",
|
||
},
|
||
)
|
||
|
||
def test_local_video_readers_match_recorded_tags_and_server(self):
|
||
files = {name: VIDEO_FIXTURES / name for name in RECORDED_TAGS}
|
||
observed = _run_oracle(
|
||
self, [{"name": name, "path": str(path)} for name, path in files.items()]
|
||
)
|
||
for name, path in files.items():
|
||
with self.subTest(fixture=name):
|
||
self.assertLess(path.stat().st_size, 16 * 1024)
|
||
container = VIDEO_CONTAINERS[path.suffix]
|
||
self.assertEqual(
|
||
observed[name]["classified"], {"media": "video", "kind": container}
|
||
)
|
||
read = observed[name]["result"]["read"]
|
||
self.assertEqual(read["status"], "complete", read)
|
||
self.assertEqual(read["container"], container)
|
||
self.assertEqual(read["tags"], _known_tags(RECORDED_TAGS[name]))
|
||
self.assertEqual(read["limited"], [])
|
||
local = _local_shape(observed[name]["result"]["selection"])
|
||
self.assertEqual(local, _server_selection(RECORDED_TAGS[name]))
|
||
if HAVE_FFMPEG:
|
||
self.assertEqual(local, _read_shape(read_inspection_file(path)))
|
||
# The corpus covers per-key carriers in every container, an envelope and a
|
||
# Matroska COMMENT that the exact-case envelope rule leaves unread.
|
||
self.assertTrue(
|
||
_server_selection(RECORDED_TAGS["legacy-comment.mp4"])["carrierText"]
|
||
)
|
||
self.assertEqual(
|
||
_server_selection(RECORDED_TAGS["legacy-comment.webm"])["records"], {}
|
||
)
|
||
|
||
@unittest.skipUnless(HAVE_FFMPEG, "ffmpeg/ffprobe not installed")
|
||
def test_large_per_key_videos_read_the_same_as_the_server(self):
|
||
base = self.directory / "base.mp4"
|
||
write_video(base)
|
||
source = base.read_bytes()
|
||
shapes = {
|
||
"available": {
|
||
"prompt": _text_of(
|
||
MiB, '{"1":{"class_type":"Note","inputs":{"text":"a"}},'
|
||
),
|
||
"workflow": _text_of(5 * MiB // 2, '{"nodes":[],"version":0.4,'),
|
||
},
|
||
"workflow_limited": {
|
||
"prompt": GRAPH,
|
||
"workflow": "[" + "0," * (WORKFLOW_METADATA_BYTES // 2) + "]",
|
||
},
|
||
"roles_over_budget": {"prompt": "A" * (ROLE_METADATA_BYTES + 1)},
|
||
}
|
||
files = {}
|
||
for name, tags in shapes.items():
|
||
path = self.directory / f"{name}.mp4"
|
||
path.write_bytes(_with_udta(source, _mdta_udta(tags)))
|
||
files[name] = path
|
||
observed = _run_oracle(
|
||
self, [{"name": name, "path": str(path)} for name, path in files.items()]
|
||
)
|
||
available = observed["available"]["result"]
|
||
self.assertEqual(available["read"]["status"], "complete")
|
||
self.assertEqual(available["read"]["limited"], [])
|
||
self.assertEqual(
|
||
available["selection"]["records"], shapes["available"], "1 MiB + 2.5 MiB"
|
||
)
|
||
self.assertEqual(
|
||
_local_shape(available["selection"]),
|
||
_read_shape(read_inspection_file(files["available"])),
|
||
)
|
||
limited = observed["workflow_limited"]["result"]
|
||
self.assertEqual(limited["read"]["limited"], ["workflow"])
|
||
# The unread workflow keeps its place in the carrier choice as a marker.
|
||
self.assertEqual(limited["read"]["tags"]["workflow"], {"oversized": "workflow"})
|
||
self.assertEqual(
|
||
_local_shape(limited["selection"]),
|
||
_read_shape(read_inspection_file(files["workflow_limited"])),
|
||
)
|
||
self.assertEqual(
|
||
limited["selection"]["record_formats"],
|
||
{"prompt": "original-text", "workflow": "limited"},
|
||
)
|
||
over = observed["roles_over_budget"]["result"]
|
||
self.assertEqual(
|
||
over["read"], {"status": "limited", "reason": "metadata-limit"}
|
||
)
|
||
self.assertIsNone(over["selection"])
|
||
with self.assertRaises(MetadataReadError):
|
||
read_inspection_file(files["roles_over_budget"])
|
||
|
||
@unittest.skipUnless(HAVE_FFMPEG, "ffmpeg/ffprobe not installed")
|
||
def test_byte_level_container_shapes_read_the_same_as_the_server(self):
|
||
"""Shapes where ffprobe is lenient: invalid UTF-8 is replaced (one U+FFFD
|
||
per failing decode step), names that do not decode are no carrier, stray
|
||
``ilst`` indexes are skipped, the first ``keys`` box wins, and an
|
||
oversized per-key workflow still takes part in the carrier choice."""
|
||
base = self.directory / "base.mp4"
|
||
write_video(base)
|
||
source = base.read_bytes()
|
||
graph, workflow = GRAPH.encode(), WORKFLOW.encode()
|
||
envelope = json.dumps({"workflow": {"nodes": [1]}}).encode()
|
||
oversized = b'{"nodes":[' + b"0," * (5 * MiB // 2) + b"0]}"
|
||
shapes = {
|
||
"latin1_byte_in_workflow": _mdta_items(
|
||
[b"prompt", b"workflow"], [(1, graph), (2, b'{"a":"caf\xe9"}')]
|
||
),
|
||
"undecodable_key_name": _mdta_items(
|
||
[b"prom\xffpt", b"prompt"], [(1, b"junk"), (2, graph)]
|
||
),
|
||
"stray_ilst_indexes": _mdta_items(
|
||
[b"prompt", b"workflow"],
|
||
[(1, graph), (5, b"junk"), (0, b"junk"), (2, workflow)],
|
||
),
|
||
"second_keys_box": _mdta_items(
|
||
[b"prompt"], [(1, graph)], keys_boxes=[[b"workflow"]]
|
||
),
|
||
"oversized_workflow_beside_envelope": _mdta_items(
|
||
[b"workflow"], [(1, oversized)], legacy=[("\u00a9cmt", envelope)]
|
||
),
|
||
# A repeat joined onto it makes the first probe overflow: the server's
|
||
# role-only re-probe reports no per-key workflow.
|
||
"first_probe_overflows": _mdta_items(
|
||
[b"workflow", b"prompt"],
|
||
[(1, oversized), (1, b"A" * (8 * MiB)), (2, graph)],
|
||
),
|
||
}
|
||
files = {}
|
||
for name, udta in shapes.items():
|
||
files[name] = self.directory / f"{name}.mp4"
|
||
files[name].write_bytes(_with_udta(source, udta))
|
||
# One byte of core.mkv's PROMPT TagString patched to 0xFF.
|
||
mkv = bytearray((VIDEO_FIXTURES / "core.mkv").read_bytes())
|
||
at = mkv.index(b'"KSampler"', mkv.index(b"PROMPT")) + 1
|
||
mkv[at] = 0xFF
|
||
files["matroska_invalid_byte"] = self.directory / "invalid-byte.mkv"
|
||
files["matroska_invalid_byte"].write_bytes(bytes(mkv))
|
||
observed = _run_oracle(
|
||
self, [{"name": name, "path": str(path)} for name, path in files.items()]
|
||
)
|
||
for name, path in files.items():
|
||
with self.subTest(shape=name):
|
||
server = read_inspection_file(path)
|
||
self.assertEqual(
|
||
_local_shape(observed[name]["result"]["selection"]),
|
||
_read_shape(server),
|
||
)
|
||
metadata = _metadata(server)
|
||
self.assertEqual(
|
||
_canonical_bytes(observed[name]["metadata"]),
|
||
_canonical_bytes(metadata),
|
||
)
|
||
self.assertEqual(metadata["status"], "available")
|
||
# The re-probe has its own bound: an 8 MiB envelope beside the oversized
|
||
# workflow fails the server's read and stops the local one.
|
||
wide = self.directory / "reprobe_overflows.mp4"
|
||
wide.write_bytes(
|
||
_with_udta(
|
||
source,
|
||
_mdta_items(
|
||
[b"workflow"],
|
||
[(1, oversized)],
|
||
legacy=[("\u00a9cmt", b'{"n":"' + b"x" * (8 * MiB) + b'"}')],
|
||
),
|
||
)
|
||
)
|
||
local = _run_oracle(self, [{"name": "wide", "path": str(wide)}])["wide"]
|
||
self.assertEqual(
|
||
local["result"]["read"], {"status": "limited", "reason": "metadata-limit"}
|
||
)
|
||
with self.assertRaises(MetadataReadError):
|
||
read_inspection_file(wide)
|
||
# The server outcomes these shapes pin.
|
||
latin1 = read_inspection_file(files["latin1_byte_in_workflow"])
|
||
self.assertEqual(latin1.metadata["workflow"], '{"a":"caf\ufffd"}')
|
||
self.assertEqual(
|
||
read_inspection_file(files["stray_ilst_indexes"]).metadata,
|
||
{"prompt": GRAPH, "workflow": WORKFLOW},
|
||
)
|
||
self.assertEqual(
|
||
read_inspection_file(files["second_keys_box"]).metadata, {"prompt": GRAPH}
|
||
)
|
||
patched = read_inspection_file(files["matroska_invalid_byte"])
|
||
self.assertIn('"\ufffdSampler"', patched.metadata["prompt"])
|
||
over = read_inspection_file(files["oversized_workflow_beside_envelope"])
|
||
self.assertEqual(
|
||
(over.record_formats, over.video_carrier_text),
|
||
({"workflow": "limited"}, None),
|
||
)
|
||
dropped = read_inspection_file(files["first_probe_overflows"])
|
||
self.assertEqual(
|
||
(dropped.metadata, dropped.record_formats),
|
||
({"prompt": GRAPH}, {"prompt": "original-text", "workflow": "limited"}),
|
||
)
|
||
|
||
|
||
def _canonical(value):
|
||
"""``value`` with every dict's keys sorted, so two orders compare as one."""
|
||
if isinstance(value, dict):
|
||
return {key: _canonical(value[key]) for key in sorted(value)}
|
||
if isinstance(value, list):
|
||
return [_canonical(item) for item in value]
|
||
return value
|
||
|
||
|
||
def _canonical_bytes(value):
|
||
return _bytes(_canonical(value))
|
||
|
||
|
||
def _server_metadata(path):
|
||
try:
|
||
return _metadata(read_inspection_file(path))
|
||
except MetadataReadError:
|
||
return _metadata(None)
|
||
|
||
|
||
def _synthetic_read(case):
|
||
if case["records"] is None:
|
||
return None
|
||
formats = case.get("record_formats")
|
||
if formats is None:
|
||
formats = {key: "original-text" for key in case["records"]}
|
||
return InspectionRead(
|
||
case["records"],
|
||
formats,
|
||
case.get("carrier_text"),
|
||
None,
|
||
None,
|
||
None,
|
||
frozenset(case.get("limited", ())),
|
||
)
|
||
|
||
|
||
def _case_supported(case):
|
||
"""Whether this interpreter has the feature a case depends on."""
|
||
if case.get("needs") == "int-max-str-digits":
|
||
limit = getattr(sys, "get_int_max_str_digits", None)
|
||
return limit is not None and limit() == 4300
|
||
return True
|
||
|
||
|
||
def _node_ceiling_cases():
|
||
"""Record sets at and past the 20,000-node ceiling, too large to commit."""
|
||
api = lambda count: "{%s}" % ",".join( # noqa: E731
|
||
'"%d":{"class_type":"A","inputs":{}}' % index for index in range(count)
|
||
)
|
||
ui = lambda count: "[%s]" % ",".join("{}" for _ in range(count)) # noqa: E731
|
||
half = GRAPH_NODE_CEILING // 2
|
||
return [
|
||
{"name": "prompt_at_ceiling", "records": {"prompt": api(GRAPH_NODE_CEILING)}},
|
||
{
|
||
"name": "prompt_over_ceiling",
|
||
"records": {"prompt": api(GRAPH_NODE_CEILING + 1)},
|
||
},
|
||
{
|
||
"name": "workflow_at_ceiling",
|
||
"records": {"workflow": '{"nodes":%s}' % ui(GRAPH_NODE_CEILING)},
|
||
},
|
||
{
|
||
"name": "workflow_over_ceiling",
|
||
"records": {"workflow": '{"nodes":%s}' % ui(GRAPH_NODE_CEILING + 1)},
|
||
},
|
||
{
|
||
"name": "workflow_subgraphs_over_ceiling",
|
||
"records": {
|
||
"workflow": '{"nodes":%s,"definitions":{"subgraphs":[{"nodes":%s},'
|
||
'5,{"nodes":3}]}}' % (ui(half + 1), ui(half))
|
||
},
|
||
},
|
||
{
|
||
"name": "workflow_subgraphs_at_ceiling",
|
||
"records": {
|
||
"workflow": '{"nodes":%s,"definitions":{"subgraphs":[{"nodes":%s}]}}'
|
||
% (ui(half), ui(half))
|
||
},
|
||
},
|
||
{
|
||
"name": "carrier_prompt_over_ceiling",
|
||
"records": {"prompt": "x"},
|
||
"record_formats": {"prompt": "reserialized-json"},
|
||
"carrier_text": '{"prompt":%s}' % api(GRAPH_NODE_CEILING + 1),
|
||
},
|
||
]
|
||
|
||
|
||
class LocalProjectionParityTests(unittest.TestCase):
|
||
"""The local projection answers byte for byte as ``_metadata`` does."""
|
||
|
||
def setUp(self):
|
||
temporary = tempfile.TemporaryDirectory()
|
||
self.addCleanup(temporary.cleanup)
|
||
self.addCleanup(assert_closed, self, temporary.name)
|
||
self.directory = Path(temporary.name)
|
||
|
||
def _files(self):
|
||
files = {}
|
||
for name in parity.COMMITTED:
|
||
files[name] = parity.fixture_path(name)
|
||
for name in parity.TEST_TIME:
|
||
path = parity.fixture_path(name, self.directory)
|
||
parity.write_png(path, parity.build(name))
|
||
files[name] = path
|
||
for name, write in EXTRA_SHAPES.items():
|
||
path = self.directory / f"{name}.png"
|
||
write(path)
|
||
files[name] = path
|
||
for header in ("tiny.jpg", "tiny.webp"):
|
||
files[header] = HEADER_FIXTURES / header
|
||
for name in RECORDED_TAGS:
|
||
files[name] = VIDEO_FIXTURES / name
|
||
return files
|
||
|
||
def test_local_metadata_matches_route_for_parity_corpus(self):
|
||
files = self._files()
|
||
cases = [
|
||
case
|
||
for case in PROJECTION_CASES + _node_ceiling_cases()
|
||
if _case_supported(case)
|
||
]
|
||
self.assertGreater(len(cases), len(PROJECTION_CASES) - 3)
|
||
entries = [{"name": name, "path": str(path)} for name, path in files.items()]
|
||
entries += [
|
||
{
|
||
"name": "case:" + case["name"],
|
||
"records": case["records"],
|
||
"record_formats": case.get("record_formats"),
|
||
"carrier_text": case.get("carrier_text"),
|
||
"limited": case.get("limited", []),
|
||
}
|
||
for case in cases
|
||
]
|
||
observed = _run_oracle(self, entries)
|
||
outcomes = set()
|
||
for name, path in files.items():
|
||
with self.subTest(fixture=name):
|
||
self.assertIsNotNone(observed[name]["classified"], name)
|
||
if path.suffix in VIDEO_CONTAINERS and not HAVE_FFMPEG:
|
||
# Without ffprobe the server cannot read the file; the tag-map
|
||
# vectors below cover the carrier projection instead.
|
||
continue
|
||
server = _server_metadata(path)
|
||
local = observed[name]["metadata"]
|
||
self.assertEqual(_canonical_bytes(local), _canonical_bytes(server))
|
||
outcomes.add(server["status"])
|
||
for field in ("prompt_graph", "workflow_graph"):
|
||
outcomes.add(field + ":" + server[field]["status"])
|
||
for case in cases:
|
||
with self.subTest(case=case["name"]):
|
||
server = _metadata(_synthetic_read(case))
|
||
local = observed["case:" + case["name"]]["metadata"]
|
||
self.assertEqual(_canonical_bytes(local), _canonical_bytes(server))
|
||
expect = case.get("expect") or {}
|
||
for field, status in expect.items():
|
||
self.assertEqual(server[field]["status"], status, field)
|
||
outcomes.add(server["status"])
|
||
for field in ("prompt_graph", "workflow_graph"):
|
||
outcomes.add(field + ":" + server[field]["status"])
|
||
# Every projection outcome must occur, or agreement proves little.
|
||
self.assertEqual(
|
||
outcomes,
|
||
{
|
||
"available",
|
||
"absent",
|
||
"unreadable",
|
||
"prompt_graph:parsed",
|
||
"prompt_graph:limited",
|
||
"prompt_graph:malformed",
|
||
"prompt_graph:absent",
|
||
"workflow_graph:parsed",
|
||
"workflow_graph:limited",
|
||
"workflow_graph:malformed",
|
||
"workflow_graph:absent",
|
||
},
|
||
)
|
||
|
||
def test_recorded_video_tag_maps_project_as_the_server(self):
|
||
"""The carrier projection on the recorded maps, with or without ffprobe."""
|
||
vectors = [(name, tags, []) for name, tags in RECORDED_TAGS.items()]
|
||
vectors += [
|
||
(name, tags, limited)
|
||
for name, tags, limited in TAG_MAP_VECTORS
|
||
if name
|
||
in {
|
||
"per_key_wins_over_comment",
|
||
"object_prompt_reserialized",
|
||
"python_number_spellings",
|
||
"integer_like_keys_keep_order",
|
||
"invalid_envelope_workflow_keeps_roles",
|
||
"workflow_only_invalid_envelope",
|
||
"reader_limited_workflow_with_per_key",
|
||
"envelope_workflow_over_budget_limited",
|
||
}
|
||
]
|
||
observed = _run_oracle(
|
||
self,
|
||
[
|
||
{"name": name, "tags": tags, "limited": limited}
|
||
for name, tags, limited in vectors
|
||
],
|
||
)
|
||
seen = set()
|
||
for name, tags, limited in vectors:
|
||
with self.subTest(vector=name):
|
||
with (
|
||
patch(
|
||
"utils.prompt_inspector.ffprobe_path",
|
||
return_value="/tool/ffprobe",
|
||
),
|
||
patch(
|
||
"utils.prompt_inspector._probe_video_tags",
|
||
side_effect=_fake_probe(
|
||
tags, drop_workflow="workflow" in limited
|
||
),
|
||
),
|
||
):
|
||
try:
|
||
server = _metadata(read_inspection_file(Path("vector.mp4")))
|
||
except MetadataReadError:
|
||
server = _metadata(None)
|
||
local = observed[name]["metadata"]
|
||
self.assertEqual(_canonical_bytes(local), _canonical_bytes(server))
|
||
seen.update(server["record_formats"].values())
|
||
seen.add(server["workflow_graph"]["status"])
|
||
self.assertTrue(
|
||
{"reserialized-json", "original-text", "invalid", "limited"} <= seen
|
||
)
|
||
self.assertTrue({"parsed", "malformed", "limited"} <= seen)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|