359 lines
13 KiB
Python
359 lines
13 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, NodeDef
|
|
from fluksio.flow.state import MemoryState
|
|
from fluksio.flow.store import FlowStore
|
|
from fluksio.models import Run, RunArtifact
|
|
from tests.utils.stores import created
|
|
|
|
|
|
@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_id="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)])
|
|
)
|
|
created(flows, FlowDef(name="g", inputs=[FlowInput(spec=clip, initial=drafted)]))
|
|
|
|
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"<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")
|
|
|
|
|
|
def test_a_superuser_deletes_an_artifact_nothing_references(
|
|
client: TestClient,
|
|
superuser_token_headers: dict[str, str],
|
|
normal_user_token_headers: dict[str, str],
|
|
) -> None:
|
|
url = f"{settings.API_V1_STR}/artifacts"
|
|
digest = client.put(
|
|
f"{url}?name=x.bin", headers=superuser_token_headers, content=b"delete me"
|
|
).json()["digest"]
|
|
|
|
refused = client.delete(f"{url}/{digest}", headers=normal_user_token_headers)
|
|
assert refused.status_code == 403
|
|
|
|
run_id = _run_row("ok", digest)
|
|
try:
|
|
held = client.delete(f"{url}/{digest}", headers=superuser_token_headers)
|
|
assert held.status_code == 409
|
|
finally:
|
|
_forget(run_id)
|
|
|
|
gone = client.delete(f"{url}/{digest}", headers=superuser_token_headers)
|
|
assert gone.status_code == 204
|
|
assert (
|
|
client.get(f"{url}/{digest}", headers=superuser_token_headers).status_code
|
|
== 404
|
|
)
|
|
again = client.delete(f"{url}/{digest}", headers=superuser_token_headers)
|
|
assert again.status_code == 404
|