234 lines
8.6 KiB
Python
234 lines
8.6 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import sqlite3
|
|
from collections.abc import Iterator
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from game_data_capture.pipeline.jobs import JobLeases
|
|
from game_data_capture.storage.blobs import ContentAddressedBlobStore
|
|
from game_data_capture.storage.database import StorageDatabase
|
|
from game_data_capture.storage.repository import CaptureRepository, StorageError
|
|
|
|
|
|
@pytest.fixture
|
|
def storage(tmp_path: Path) -> tuple[StorageDatabase, CaptureRepository]:
|
|
root = (tmp_path / "utility").resolve()
|
|
migrations = Path(__file__).parents[3] / "migrations"
|
|
database = StorageDatabase(root, migrations.resolve())
|
|
database.migrate()
|
|
return database, CaptureRepository(database, ContentAddressedBlobStore(root))
|
|
|
|
|
|
def test_identical_occurrences_are_deduplicated_and_receipts_survive_restart(
|
|
storage: tuple[StorageDatabase, CaptureRepository],
|
|
) -> None:
|
|
database, repository = storage
|
|
assert (
|
|
repository.receive(
|
|
"capture-1", "receipt-1", [b"same"], media_type="image/png", max_bytes=100
|
|
)
|
|
== "capture-1"
|
|
)
|
|
assert (
|
|
repository.receive(
|
|
"capture-2", "receipt-2", [b"same"], media_type="image/png", max_bytes=100
|
|
)
|
|
== "capture-2"
|
|
)
|
|
# A simulated restart uses fresh repository objects and a repeated intake receipt.
|
|
restarted = CaptureRepository(database, ContentAddressedBlobStore(database.root))
|
|
assert (
|
|
restarted.receive(
|
|
"capture-other", "receipt-2", [b"same"], media_type="image/png", max_bytes=100
|
|
)
|
|
== "capture-2"
|
|
)
|
|
assert restarted.report()["blobs"] == 1
|
|
assert restarted.report()["captures"] == 2
|
|
|
|
|
|
def test_capture_never_references_missing_blob_and_cas_rejects_stale(
|
|
storage: tuple[StorageDatabase, CaptureRepository],
|
|
) -> None:
|
|
database, repository = storage
|
|
repository.receive(
|
|
"capture-1", "receipt-1", [b"complete"], media_type="image/png", max_bytes=100
|
|
)
|
|
row = database.connect().execute("SELECT blob_sha256 FROM captures").fetchone()
|
|
assert row is not None and repository.blobs.verify(str(row["blob_sha256"])) == len(b"complete")
|
|
assert repository.transition("capture-1", "queued", 0) == 1
|
|
with pytest.raises(StorageError):
|
|
repository.transition("capture-1", "failed", 0)
|
|
|
|
|
|
def test_reprocessing_is_append_only_and_leases_enforce_owner(
|
|
storage: tuple[StorageDatabase, CaptureRepository],
|
|
) -> None:
|
|
database, repository = storage
|
|
repository.receive(
|
|
"capture-1", "receipt-1", [b"capture"], media_type="image/png", max_bytes=100
|
|
)
|
|
repository.append_ocr_run("run-1", "capture-1", "ocr", "1", {"raw": "one"})
|
|
repository.append_ocr_run("run-2", "capture-1", "ocr", "2", {"raw": "two"})
|
|
assert database.connect().execute("SELECT count(*) FROM ocr_runs").fetchone()[0] == 2
|
|
jobs = JobLeases(database)
|
|
jobs.create("job-1", "ocr", "capture-1")
|
|
assert jobs.acquire("worker-a", 30) == "job-1"
|
|
with pytest.raises(StorageError):
|
|
jobs.finish("job-1", "worker-b", "succeeded")
|
|
jobs.finish("job-1", "worker-a", "succeeded")
|
|
|
|
|
|
def test_replayed_receipt_does_not_consume_payload(
|
|
storage: tuple[StorageDatabase, CaptureRepository],
|
|
) -> None:
|
|
database, repository = storage
|
|
repository.receive(
|
|
"capture-1", "receipt-1", [b"capture"], media_type="image/png", max_bytes=100
|
|
)
|
|
|
|
consumed = [False]
|
|
|
|
class MustNotConsume:
|
|
def __iter__(self) -> Iterator[bytes]:
|
|
consumed[0] = True
|
|
raise AssertionError("receipt replay consumed payload")
|
|
|
|
assert (
|
|
repository.receive(
|
|
"capture-2", "receipt-1", MustNotConsume(), media_type="image/png", max_bytes=100
|
|
)
|
|
== "capture-1"
|
|
)
|
|
assert not consumed[0]
|
|
|
|
|
|
def test_every_provenance_kind_is_insert_only(
|
|
storage: tuple[StorageDatabase, CaptureRepository],
|
|
) -> None:
|
|
database, repository = storage
|
|
repository.receive(
|
|
"capture-1", "receipt-1", [b"capture"], media_type="image/png", max_bytes=100
|
|
)
|
|
repository.append_ocr_run("run-1", "capture-1", "ocr", "1", {})
|
|
repository.append_region("region-1", "run-1", {})
|
|
repository.append_classification("class-1", "run-1", {})
|
|
repository.append_candidate("candidate-1", "run-1", {})
|
|
repository.append_routing("route-1", "candidate-1", {})
|
|
repository.append_review("review-1", "candidate-1", {})
|
|
repository.append_lineage("lineage-1", "child-1", "parent-1", "derived")
|
|
repository.append_export("export-1", "ready", {})
|
|
with pytest.raises(sqlite3.IntegrityError):
|
|
repository.append_ocr_run("run-1", "capture-1", "ocr", "1", {})
|
|
payloads = {
|
|
"ocr_runs": "payload_json",
|
|
"regions": "region_json",
|
|
"classifications": "classification_json",
|
|
"extraction_candidates": "candidate_json",
|
|
"routing_decisions": "decision_json",
|
|
"review_events": "event_json",
|
|
"exports": "manifest_json",
|
|
}
|
|
with database.connect() as connection:
|
|
for table in (
|
|
"ocr_runs",
|
|
"regions",
|
|
"classifications",
|
|
"extraction_candidates",
|
|
"routing_decisions",
|
|
"review_events",
|
|
"lineage",
|
|
"exports",
|
|
):
|
|
assert connection.execute(f"SELECT count(*) FROM {table}").fetchone()[0] == 1
|
|
for table, column in payloads.items():
|
|
json.loads(connection.execute(f"SELECT {column} FROM {table}").fetchone()[0])
|
|
assert (
|
|
connection.execute("SELECT count(*) FROM ocr_runs WHERE run_id='run-1'").fetchone()[0]
|
|
== 1
|
|
)
|
|
|
|
|
|
def test_failure_sanitization_and_provenance_bounds(
|
|
storage: tuple[StorageDatabase, CaptureRepository],
|
|
) -> None:
|
|
database, repository = storage
|
|
repository.receive(
|
|
"capture-1", "receipt-1", [b"capture"], media_type="image/png", max_bytes=100
|
|
)
|
|
with pytest.raises(ValueError):
|
|
repository.append_ocr_run("run", "capture-1", "", "1", {})
|
|
with pytest.raises(ValueError):
|
|
repository.append_ocr_run("run", "capture-1", "x" * 101, "1", {})
|
|
with pytest.raises(ValueError):
|
|
repository.append_lineage("line", "child", "parent", "")
|
|
with pytest.raises(ValueError):
|
|
repository.append_export("export", "", {})
|
|
with pytest.raises(ValueError):
|
|
repository.append_candidate("candidate", "run", "x" * 1_000_001)
|
|
long_token = "x" * 600
|
|
short_message = (
|
|
"sk-secret123 Authorization: Bearer auth-secret token=token-secret secret=hidden "
|
|
"password=pw-secret api_key=key-secret /home/agent/private.png\n"
|
|
)
|
|
repository.record_failure("failure-1", "capture", "capture-1", "failure-code", short_message)
|
|
repository.record_failure("failure-2", "capture", "capture-1", "failure-code", long_token)
|
|
with database.connect() as connection:
|
|
rows = list(
|
|
connection.execute("SELECT code, message FROM runtime_failures ORDER BY failure_id")
|
|
)
|
|
assert len(rows) == 2
|
|
for row in rows:
|
|
assert len(row["message"]) <= 512 and len(row["code"]) <= 64
|
|
assert "\n" not in row["message"]
|
|
for value in (
|
|
"sk-secret123",
|
|
"auth-secret",
|
|
"token-secret",
|
|
"hidden",
|
|
"pw-secret",
|
|
"key-secret",
|
|
"/home/agent/private.png",
|
|
"\n",
|
|
):
|
|
assert value not in rows[0]["message"]
|
|
assert long_token not in rows[1]["message"]
|
|
with pytest.raises(ValueError):
|
|
repository.record_failure("failure-2", "BAD", "capture-1", "code", "message")
|
|
|
|
|
|
def test_complete_blob_after_capture_insert_failure_is_purgeable_orphan(
|
|
storage: tuple[StorageDatabase, CaptureRepository],
|
|
) -> None:
|
|
database, repository = storage
|
|
repository.receive(
|
|
"capture-1", "receipt-1", [b"payload-a"], media_type="image/png", max_bytes=100
|
|
)
|
|
payload_b = b"payload-b"
|
|
digest_b = hashlib.sha256(payload_b).hexdigest()
|
|
with pytest.raises(sqlite3.IntegrityError):
|
|
repository.receive(
|
|
"capture-1", "receipt-2", [payload_b], media_type="image/png", max_bytes=100
|
|
)
|
|
assert repository.blobs.verify(digest_b) == len(payload_b)
|
|
with database.connect() as connection:
|
|
assert (
|
|
connection.execute(
|
|
"SELECT count(*) FROM blobs WHERE sha256 = ?", (digest_b,)
|
|
).fetchone()[0]
|
|
== 0
|
|
)
|
|
assert (
|
|
connection.execute(
|
|
"SELECT count(*) FROM captures WHERE blob_sha256 = ?", (digest_b,)
|
|
).fetchone()[0]
|
|
== 0
|
|
)
|
|
assert repository.blobs._path(digest_b).relative_to(
|
|
repository.blobs.root
|
|
) in repository.blobs.purge_dry_run(set())
|