tests_python/test_codebase_executor.py

from pathlib import Path

from deepagents.backends import LocalShellBackend

from arena.codebase_executor import (
    CodebaseSandboxExecutor,
    LocalCodebaseExecutor,
)
from arena.codebase_archive import RedisCodebaseArchiveStore
from arena.codebase_handoff import CodebaseHandoff


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 = {
            "/workspace/app.py": b"print('ok')\n",
            "/workspace/pkg/data.txt": b"data",
        }
        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]


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 FakeUploadResult:
    error = None


class ReopenableBackend:
    def __init__(self, name: str, closed: list[str]) -> None:
        self.name = name
        self.closed = closed
        self.uploaded: list[tuple[str, bytes]] = []
        self._validator_workspace = "/workspace"

    def upload_files(self, files: list[tuple[str, bytes]]) -> list[FakeUploadResult]:
        self.uploaded.extend(files)
        return [FakeUploadResult() for _ in files]

    def execute(self, command: str, timeout: int):
        class Result:
            exit_code = 0
            output = ""

        assert command.startswith("cd /workspace && ")
        return Result()

    def delete(self) -> None:
        self.closed.append(self.name)


def test_local_codebase_executor_runs_command_in_workspace(tmp_path) -> None:
    source_root = tmp_path / "sandbox"
    workspace = source_root / "workspace"
    workspace.mkdir(parents=True)
    (workspace / "answer.txt").write_text("ok")
    pointer = CodebaseHandoff(root_dir=tmp_path / "snapshots").snapshot_local_workspace(source_root, "run-1")

    passed, output = LocalCodebaseExecutor(root_dir=tmp_path / "snapshots").run(pointer, "test -f answer.txt")

    assert passed is True
    assert "exit code 0" in output


def test_local_codebase_executor_captures_failed_command(tmp_path) -> None:
    source_root = tmp_path / "sandbox"
    (source_root / "workspace").mkdir(parents=True)
    pointer = CodebaseHandoff(root_dir=tmp_path / "snapshots").snapshot_local_workspace(source_root, "run-1")

    passed, output = LocalCodebaseExecutor(root_dir=tmp_path / "snapshots").run(
        pointer,
        "printf 'bad' >&2; exit 7",
    )

    assert passed is False
    assert "bad" in output
    assert "exit code 7" in output


def test_local_codebase_executor_rejects_unknown_or_unsafe_pointer(tmp_path) -> None:
    executor = LocalCodebaseExecutor(root_dir=tmp_path / "snapshots")

    unsupported = executor.run("daytona://sandbox/x/workspace", "echo no")
    traversal = executor.run("local-snapshot://../secret", "echo no")

    assert unsupported[0] is False
    assert "unsupported archive pointer" in unsupported[1]
    assert traversal[0] is False
    assert "invalid snapshot id" in traversal[1]


def test_snapshot_replaces_previous_copy_without_mutating_source(tmp_path) -> None:
    source_root = tmp_path / "sandbox"
    workspace = source_root / "workspace"
    workspace.mkdir(parents=True)
    (workspace / "answer.txt").write_text("v1")
    snapshot_root = tmp_path / "snapshots"
    handoff = CodebaseHandoff(root_dir=snapshot_root)

    pointer = handoff.snapshot_local_workspace(source_root, "run-1")
    (workspace / "answer.txt").write_text("v2")
    handoff.snapshot_local_workspace(source_root, "run-1")

    snapshot_file = snapshot_root / pointer.removeprefix("local-snapshot://") / "workspace" / "answer.txt"
    assert snapshot_file.read_text() == "v2"
    assert (workspace / "answer.txt").read_text() == "v2"


def test_backend_snapshot_uses_deepagents_filesystem_interface(tmp_path) -> None:
    pointer = CodebaseHandoff(root_dir=tmp_path / "snapshots").snapshot_backend_workspace(FakeBackend(), "run-1")

    passed, output = LocalCodebaseExecutor(root_dir=tmp_path / "snapshots").run(
        pointer,
        "python app.py && test -f pkg/data.txt",
    )

    assert passed is True
    assert "exit code 0" in output


