litellm/tests/unit/harness/sandbox/test_docker.py

import asyncio
import hashlib
import shutil
import subprocess
from typing import Optional

import pytest

from litellm import sandbox
from litellm.harness.errors import SandboxError
from litellm.harness.sandbox import DockerSandbox, Sandbox
from litellm.harness.sandbox.docker import parse_sha256sum

DOCKER_IMAGE = "alpine:3.20"
CID = "cid123"


class FakeStdin:
    def __init__(self) -> None:
        self.data = b""
        self.closed = False

    def write(self, data: bytes) -> None:
        self.data += data

    async def drain(self) -> None:
        return None

    def close(self) -> None:
        self.closed = True


class FakeHandle:
    def __init__(self, stdout: bytes = b"", stderr: bytes = b"", code: int = 0):
        self.stdin = FakeStdin()
        self.stdout = asyncio.StreamReader()
        self.stdout.feed_data(stdout)
        self.stdout.feed_eof()
        self.stderr = asyncio.StreamReader()
        self.stderr.feed_data(stderr)
        self.stderr.feed_eof()
        self.returncode: Optional[int] = code
        self._code = code

    async def wait(self) -> int:
        return self._code

    async def kill(self) -> None:
        return None


class Recorder:
    """Stands in for DockerSandbox._spawn; scripted responses by docker subcommand."""

    def __init__(self) -> None:
        self.calls: list[list[str]] = []
        self.handles: list[FakeHandle] = []
        self.responses: dict[str, FakeHandle] = {}

    async def __call__(self, args: list[str]) -> FakeHandle:
        self.calls.append(args)
        handle = self.responses.pop(args[0], None)
        if handle is None:
            handle = FakeHandle(stdout=f"{CID}\n".encode() if args[0] == "run" else b"")
        self.handles.append(handle)
        return handle


@pytest.fixture
def fake(monkeypatch):
    rec = Recorder()
    monkeypatch.setattr(DockerSandbox, "_spawn", lambda self, args: rec(args))
    return rec


def test_run_args():
    box = sandbox.docker(
        "img:1",
        mounts={"/host/src": "/workspace"},
        env={"A": "1"},
        name="h1",
    )
    assert isinstance(box, Sandbox)
    assert box.run_args() == [
        "run",
        "-d",
        "--rm",
        "--add-host=host.docker.internal:host-gateway",
        "--name",
        "h1",
        "-v",
        "/host/src:/workspace",
        "-e",
        "A=1",
        "-w",
        "/workspace",
        "img:1",
        "sleep",
        "infinity",
    ]
    assert box.host_url(8080) == "http://host.docker.internal:8080"


def test_relative_workdir_rejected():
    with pytest.raises(SandboxError):
        sandbox.docker("img", workdir="rel")


async def test_lazy_start_and_exec_args(fake):
    box = sandbox.docker("img")
    assert fake.calls == []
    await box.exec(["echo", "hi"], env={"K": "V"}, cwd="sub")
    await box.exec(["true"])
    assert fake.calls[0][0] == "run"
    assert [c for c in fake.calls if c[0] == "run"] == [fake.calls[0]]
    assert fake.calls[1] == [
        "exec",
        "-i",
        "-w",
        "/workspace/sub",
        "-e",
        "K=V",
        CID,
        "echo",
        "hi",
    ]
    assert fake.calls[2] == ["exec", "-i", "-w", "/workspace", CID, "true"]


async def test_run_start_failure(fake):
    fake.responses["run"] = FakeHandle(stderr=b"no such image", code=125)
    box = sandbox.docker("img")
    with pytest.raises(SandboxError, match="no such image"):
        await box.run(["echo"])


async def test_read_write_which_tempdir(fake):
    box = sandbox.docker("img")
    await box.start()

    fake.responses["exec"] = FakeHandle(stdout=b"content")
    assert await box.read("a.txt") == b"content"
    assert fake.calls[-1][-2:] == ["cat", "/workspace/a.txt"]

    await box.write("d/b.txt", b"payload")
    assert fake.calls[-1][-5:-1] == ["sh", "-c", fake.calls[-1][-3], "sh"]
    assert fake.calls[-1][-1] == "/workspace/d/b.txt"
    assert fake.handles[-1].stdin.data == b"payload"
    assert fake.handles[-1].stdin.closed

    fake.responses["exec"] = FakeHandle(stdout=b"/usr/bin/codex\n")
    assert await box.which("codex") == "/usr/bin/codex"
    assert fake.calls[-1][-5:] == ["sh", "-lc", 'command -v "$1"', "sh", "codex"]

    fake.responses["exec"] = FakeHandle(code=1)
    assert await box.which("nope") is None

    fake.responses["exec"] = FakeHandle(stdout=b"/tmp/tmp.abc\n")
    assert await box.tempdir() == "/tmp/tmp.abc"

    fake.responses["exec"] = FakeHandle(stderr=b"No such file", code=1)
    with pytest.raises(SandboxError, match="No such file"):
        await box.read("missing")


async def test_snapshot_parses_output(fake):
    box = sandbox.docker("img")
    digest = "a" * 64
    fake.responses["exec"] = FakeHandle(
        stdout=f"{digest}  ./x.txt\n{digest}  ./dir/with space.txt\n".encode()
    )
    snap = await box.snapshot()
    assert snap == {"x.txt": digest, "dir/with space.txt": digest}
    script = fake.calls[-1][-3]
    assert "-name '.git'" in script and "-prune" in script
    assert fake.calls[-1][-1] == "/workspace"


async def test_close_removes_container(fake):
    box = sandbox.docker("img")
    await box.start()
    await box.close()
    assert fake.calls[-1] == ["rm", "-f", CID]
    with pytest.raises(SandboxError):
        await box.start()


async def test_close_without_start_is_noop(fake):
    await sandbox.docker("img").close()
    assert fake.calls == []


async def test_missing_docker_binary(monkeypatch):
    monkeypatch.setattr(shutil, "which", lambda name, *a, **k: None)
    with pytest.raises(SandboxError, match="docker"):
        await sandbox.docker("img").start()


def test_parse_sha256sum_ignores_junk():
    assert parse_sha256sum("garbage\n\n") == {}


def _docker_usable() -> bool:
    if shutil.which("docker") is None:
        return False
    try:
        return (
            subprocess.run(
                ["docker", "info"], capture_output=True, timeout=20
            ).returncode
            == 0
        )
    except (OSError, subprocess.SubprocessError):
        return False


@pytest.mark.skipif(not _docker_usable(), reason="docker daemon not available")
async def test_real_docker_roundtrip():
    box = sandbox.docker(DOCKER_IMAGE, workdir="/workspace")
    try:
        result = await box.run(["echo", "hello"], timeout=120)
        assert result.stdout.strip() == "hello"
        assert result.exit_code == 0

        await box.write("seed.txt", b"seed")
        await box.write("sub/out.txt", b"from host")
        assert await box.read("sub/out.txt") == b"from host"
        await box.write("node_modules/skip.js", b"x")

        assert await box.which("sh") is not None
        assert await box.which("definitely-not-a-binary-xyz") is None
        tmp = await box.tempdir()
        assert tmp.startswith("/")

        snap = await box.snapshot()
        assert snap == {
            "seed.txt": hashlib.sha256(b"seed").hexdigest(),
            "sub/out.txt": hashlib.sha256(b"from host").hexdigest(),
        }
    finally:
        await box.close()