tests_python/test_codebase_handoff.py

"""Direct tests for the Codebase handoff deep module."""

from pathlib import Path

from arena.codebase_archive import LocalCodebaseArchiveStore, RedisCodebaseArchiveStore
from arena.codebase_handoff import (
    CodebaseHandoff,
    pointer_error,
)


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}")


class FakeDownload:
    def __init__(self, path: str, content: bytes | None = None, error: str | None = None) -> None:
        self.path = path
        self.content = content
        self.error = error


class FakeGlob:
    def __init__(self, matches: list[dict], error: str | None = None) -> None:
        self.matches = matches
        self.error = error


class FakeBackend:
    def __init__(self) -> None:
        self.files: dict[str, bytes] = {}
        self.expected_glob = "/workspace/**"

    def glob(self, pattern: str) -> FakeGlob:
        assert pattern == self.expected_glob
        return FakeGlob([{"path": path, "is_dir": False} for path in self.files])

    def download_files(self, paths: list[str]) -> list[FakeDownload]:
        return [FakeDownload(path, self.files[path]) for path in paths]

    def upload_files(self, files: list[tuple[str, bytes]]) -> list[FakeDownload]:
        for path, content in files:
            self.files[path] = content
        return [FakeDownload(path=path) for path, _ in files]


def test_handoff_snapshot_local_workspace_round_trip(tmp_path) -> None:
    source = tmp_path / "sandbox"
    (source / "workspace").mkdir(parents=True)
    (source / "workspace" / "main.py").write_text("print('hi')")

    handoff = CodebaseHandoff(root_dir=tmp_path / "snapshots")
    pointer = handoff.snapshot_local_workspace(source, "run-1")

    snapshot_workspace = tmp_path / "snapshots" / "run-1" / "workspace"
    assert pointer == "local-snapshot://run-1"
    assert (snapshot_workspace / "main.py").read_text() == "print('hi')"


def test_handoff_snapshot_backend_workspace(tmp_path) -> None:
    backend = FakeBackend()
    backend.files = {"/workspace/main.py": b"x = 1\n", "/workspace/readme.md": b"# r"}
    handoff = CodebaseHandoff(root_dir=tmp_path / "snapshots")

    pointer = handoff.snapshot_backend_workspace(backend, "run-1")

    snapshot_workspace = tmp_path / "snapshots" / "run-1" / "workspace"
    assert pointer == "local-snapshot://run-1"
    assert (snapshot_workspace / "main.py").read_text() == "x = 1\n"
    assert (snapshot_workspace / "readme.md").read_text() == "# r"


def test_handoff_snapshot_uses_redis_archive_store(tmp_path) -> None:
    redis_client = FakeRedisClient()
    archive_store = RedisCodebaseArchiveStore(client=redis_client)
    backend = FakeBackend()
    backend.files = {"/workspace/main.py": b"x = 1\n"}

    handoff = CodebaseHandoff(archive_store=archive_store, root_dir=tmp_path / "snapshots")
    pointer = handoff.snapshot_backend_workspace(backend, "run-1")

    assert pointer == "redis-archive://run-1"
    assert "arena:codebase-archive:run-1" in redis_client.data


def test_handoff_restore_to_backend_writes_under_workspace_dir(tmp_path) -> None:
    source = tmp_path / "sandbox"
    (source / "workspace").mkdir(parents=True)
    (source / "workspace" / "main.py").write_text("print('hi')")
    handoff = CodebaseHandoff(root_dir=tmp_path / "snapshots")
    pointer = handoff.snapshot_local_workspace(source, "run-1")

    target = FakeBackend()
    count = handoff.restore_to_backend(pointer, target, workspace_dir="/sandbox/work")

    assert count == 1
    assert target.files == {"/sandbox/work/main.py": b"print('hi')"}


def test_handoff_snapshot_rejects_unsafe_codebase_id(tmp_path) -> None:
    handoff = CodebaseHandoff(root_dir=tmp_path / "snapshots")

    try:
        handoff.snapshot_local_workspace(tmp_path, "../escape")
    except ValueError as exc:
        assert "invalid codebase id" in str(exc).lower()
    else:
        raise AssertionError("expected ValueError for unsafe codebase id")


def test_handoff_restore_rejects_unsupported_pointer(tmp_path) -> None:
    handoff = CodebaseHandoff(root_dir=tmp_path / "snapshots")

    try:
        handoff.restore_to_backend("daytona://sandbox/x", FakeBackend(), workspace_dir="/sandbox/work")
    except ValueError as exc:
        assert str(exc) == pointer_error("daytona://sandbox/x")
    else:
        raise AssertionError("expected ValueError for unsupported pointer")


def test_handoff_full_snapshot_store_restore_execute_cycle(tmp_path) -> None:
    """End-to-end behavior: a Validator sandbox can run a check against a
    snapshotted Workspace restored through the handoff deep module."""

    source = tmp_path / "sandbox"
    (source / "workspace").mkdir(parents=True)
    (source / "workspace" / "answer.txt").write_text("ok")

    handoff = CodebaseHandoff(root_dir=tmp_path / "snapshots")
    pointer = handoff.snapshot_local_workspace(source, "run-1")

    target = FakeBackend()
    count = handoff.restore_to_backend(pointer, target, workspace_dir="/vfs/work")

    assert count == 1
    assert target.files == {"/vfs/work/answer.txt": b"ok"}
    assert "answer.txt" in {path.rsplit("/", 1)[-1] for path in target.files}


def test_handoff_restore_is_idempotent_within_one_archive_pointer(tmp_path) -> None:
    source = tmp_path / "sandbox"
    (source / "workspace").mkdir(parents=True)
    (source / "workspace" / "data.txt").write_text("payload")

    handoff = CodebaseHandoff(root_dir=tmp_path / "snapshots")
    pointer = handoff.snapshot_local_workspace(source, "run-1")

    first = FakeBackend()
    handoff.restore_to_backend(pointer, first, workspace_dir="/vfs/work")

    first.files["/vfs/work/extra.txt"] = b"leftover"
    second = FakeBackend()
    handoff.restore_to_backend(pointer, second, workspace_dir="/vfs/work")

    assert "/vfs/work/extra.txt" not in second.files
    assert second.files == {"/vfs/work/data.txt": b"payload"}