def test_backend_snapshot_supports_custom_workspace_dir(tmp_path) -> None:
    backend = FakeBackend()
    backend.files = {"/home/daytona/workspace/app.py": b"print('ok')\n"}
    backend.expected_glob = "/home/daytona/workspace/**"

    pointer = CodebaseHandoff(root_dir=tmp_path / "snapshots").snapshot_backend_workspace(
        backend,
        "run-1",
        workspace_dir="/home/daytona/workspace",
    )

    passed, output = LocalCodebaseExecutor(root_dir=tmp_path / "snapshots").run(pointer, "python app.py")

    assert passed is True
    assert "exit code 0" in output


def test_codebase_sandbox_executor_restores_snapshot_into_backend(tmp_path) -> None:
    source_root = tmp_path / "sandbox"
    workspace = source_root / "workspace"
    workspace.mkdir(parents=True)
    (workspace / "answer.txt").write_text("ok")
    pointer = CodebaseHandoff(root_dir=tmp_path / "snapshots").snapshot_local_workspace(source_root, "run-1")
    validator_root = tmp_path / "validator"

    executor = CodebaseSandboxExecutor(
        root_dir=tmp_path / "snapshots",
        backend_factory=lambda: LocalShellBackend(root_dir=validator_root, virtual_mode=True, inherit_env=False),
    )

    passed, output = executor.run(pointer, "test -f answer.txt")

    assert passed is True
    assert "exit code 0" in output
    assert (validator_root / "workspace" / "answer.txt").read_text() == "ok"


def test_codebase_sandbox_executor_opens_backend_lazily(tmp_path) -> None:
    calls: list[bool] = []

    executor = CodebaseSandboxExecutor(
        root_dir=tmp_path / "snapshots",
        backend_factory=lambda: calls.append(True) or ReopenableBackend("backend-1", []),
    )

    assert calls == []
    assert executor.backend is None


def test_codebase_sandbox_executor_restores_redis_archive_pointer(tmp_path) -> None:
    archive_store = RedisCodebaseArchiveStore(client=FakeRedisClient())
    pointer = CodebaseHandoff(archive_store=archive_store).snapshot_backend_workspace(FakeBackend(), "run-1")
    validator_root = tmp_path / "validator"

    executor = CodebaseSandboxExecutor(
        archive_store=archive_store,
        backend_factory=lambda: LocalShellBackend(root_dir=validator_root, virtual_mode=True, inherit_env=False),
    )

    passed, output = executor.run(pointer, "test -f app.py && test -f pkg/data.txt")

    assert pointer == "redis-archive://run-1"
    assert passed is True
    assert "exit code 0" in output


def test_codebase_sandbox_executor_reports_failed_command(tmp_path) -> None:
    source_root = tmp_path / "sandbox"
    (source_root / "workspace").mkdir(parents=True)
    pointer = CodebaseHandoff(root_dir=tmp_path / "snapshots").snapshot_local_workspace(source_root, "run-1")

    executor = CodebaseSandboxExecutor(
        root_dir=tmp_path / "snapshots",
        backend_factory=lambda: LocalShellBackend(
            root_dir=tmp_path / "validator",
            virtual_mode=True,
            inherit_env=False,
        ),
    )

    passed, output = executor.run(pointer, "test -f missing.txt")

    assert passed is False
    assert "exit code" in output.lower()


def test_codebase_sandbox_executor_reopens_backend_after_close(tmp_path) -> None:
    source_root = tmp_path / "sandbox"
    workspace = source_root / "workspace"
    workspace.mkdir(parents=True)
    (workspace / "answer.txt").write_text("ok")
    pointer = CodebaseHandoff(root_dir=tmp_path / "snapshots").snapshot_local_workspace(source_root, "run-1")
    closed: list[str] = []
    created: list[ReopenableBackend] = []

    def backend_factory() -> ReopenableBackend:
        backend = ReopenableBackend(f"backend-{len(created) + 1}", closed)
        created.append(backend)
        return backend

    executor = CodebaseSandboxExecutor(
        root_dir=tmp_path / "snapshots",
        backend_factory=backend_factory,
        sandbox_ttl_seconds=60,
    )

    first = executor.run(pointer, "test -f answer.txt")
    executor.close()
    second = executor.run(pointer, "test -f answer.txt")

    assert first[0] is True
    assert second[0] is True
    assert closed == ["backend-1"]
    assert [backend.name for backend in created] == ["backend-1", "backend-2"]