File size: 3,719 Bytes
31c8075 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | import pytest
from arena.codebase_archive import (
ArchiveFile,
CodebaseArchiveSnapshot,
LocalCodebaseArchiveStore,
RedisCodebaseArchiveStore,
archive_from_directory,
archive_from_downloads,
archive_to_zip_bytes,
)
from arena.redis_client import RedisStoreError
class FakeRedisClient:
def __init__(self) -> None:
self.data: dict[str, str] = {}
def execute(self, *args: str):
command = args[0].upper()
if command == "GET":
return self.data.get(args[1])
if command == "SET":
self.data[args[1]] = args[2]
return "OK"
raise AssertionError(f"unexpected command: {args}")
def test_local_codebase_archive_round_trips_workspace_files(tmp_path) -> None:
workspace = tmp_path / "workspace"
workspace.mkdir()
(workspace / "app.py").write_text("print('ok')\n")
store = LocalCodebaseArchiveStore(root_dir=tmp_path / "archives")
pointer = store.save(archive_from_directory("run-1", workspace))
restored = tmp_path / "restored"
store.materialize(pointer, restored)
assert pointer == "local-snapshot://run-1"
assert (restored / "app.py").read_text() == "print('ok')\n"
def test_redis_codebase_archive_uses_codebase_namespace(tmp_path) -> None:
client = FakeRedisClient()
store = RedisCodebaseArchiveStore(client=client)
archive = CodebaseArchiveSnapshot(
archive_id="run-1",
files=[ArchiveFile(path="answer.txt", content_base64="b2s=", size_bytes=2)],
)
pointer = store.save(archive)
restored = tmp_path / "restored"
store.materialize(pointer, restored)
assert pointer == "redis-archive://run-1"
assert "arena:codebase-archive:run-1" in client.data
assert (restored / "answer.txt").read_text() == "ok"
def test_codebase_archive_rejects_unsafe_paths() -> None:
with pytest.raises(ValueError):
CodebaseArchiveSnapshot(
archive_id="run-1",
files=[ArchiveFile(path="../secret.txt", content_base64="", size_bytes=0)],
)
def test_redis_codebase_archive_reports_invalid_payload() -> None:
client = FakeRedisClient()
client.data["arena:codebase-archive:run-1"] = "not-json"
with pytest.raises(RedisStoreError):
RedisCodebaseArchiveStore(client=client).load("redis-archive://run-1")
def test_codebase_archive_can_be_rendered_as_zip_bytes(tmp_path) -> None:
workspace = tmp_path / "workspace"
workspace.mkdir()
(workspace / "app.py").write_text("print('ok')\n")
archive = archive_from_directory("run-1", workspace)
payload = archive_to_zip_bytes(archive)
assert b"PK" in payload[:4]
def test_codebase_archive_excludes_python_cache_files_from_directory(tmp_path) -> None:
workspace = tmp_path / "workspace"
workspace.mkdir()
(workspace / "app.py").write_text("print('ok')\n")
pycache = workspace / "taskledger" / "__pycache__"
pycache.mkdir(parents=True)
(pycache / "cli.cpython-314.pyc").write_bytes(b"bytecode")
(workspace / "taskledger" / "models.pyo").write_bytes(b"optimized bytecode")
archive = archive_from_directory("run-1", workspace)
payload = archive_to_zip_bytes(archive)
assert [file.path for file in archive.files] == ["app.py"]
assert b"cli.cpython" not in payload
def test_codebase_archive_excludes_python_cache_files_from_downloads() -> None:
archive = archive_from_downloads(
"run-1",
[
("README.md", b"# ok\n"),
("taskledger/__pycache__/cli.cpython-314.pyc", b"bytecode"),
("taskledger/storage.pyo", b"optimized bytecode"),
],
)
assert [file.path for file in archive.files] == ["README.md"]
|