tests_python/test_codebase_archive.py

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"]