Files
opensquilla--opensquilla/tests/test_attachment_workspace.py
2026-07-13 13:12:33 +08:00

240 lines
8.6 KiB
Python

from __future__ import annotations
import hashlib
from pathlib import Path
from opensquilla.attachment_refs import make_attachment_ref, write_transcript_material
from opensquilla.attachment_workspace import (
AttachmentWorkspaceMaterializer,
is_materializable_attachment_mime,
)
_MATERIALIZABLE_MIMES = frozenset({"application/pdf", "text/plain"})
def test_unsupported_mime_is_not_materialized(tmp_path: Path) -> None:
media_root = tmp_path / "media"
workspace = tmp_path / "workspace"
materializer = AttachmentWorkspaceMaterializer(
media_root=media_root,
workspace_dir=workspace,
materializable_mimes=_MATERIALIZABLE_MIMES,
)
result = materializer.materialize_bytes(
b"not an image",
name="photo.png",
mime="image/png",
session_id="session-a",
)
assert result.available is False
assert result.error == "attachment type is not materializable"
assert not (workspace / ".opensquilla").exists()
assert not is_materializable_attachment_mime("image/png", _MATERIALIZABLE_MIMES)
def test_materializes_transcript_ref_inside_workspace(tmp_path: Path) -> None:
media_root = tmp_path / "media"
workspace = tmp_path / "workspace"
payload = b"%PDF-1.4\nminimal\n%%EOF\n"
sha, _path, _wrote = write_transcript_material(
media_root=media_root,
session_id="session-a",
payload=payload,
)
ref = make_attachment_ref(
sha256=sha,
name="../../report.pdf",
mime="application/pdf",
size=len(payload),
session_id="session-a",
source="transcript",
)
result = AttachmentWorkspaceMaterializer(
media_root=media_root,
workspace_dir=workspace,
materializable_mimes=_MATERIALIZABLE_MIMES,
).materialize(ref)
assert result.available is True
assert result.rel_path is not None
assert result.rel_path.startswith(".opensquilla/attachments/session-a/")
assert ".." not in result.rel_path
materialized = (workspace / result.rel_path).resolve()
materialized.relative_to(workspace.resolve())
assert materialized.read_bytes() == payload
assert materialized.name == f"{sha[:12]}-report.pdf"
def test_existing_materialized_file_is_reused_when_hash_matches(tmp_path: Path) -> None:
workspace = tmp_path / "workspace"
payload = b"hello,world\n"
sha = hashlib.sha256(payload).hexdigest()
rel_dir = workspace / ".opensquilla" / "attachments" / "session-a"
rel_dir.mkdir(parents=True)
existing = rel_dir / f"{sha[:12]}-notes.txt"
existing.write_bytes(payload)
before_mtime_ns = existing.stat().st_mtime_ns
result = AttachmentWorkspaceMaterializer(
media_root=tmp_path / "media",
workspace_dir=workspace,
materializable_mimes=_MATERIALIZABLE_MIMES,
).materialize_bytes(
payload,
name="notes.txt",
mime="text/plain",
session_id="session-a",
)
assert result.available is True
assert existing.stat().st_mtime_ns == before_mtime_ns
assert existing.read_bytes() == payload
def _budgeted(tmp_path: Path, budget: int | None) -> AttachmentWorkspaceMaterializer:
return AttachmentWorkspaceMaterializer(
media_root=tmp_path / "media",
workspace_dir=tmp_path / "workspace",
materializable_mimes=None,
disk_budget_bytes=budget,
)
def test_budget_rejects_materialization_that_would_exceed_it(tmp_path: Path) -> None:
materializer = _budgeted(tmp_path, budget=10)
result = materializer.materialize_bytes(
b"x" * 11,
name="big.bin",
mime="application/octet-stream",
session_id="session-a",
)
assert result.available is False
assert result.error is not None
assert "workspace attachment budget exceeded" in result.error
# The marker text names the remedy for the operator/model to see.
assert "workspace_attachment_disk_budget_bytes" in result.error
files = list((tmp_path / "workspace" / ".opensquilla" / "attachments").rglob("*-big.bin"))
assert files == []
def test_budget_counts_existing_workspace_files(tmp_path: Path) -> None:
materializer = _budgeted(tmp_path, budget=16)
first = materializer.materialize_bytes(
b"a" * 10, name="a.bin", mime="application/octet-stream", session_id="session-a"
)
assert first.available is True
# A fresh instance re-scans the tree, so the 10 existing bytes count.
second = _budgeted(tmp_path, budget=16).materialize_bytes(
b"b" * 10, name="b.bin", mime="application/octet-stream", session_id="session-a"
)
assert second.available is False
assert "workspace attachment budget exceeded" in (second.error or "")
# A smaller payload that fits the remaining headroom is admitted.
third = _budgeted(tmp_path, budget=16).materialize_bytes(
b"c" * 6, name="c.bin", mime="application/octet-stream", session_id="session-a"
)
assert third.available is True
def test_budget_reuse_of_existing_file_is_always_free(tmp_path: Path) -> None:
payload = b"d" * 12
first = _budgeted(tmp_path, budget=12).materialize_bytes(
payload, name="d.bin", mime="application/octet-stream", session_id="session-a"
)
assert first.available is True
# At-budget workspace: re-materializing the SAME content must stay
# available (reuse short-circuits before the budget check) so replay of
# already-materialized history never degrades when the budget fills.
again = _budgeted(tmp_path, budget=12).materialize_bytes(
payload, name="d.bin", mime="application/octet-stream", session_id="session-a"
)
assert again.available is True
assert again.rel_path == first.rel_path
def test_budget_none_is_unbounded(tmp_path: Path) -> None:
result = _budgeted(tmp_path, budget=None).materialize_bytes(
b"e" * 4096, name="e.bin", mime="application/octet-stream", session_id="session-a"
)
assert result.available is True
def test_budget_shared_across_one_instance_batch(tmp_path: Path) -> None:
# One materializer instance (one turn) tracks its own writes against the
# budget without re-walking the tree.
materializer = _budgeted(tmp_path, budget=16)
first = materializer.materialize_bytes(
b"f" * 10, name="f.bin", mime="application/octet-stream", session_id="session-a"
)
second = materializer.materialize_bytes(
b"g" * 10, name="g.bin", mime="application/octet-stream", session_id="session-a"
)
assert first.available is True
assert second.available is False
def test_workspace_attachment_budget_from_config_guards() -> None:
from types import SimpleNamespace
from opensquilla.attachment_workspace import workspace_attachment_budget_from_config
good = SimpleNamespace(
attachments=SimpleNamespace(workspace_attachment_disk_budget_bytes=123)
)
assert workspace_attachment_budget_from_config(good) == 123
assert workspace_attachment_budget_from_config(None) is None
assert (
workspace_attachment_budget_from_config(
SimpleNamespace(
attachments=SimpleNamespace(workspace_attachment_disk_budget_bytes=0)
)
)
is None
)
assert (
workspace_attachment_budget_from_config(
SimpleNamespace(
attachments=SimpleNamespace(workspace_attachment_disk_budget_bytes="2")
)
)
is None
)
def test_budget_overwrite_of_mismatched_file_frees_replaced_bytes(tmp_path: Path) -> None:
workspace = tmp_path / "workspace"
payload = b"h" * 10
sha = hashlib.sha256(payload).hexdigest()
target_dir = workspace / ".opensquilla" / "attachments" / "session-a"
target_dir.mkdir(parents=True)
stale = target_dir / f"{sha[:12]}-h.bin"
stale.write_bytes(b"stale-different-content-12345") # 29 bytes at the target path
materializer = _budgeted(tmp_path, budget=16)
# 29 stale bytes alone exceed the budget, but the overwrite frees them:
# (29 - 29) + 10 <= 16 must be admitted.
result = materializer.materialize_bytes(
payload, name="h.bin", mime="application/octet-stream", session_id="session-a"
)
assert result.available is True
assert stale.read_bytes() == payload
# Cached usage after the overwrite must be 10 (not 29 or 39): a 6-byte
# payload fits (10 + 6 <= 16), one more byte does not.
ok = materializer.materialize_bytes(
b"i" * 6, name="i.bin", mime="application/octet-stream", session_id="session-a"
)
assert ok.available is True
over = materializer.materialize_bytes(
b"j", name="j.bin", mime="application/octet-stream", session_id="session-a"
)
assert over.available is False