import os import time from datetime import UTC, datetime import pytest from fastapi.testclient import TestClient from sqlmodel import Session, col, select from fluksio.core.config import settings from fluksio.core.db import engine as db_engine from fluksio.flow.artifacts import ArtifactStore, VolatileStore from fluksio.flow.messages import MessageSpec from fluksio.flow.runs import new_run_id, sweep_artifacts from fluksio.flow.schemas import FlowDef, FlowInput, NodeDef from fluksio.flow.state import MemoryState from fluksio.flow.store import FlowStore from fluksio.models import Run, RunArtifact @pytest.fixture def store(tmp_path) -> ArtifactStore: return ArtifactStore(tmp_path / "artifacts") def _run_row(status: str, digest: str = "") -> str: """A run, and optionally the artifact it recorded. Returns its id.""" run_id = new_run_id() with Session(db_engine) as session: session.add( Run(id=run_id, flow="f", status=status, created_at=datetime.now(UTC)) ) if digest: session.add( RunArtifact( run_id=run_id, name="f.out", node="n", digest=digest, size=3, ) ) session.commit() return run_id def _forget(run_id: str) -> None: with Session(db_engine) as session: rows = session.exec( select(RunArtifact).where(col(RunArtifact.run_id) == run_id) ).all() for row in rows: session.delete(row) session.delete(session.get(Run, run_id)) session.commit() def test_the_grace_window_spares_a_fresh_artifact(store): ref = store.put([b"new"]) assert store.collect(set(), grace_s=3600) == 0 assert store.path(ref["digest"]) is not None assert store.collect(set(), grace_s=0) == 1 assert store.path(ref["digest"]) is None def test_a_sweep_keeps_what_a_run_recorded_and_what_a_message_holds(store): recorded = store.put([b"kept by a run"]) held = store.put([b"kept by a message"]) nested = store.put([b"kept inside a payload"]) orphan = store.put([b"referred to by nothing"]) run_id = _run_row("ok", recorded["digest"]) state = MemoryState() state.update( { "cam.frame": held, # A reference inside a json payload counts as much as a bare one. "cam.report": {"latest": {"clip": nested}}, "cam.count": 3, } ) try: assert sweep_artifacts(store, state) == 1 finally: _forget(run_id) assert store.path(orphan["digest"]) is None for kept in (recorded, held, nested): assert store.path(kept["digest"]) is not None def test_a_sweep_stands_aside_while_a_run_is_in_flight(store): orphan = store.put([b"mid-run checkpoint"]) run_id = _run_row("running") try: assert sweep_artifacts(store, MemoryState()) == 0 finally: _forget(run_id) assert store.path(orphan["digest"]) is not None def test_a_streamed_chunk_falls_out_once_the_message_moves_on(store): """What makes a media stream affordable: only the current frame is held.""" state = MemoryState() first = store.put([b"frame one"], media_type="image/png") state.update({"cam.frame": first}) assert sweep_artifacts(store, state) == 0 second = store.put([b"frame two"], media_type="image/png") state.update({"cam.frame": second}) # Old enough to be swept: the sweep only spares what grace covers. os.utime(store.path(first["digest"]), (time.time() - 10, time.time() - 10)) assert sweep_artifacts(store, state, grace_s=5) == 1 assert store.path(first["digest"]) is None assert store.path(second["digest"]) is not None def test_a_sweep_keeps_what_a_flow_input_starts_with(store, tmp_path): """A flow that has never run holds its `initial` in no message yet.""" published = store.put([b"published initial"]) drafted = store.put([b"draft initial"]) orphan = store.put([b"referred to by nothing"]) flows = FlowStore(tmp_path / "flows") clip = MessageSpec(name="clip", dtype="artifact") flows.write_flow( FlowDef(name="f", inputs=[FlowInput(spec=clip, initial=published)]) ) flows.write_draft( FlowDef(name="g", inputs=[FlowInput(spec=clip, initial=drafted)]), 0 ) assert sweep_artifacts(store, MemoryState(), grace_s=0, flows=flows) == 1 assert store.path(orphan["digest"]) is None for kept in (published, drafted): assert store.path(kept["digest"]) is not None def test_a_sweep_keeps_what_a_node_setting_names(store, tmp_path): """A Service's weights outlive the run that produced them.""" weights = store.put([b"weights"]) orphan = store.put([b"referred to by nothing"]) flows = FlowStore(tmp_path / "flows") params = {"nested": [{"weights": weights["digest"]}]} flows.write_flow(FlowDef(name="f", nodes=[NodeDef(id="n", params=params)])) assert sweep_artifacts(store, MemoryState(), grace_s=0, flows=flows) == 1 assert store.path(orphan["digest"]) is None assert store.path(weights["digest"]) is not None def test_a_ring_drops_the_oldest_once_it_is_full(tmp_path): """What makes a camera affordable: the room is fixed, not the history.""" ring = VolatileStore(tmp_path / "ring", limit_bytes=3000) refs = [ring.put([bytes([i]) * 1000], name=f"{i}.bin") for i in range(5)] held = [ref for ref in refs if ring.has(ref["digest"])] assert len(held) == 3 # The newest three, in order: eviction is by age, not by chance. assert held == refs[2:] assert ring.size == 3000 # Held in memory: nothing reached the directory. assert not any(ring.root.iterdir()) def test_the_newest_frame_stays_whatever_its_size(tmp_path): """A frame larger than the bound is still there for its own publish.""" ring = VolatileStore(tmp_path / "ring", limit_bytes=100) small = ring.put([b"s" * 50]) big = ring.put([b"b" * 500]) assert ring.get(big["digest"]) == b"b" * 500 assert not ring.has(small["digest"]) def test_a_workers_frame_is_read_from_its_directory_once(tmp_path): """A worker in the container writes its frames to the ring's directory.""" import hashlib ring = VolatileStore(tmp_path / "ring", limit_bytes=1_000_000) data = b"a worker's frame" body = hashlib.sha256(data).hexdigest() path = ring.root / body[:2] / body path.parent.mkdir(parents=True) path.write_bytes(data) assert ring.has("sha256:" + body) assert ring.get("sha256:" + body) == data # Held after the first read, so every screen is served from memory. path.unlink() assert ring.get("sha256:" + body) == data def test_a_volatile_reference_resolves_like_any_other(tmp_path): """A frame is an ordinary reference, which is what keeps the rest honest. Everything that asks whether a digest resolves — checking a run's input, a cached result — goes through ``has``, and reading one through ``read``, so the ring has to answer there. """ store = ArtifactStore(tmp_path / "artifacts") store.volatile = VolatileStore(tmp_path / "ring", limit_bytes=1_000_000) ref = store.put([b"frame"], name="f.png", media_type="image/png", volatile=True) assert store.has(ref["digest"]) assert b"".join(store.read(ref["digest"])) == b"frame" # In the ring, and not on the volume the sweep is about. assert store.path(ref["digest"]) is None assert not (store.root / ref["digest"][7:9]).exists() def test_a_recorded_frame_is_copied_out_of_the_ring(tmp_path): """Emitted media is not kept; returned media is.""" store = ArtifactStore(tmp_path / "artifacts") store.volatile = VolatileStore(tmp_path / "ring", limit_bytes=1_000_000) ref = store.put([b"kept"], name="f.png", volatile=True) assert store.adopt(ref["digest"]) is True store.volatile = VolatileStore(tmp_path / "ring2", limit_bytes=1_000_000) assert store.path(ref["digest"]) is not None assert store.adopt("sha256:" + "0" * 64) is False def test_an_upload_may_ask_for_the_ring( client: TestClient, superuser_token_headers: dict[str, str] ) -> None: """A camera on another host publishes over HTTP and wants the ring too.""" store = client.app.state.artifact_store previous = store.volatile store.volatile = VolatileStore( store.root.parent / "test-ring", limit_bytes=1_000_000 ) try: url = f"{settings.API_V1_STR}/artifacts" answer = client.put( f"{url}?name=f.png&media_type=image%2Fpng&volatile=1", headers=superuser_token_headers, content=b"\x89PNG-frame", ) assert answer.status_code == 200, answer.text digest = answer.json()["digest"] assert store.volatile.has(digest) assert not (store.root / digest[7:9]).exists() # And it is served back like anything else, which is what the widget # falls back on when the socket did not push it. served = client.get(f"{url}/{digest}", headers=superuser_token_headers) assert served.status_code == 200 assert served.content == b"\x89PNG-frame" finally: store.volatile = previous def test_a_ring_never_holds_more_than_its_bound_and_one_frame(tmp_path): """Frames of every size, the bound never passed by more than the newest.""" ring = VolatileStore(tmp_path / "ring", limit_bytes=10_000) sizes = [700, 9_000, 30, 12_000, 4_000, 4_000, 2_500, 9_999, 1, 6_000] * 3 for index, size in enumerate(sizes): ring.put([bytes([index % 256]) * size]) assert ring.size <= ring.limit_bytes or len(ring._frames) == 1 assert ring.size <= ring.limit_bytes + size def test_a_volatile_upload_is_held_to_the_frame_cap( client: TestClient, superuser_token_headers: dict[str, str], monkeypatch: pytest.MonkeyPatch, ) -> None: """A frame is held in memory whole, so it is refused past a few MB before any of it is read, and a chunked one as soon as it passes the cap.""" store = client.app.state.artifact_store previous = store.volatile store.volatile = VolatileStore( store.root.parent / "test-ring", limit_bytes=1_000_000 ) monkeypatch.setattr(settings, "MAX_FRAME_BYTES", 1000) url = f"{settings.API_V1_STR}/artifacts?volatile=1&media_type=image%2Fjpeg" try: declared = client.put(url, headers=superuser_token_headers, content=b"x" * 1001) assert declared.status_code == 413 def chunks(): for _ in range(5): yield b"x" * 400 chunked = client.put(url, headers=superuser_token_headers, content=chunks()) assert chunked.status_code == 413 assert store.volatile.size == 0 fits = client.put(url, headers=superuser_token_headers, content=b"x" * 1000) assert fits.status_code == 200, fits.text # The durable store keeps its own, larger cap. kept = client.put( f"{settings.API_V1_STR}/artifacts", headers=superuser_token_headers, content=b"y" * 5000, ) assert kept.status_code == 200, kept.text finally: store.volatile = previous def test_a_type_that_can_run_script_is_served_sandboxed( client: TestClient, superuser_token_headers: dict[str, str] ) -> None: """Any member can upload, and the caller names the type it is served as.""" url = f"{settings.API_V1_STR}/artifacts" page = client.put( f"{url}?name=x.html", headers=superuser_token_headers, content=b"