233 lines
8.3 KiB
Python
233 lines
8.3 KiB
Python
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
|
|
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_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.path(ref["digest"]) is not None]
|
|
assert len(held) == 3
|
|
# The newest three, in order: eviction is by age, not by chance.
|
|
assert held == refs[2:]
|
|
|
|
|
|
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 resolves a digest — serving one, checking a run's input,
|
|
a panel's scope — goes through ``path``, 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.path(ref["digest"]) is not None
|
|
# In the ring, and not on the volume the sweep is about.
|
|
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.collect(set())
|
|
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.path(digest) is not None
|
|
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_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"<script>"
|
|
).json()["digest"]
|
|
image = client.put(
|
|
f"{url}?name=x.png", headers=superuser_token_headers, content=b"\x89PNG"
|
|
).json()["digest"]
|
|
|
|
def served(digest: str, media_type: str) -> dict[str, str]:
|
|
answer = client.get(
|
|
f"{url}/{digest}?media_type={media_type}", headers=superuser_token_headers
|
|
)
|
|
assert answer.status_code == 200
|
|
return dict(answer.headers)
|
|
|
|
for media_type in ("text/html", "image/svg+xml", "application/xml"):
|
|
headers = served(page, media_type)
|
|
assert headers["content-security-policy"] == "sandbox"
|
|
assert headers["x-content-type-options"] == "nosniff"
|
|
assert "content-security-policy" not in served(image, "image/png")
|