wehub-resource-sync
3642 lines
139 KiB
Python
3642 lines
139 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import io
|
|
import json
|
|
import tarfile
|
|
import time
|
|
import uuid
|
|
from dataclasses import FrozenInstanceError
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
|
|
from agents.sandbox import Manifest, SandboxPathGrant
|
|
from agents.sandbox.config import DEFAULT_PYTHON_SANDBOX_IMAGE
|
|
from agents.sandbox.errors import (
|
|
ExecTimeoutError,
|
|
ExecTransportError,
|
|
ExposedPortUnavailableError,
|
|
InvalidManifestPathError,
|
|
WorkspaceArchiveReadError,
|
|
WorkspaceArchiveWriteError,
|
|
WorkspaceReadNotFoundError,
|
|
WorkspaceWriteTypeError,
|
|
)
|
|
from agents.sandbox.snapshot import NoopSnapshot
|
|
from agents.sandbox.types import ExposedPortEndpoint
|
|
from agents.sandbox.util.tar_utils import validate_tar_bytes
|
|
from tests._fake_workspace_paths import resolve_fake_workspace_path
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Package re-export test
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_blaxel_package_re_exports_backend_symbols() -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import BlaxelSandboxClient
|
|
|
|
package_module = __import__(
|
|
"agents.extensions.sandbox.blaxel", fromlist=["BlaxelSandboxClient"]
|
|
)
|
|
assert package_module.BlaxelSandboxClient is BlaxelSandboxClient
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fakes that replicate the Blaxel SDK surface used by the sandbox backend.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _FakeExecResult:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
exit_code: int = 0,
|
|
output: str = "",
|
|
stderr: str = "",
|
|
pid: str = "",
|
|
) -> None:
|
|
self.exit_code = exit_code
|
|
self.stdout = output
|
|
self.stderr = stderr
|
|
self.logs = output
|
|
self.pid = pid
|
|
|
|
|
|
def _fake_helper_exec_result(command: str, *, symlinks: dict[str, str]) -> _FakeExecResult | None:
|
|
resolved = resolve_fake_workspace_path(
|
|
command,
|
|
symlinks=symlinks,
|
|
home_dir="/workspace",
|
|
)
|
|
if resolved is not None:
|
|
return _FakeExecResult(
|
|
exit_code=resolved.exit_code,
|
|
output=resolved.stdout,
|
|
stderr=resolved.stderr,
|
|
)
|
|
|
|
if "INSTALL_RUNTIME_HELPER_V1" in command or command.startswith(
|
|
"test -x /tmp/openai-agents/bin/resolve-workspace-path-"
|
|
):
|
|
return _FakeExecResult()
|
|
|
|
return None
|
|
|
|
|
|
class _FakeProcess:
|
|
def __init__(self) -> None:
|
|
self.exec_calls: list[tuple[dict[str, Any], dict[str, object]]] = []
|
|
self.next_result = _FakeExecResult()
|
|
self._results_queue: list[_FakeExecResult] = []
|
|
self.delay: float = 0.0
|
|
self.symlinks: dict[str, str] = {}
|
|
|
|
async def exec(self, config: dict[str, Any], **kwargs: object) -> _FakeExecResult:
|
|
self.exec_calls.append((config, dict(kwargs)))
|
|
helper_result = _fake_helper_exec_result(
|
|
str(config.get("command", "")),
|
|
symlinks=self.symlinks,
|
|
)
|
|
if helper_result is not None:
|
|
return helper_result
|
|
if self.delay > 0:
|
|
await asyncio.sleep(self.delay)
|
|
if self._results_queue:
|
|
return self._results_queue.pop(0)
|
|
result = self.next_result
|
|
self.next_result = _FakeExecResult()
|
|
return result
|
|
|
|
|
|
class _FakeFs:
|
|
def __init__(self) -> None:
|
|
self.files: dict[str, bytes] = {}
|
|
self.dirs: list[str] = []
|
|
self.mkdir_calls: list[str] = []
|
|
self.read_error: Exception | None = None
|
|
self.write_error: Exception | None = None
|
|
self.mkdir_error: Exception | None = None
|
|
self.return_str: bool = False
|
|
self.read_binary_calls: list[str] = []
|
|
self.write_binary_calls: list[tuple[str, bytes]] = []
|
|
|
|
async def mkdir(self, path: str, permissions: str = "0755") -> None:
|
|
self.mkdir_calls.append(path)
|
|
if self.mkdir_error is not None:
|
|
raise self.mkdir_error
|
|
self.dirs.append(path)
|
|
|
|
async def read_binary(self, path: str) -> bytes | str:
|
|
self.read_binary_calls.append(path)
|
|
if self.read_error is not None:
|
|
raise self.read_error
|
|
if path not in self.files:
|
|
raise FileNotFoundError(f"not found: {path}")
|
|
data = self.files[path]
|
|
if self.return_str:
|
|
return data.decode("utf-8")
|
|
return data
|
|
|
|
async def write_binary(self, path: str, data: bytes) -> None:
|
|
self.write_binary_calls.append((path, data))
|
|
if self.write_error is not None:
|
|
raise self.write_error
|
|
self.files[path] = data
|
|
|
|
async def ls(self, path: str) -> list[str]:
|
|
# Return files whose paths start with the given directory.
|
|
matches = [p for p in self.files if p.startswith(path.rstrip("/") + "/") or p == path]
|
|
return matches if matches else [path]
|
|
|
|
|
|
class _FakePreviewToken:
|
|
def __init__(self, value: str = "fake-token-abc123") -> None:
|
|
self.value = value
|
|
|
|
|
|
class _FakePreviewTokens:
|
|
def __init__(self) -> None:
|
|
self.create_calls: list[Any] = []
|
|
self.next_token = _FakePreviewToken()
|
|
self.error: Exception | None = None
|
|
|
|
async def create(self, expires_at: Any) -> _FakePreviewToken:
|
|
self.create_calls.append(expires_at)
|
|
if self.error is not None:
|
|
raise self.error
|
|
return self.next_token
|
|
|
|
|
|
class _FakePreview:
|
|
def __init__(self, url: str = "https://preview.example.com:443/") -> None:
|
|
self.url = url
|
|
self.tokens = _FakePreviewTokens()
|
|
|
|
|
|
class _FakePreviews:
|
|
def __init__(self) -> None:
|
|
self.calls: list[dict[str, Any]] = []
|
|
self.next_preview = _FakePreview()
|
|
self.error: Exception | None = None
|
|
|
|
async def create_if_not_exists(self, config: dict[str, Any]) -> _FakePreview:
|
|
self.calls.append(config)
|
|
if self.error is not None:
|
|
raise self.error
|
|
return self.next_preview
|
|
|
|
|
|
class _FakeMetadata:
|
|
def __init__(self, name: str = "test-sandbox", url: str = "https://test.bl.run") -> None:
|
|
self.name = name
|
|
self.url = url
|
|
|
|
|
|
class _FakeSandboxModel:
|
|
def __init__(self, name: str = "test-sandbox", url: str = "https://test.bl.run") -> None:
|
|
self.metadata = _FakeMetadata(name=name, url=url)
|
|
|
|
|
|
class _FakeDrives:
|
|
"""Fake drives API for testing Blaxel Drive mounts."""
|
|
|
|
def __init__(self) -> None:
|
|
self.mount_calls: list[tuple[str, str, str]] = []
|
|
self.unmount_calls: list[str] = []
|
|
self.mount_error: Exception | None = None
|
|
self.unmount_error: Exception | None = None
|
|
|
|
async def mount(self, drive_name: str, mount_path: str, drive_path: str) -> None:
|
|
self.mount_calls.append((drive_name, mount_path, drive_path))
|
|
if self.mount_error is not None:
|
|
raise self.mount_error
|
|
|
|
async def unmount(self, mount_path: str) -> None:
|
|
self.unmount_calls.append(mount_path)
|
|
if self.unmount_error is not None:
|
|
raise self.unmount_error
|
|
|
|
|
|
class _FakeSandboxInstance:
|
|
"""Mimics ``blaxel.core.sandbox.SandboxInstance``."""
|
|
|
|
def __init__(self, name: str = "test-sandbox", url: str = "https://test.bl.run") -> None:
|
|
self.process = _FakeProcess()
|
|
self.fs = _FakeFs()
|
|
self.previews = _FakePreviews()
|
|
self.sandbox = _FakeSandboxModel(name=name, url=url)
|
|
self.drives = _FakeDrives()
|
|
self._deleted = False
|
|
|
|
async def delete(self) -> None:
|
|
self._deleted = True
|
|
|
|
# Class-level stubs used by the client.
|
|
_instances: dict[str, _FakeSandboxInstance] = {}
|
|
_create_error: Exception | None = None
|
|
|
|
@classmethod
|
|
async def create_if_not_exists(cls, config: dict[str, Any]) -> _FakeSandboxInstance:
|
|
if cls._create_error is not None:
|
|
raise cls._create_error
|
|
name = config.get("name", "default")
|
|
inst = cls(name=name)
|
|
cls._instances[name] = inst
|
|
return inst
|
|
|
|
@classmethod
|
|
async def get(cls, name: str) -> _FakeSandboxInstance:
|
|
if name in cls._instances:
|
|
return cls._instances[name]
|
|
raise RuntimeError(f"sandbox {name} not found")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _reset_fake_instances() -> None:
|
|
_FakeSandboxInstance._instances.clear()
|
|
_FakeSandboxInstance._create_error = None
|
|
|
|
|
|
@pytest.fixture()
|
|
def fake_sandbox() -> _FakeSandboxInstance:
|
|
return _FakeSandboxInstance(name="test-sandbox")
|
|
|
|
|
|
def _make_state(
|
|
sandbox_name: str = "test-sandbox",
|
|
root: str = "/workspace",
|
|
pause_on_exit: bool = False,
|
|
sandbox_url: str | None = "https://test.bl.run",
|
|
extra_path_grants: tuple[SandboxPathGrant, ...] = (),
|
|
) -> Any:
|
|
from agents.extensions.sandbox.blaxel.sandbox import (
|
|
BlaxelSandboxSessionState,
|
|
BlaxelTimeouts,
|
|
)
|
|
|
|
return BlaxelSandboxSessionState(
|
|
session_id=uuid.uuid4(),
|
|
manifest=Manifest(root=root, extra_path_grants=extra_path_grants),
|
|
snapshot=NoopSnapshot(id="test-snapshot"),
|
|
sandbox_name=sandbox_name,
|
|
pause_on_exit=pause_on_exit,
|
|
timeouts=BlaxelTimeouts(),
|
|
sandbox_url=sandbox_url,
|
|
)
|
|
|
|
|
|
def _make_session(
|
|
fake: _FakeSandboxInstance,
|
|
state: Any | None = None,
|
|
token: str | None = "test-token",
|
|
) -> Any:
|
|
from agents.extensions.sandbox.blaxel.sandbox import BlaxelSandboxSession
|
|
|
|
if state is None:
|
|
state = _make_state()
|
|
return BlaxelSandboxSession.from_state(state, sandbox=fake, token=token)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Session tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestBlaxelSandboxSession:
|
|
@pytest.mark.asyncio
|
|
async def test_exec_success(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
fake_sandbox.process.next_result = _FakeExecResult(exit_code=0, output="hello world")
|
|
result = await session._exec_internal("echo", "hello")
|
|
assert result.exit_code == 0
|
|
assert result.stdout == b"hello world"
|
|
assert len(fake_sandbox.process.exec_calls) == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exec_success_preserves_split_stderr(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
fake_sandbox.process.next_result = _FakeExecResult(
|
|
exit_code=0,
|
|
output="hello world",
|
|
stderr="warning",
|
|
)
|
|
result = await session._exec_internal("echo", "hello")
|
|
assert result.exit_code == 0
|
|
assert result.stdout == b"hello world"
|
|
assert result.stderr == b"warning"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exec_nonzero(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
fake_sandbox.process.next_result = _FakeExecResult(
|
|
exit_code=1, output="", stderr="error msg"
|
|
)
|
|
result = await session._exec_internal("false")
|
|
assert result.exit_code == 1
|
|
assert result.stderr == b"error msg"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exec_transport_error(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
|
|
async def _raise(*args: object, **kw: object) -> None:
|
|
raise ConnectionError("transport error")
|
|
|
|
fake_sandbox.process.exec = _raise # type: ignore[assignment]
|
|
with pytest.raises(ExecTransportError) as exc_info:
|
|
await session._exec_internal("echo", "hello")
|
|
assert str(exc_info.value) == "Blaxel exec failed: ConnectionError: transport error"
|
|
assert exc_info.value.context["backend"] == "blaxel"
|
|
assert exc_info.value.context["provider_error"] == "ConnectionError: transport error"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mkdir(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
await session.mkdir("subdir")
|
|
assert len(fake_sandbox.fs.mkdir_calls) == 1
|
|
assert "/workspace/subdir" in fake_sandbox.fs.mkdir_calls[0]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
fake_sandbox.fs.files["/workspace/test.txt"] = b"file content"
|
|
result = await session.read("test.txt")
|
|
assert result.read() == b"file content"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read_not_found(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
with pytest.raises(WorkspaceReadNotFoundError):
|
|
await session.read("nonexistent.txt")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read_rejects_workspace_symlink_to_ungranted_path(
|
|
self,
|
|
fake_sandbox: _FakeSandboxInstance,
|
|
) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
fake_sandbox.process.symlinks["/workspace/link"] = "/private"
|
|
|
|
with pytest.raises(InvalidManifestPathError) as exc_info:
|
|
await session.read("link/secret.txt")
|
|
|
|
assert fake_sandbox.fs.read_binary_calls == []
|
|
assert str(exc_info.value) == "manifest path must not escape root: link/secret.txt"
|
|
assert exc_info.value.context == {
|
|
"rel": "link/secret.txt",
|
|
"reason": "escape_root",
|
|
"resolved_path": "workspace escape: /private/secret.txt",
|
|
}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_write(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
await session.write("output.txt", io.BytesIO(b"written data"))
|
|
assert fake_sandbox.fs.files["/workspace/output.txt"] == b"written data"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_write_rejects_workspace_symlink_to_read_only_extra_path_grant(
|
|
self,
|
|
fake_sandbox: _FakeSandboxInstance,
|
|
) -> None:
|
|
state = _make_state(
|
|
extra_path_grants=(SandboxPathGrant(path="/tmp/protected", read_only=True),)
|
|
)
|
|
session = _make_session(fake_sandbox, state=state)
|
|
fake_sandbox.process.symlinks["/workspace/link"] = "/tmp/protected"
|
|
|
|
with pytest.raises(WorkspaceArchiveWriteError) as exc_info:
|
|
await session.write("link/out.txt", io.BytesIO(b"blocked"))
|
|
|
|
assert fake_sandbox.fs.write_binary_calls == []
|
|
assert str(exc_info.value) == "failed to write archive for path: /workspace/link/out.txt"
|
|
assert exc_info.value.context == {
|
|
"path": "/workspace/link/out.txt",
|
|
"reason": "read_only_extra_path_grant",
|
|
"grant_path": "/tmp/protected",
|
|
"resolved_path": "/tmp/protected/out.txt",
|
|
}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mkdir_rejects_workspace_symlink_to_read_only_extra_path_grant(
|
|
self,
|
|
fake_sandbox: _FakeSandboxInstance,
|
|
) -> None:
|
|
state = _make_state(
|
|
extra_path_grants=(SandboxPathGrant(path="/tmp/protected", read_only=True),)
|
|
)
|
|
session = _make_session(fake_sandbox, state=state)
|
|
fake_sandbox.process.symlinks["/workspace/link"] = "/tmp/protected"
|
|
|
|
with pytest.raises(WorkspaceArchiveWriteError) as exc_info:
|
|
await session.mkdir("link/newdir")
|
|
|
|
assert fake_sandbox.fs.mkdir_calls == []
|
|
assert str(exc_info.value) == "failed to write archive for path: /workspace/link/newdir"
|
|
assert exc_info.value.context == {
|
|
"path": "/workspace/link/newdir",
|
|
"reason": "read_only_extra_path_grant",
|
|
"grant_path": "/tmp/protected",
|
|
"resolved_path": "/tmp/protected/newdir",
|
|
}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_running(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
assert await session.running() is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_running_when_down(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
|
|
async def _raise(*args: object, **kw: object) -> None:
|
|
raise ConnectionError("offline")
|
|
|
|
fake_sandbox.fs.ls = _raise # type: ignore[assignment]
|
|
assert await session.running() is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_shutdown_deletes(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
await session.shutdown()
|
|
assert fake_sandbox._deleted is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_shutdown_pause_on_exit(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
state = _make_state(pause_on_exit=True)
|
|
session = _make_session(fake_sandbox, state=state)
|
|
await session.shutdown()
|
|
assert fake_sandbox._deleted is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_normalize_path_relative(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
result = session.normalize_path("subdir/file.txt")
|
|
assert result.as_posix() == "/workspace/subdir/file.txt"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_normalize_path_escape_blocked(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
with pytest.raises(InvalidManifestPathError):
|
|
session.normalize_path("../../etc/passwd")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_normalize_path_absolute_blocked(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
with pytest.raises(InvalidManifestPathError):
|
|
session.normalize_path("/etc/passwd")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mkdir_root_is_noop(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
state = _make_state(root="/")
|
|
session = _make_session(fake_sandbox, state=state)
|
|
await session.mkdir("/")
|
|
# No fs.mkdir call should have been made.
|
|
assert len(fake_sandbox.fs.mkdir_calls) == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mkdir_failure(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
fake_sandbox.fs.mkdir_error = ConnectionError("fs down")
|
|
with pytest.raises(WorkspaceArchiveWriteError):
|
|
await session.mkdir("faildir")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read_returns_str(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
fake_sandbox.fs.files["/workspace/text.txt"] = b"string content"
|
|
fake_sandbox.fs.return_str = True
|
|
result = await session.read("text.txt")
|
|
assert result.read() == b"string content"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read_status_404_via_args_dict(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
# Simulate Blaxel ResponseError with status in args[0] dict.
|
|
err = Exception({"status": 404, "message": "not found"})
|
|
fake_sandbox.fs.read_error = err
|
|
with pytest.raises(WorkspaceReadNotFoundError):
|
|
await session.read("missing.txt")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read_generic_error(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
fake_sandbox.fs.read_error = RuntimeError("unexpected")
|
|
with pytest.raises(WorkspaceArchiveReadError):
|
|
await session.read("broken.txt")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read_status_attr_on_error(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
# Error with .status attribute set (e.g. Blaxel ResponseError).
|
|
session = _make_session(fake_sandbox)
|
|
err = RuntimeError("file missing")
|
|
err.status = 404 # type: ignore[attr-defined]
|
|
fake_sandbox.fs.read_error = err
|
|
with pytest.raises(WorkspaceReadNotFoundError):
|
|
await session.read("gone.txt")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_read_not_found_via_error_string(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
fake_sandbox.fs.read_error = RuntimeError("No such file or directory")
|
|
with pytest.raises(WorkspaceReadNotFoundError):
|
|
await session.read("missing.txt")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_write_str_payload(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
await session.write("text.txt", io.StringIO("hello text"))
|
|
assert fake_sandbox.fs.files["/workspace/text.txt"] == b"hello text"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_write_invalid_payload_type(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
|
|
class _BadIO(io.IOBase):
|
|
def read(self) -> int:
|
|
return 42
|
|
|
|
with pytest.raises(WorkspaceWriteTypeError):
|
|
await session.write("bad.txt", _BadIO())
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_write_fs_error(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
fake_sandbox.fs.write_error = ConnectionError("fs write failed")
|
|
with pytest.raises(WorkspaceArchiveWriteError):
|
|
await session.write("fail.txt", io.BytesIO(b"data"))
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exec_timeout(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
fake_sandbox.process.delay = 10.0
|
|
with pytest.raises(ExecTimeoutError):
|
|
await session._exec_internal("sleep", "100", timeout=0.01)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exec_timeout_reports_default_timeout(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import BlaxelTimeouts
|
|
|
|
state = _make_state()
|
|
state.timeouts = BlaxelTimeouts(exec_timeout_s=1)
|
|
session = _make_session(fake_sandbox, state=state)
|
|
fake_sandbox.process.delay = 10.0
|
|
|
|
with pytest.raises(ExecTimeoutError) as exc_info:
|
|
await session._exec_internal("sleep", "100")
|
|
|
|
assert exc_info.value.timeout_s == 1.0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stop_calls_pty_terminate(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
terminated = []
|
|
original = session.pty_terminate_all
|
|
|
|
async def _track() -> None:
|
|
terminated.append(True)
|
|
await original()
|
|
|
|
session.pty_terminate_all = _track
|
|
await session.stop()
|
|
assert len(terminated) == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_shutdown_delete_raises(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
|
|
async def _raise() -> None:
|
|
raise RuntimeError("delete failed")
|
|
|
|
fake_sandbox.delete = _raise # type: ignore[method-assign]
|
|
# Should not raise; error is suppressed.
|
|
await session.shutdown()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sandbox_name_property(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
assert session.sandbox_name == "test-sandbox"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exposed_port_invalid_url(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
fake_sandbox.previews.next_preview = _FakePreview(url="")
|
|
with pytest.raises(ExposedPortUnavailableError):
|
|
await session._resolve_exposed_port(8080)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exposed_port_bad_url_parse(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
# URL without a hostname.
|
|
fake_sandbox.previews.next_preview = _FakePreview(url="https://")
|
|
with pytest.raises(ExposedPortUnavailableError):
|
|
await session._resolve_exposed_port(8080)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exposed_port_http_scheme(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
fake_sandbox.previews.next_preview = _FakePreview(url="http://preview.example.com/")
|
|
endpoint = await session._resolve_exposed_port(80)
|
|
assert endpoint.tls is False
|
|
assert endpoint.port == 80
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exposed_port(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
endpoint = await session._resolve_exposed_port(3000)
|
|
assert isinstance(endpoint, ExposedPortEndpoint)
|
|
assert endpoint.host == "preview.example.com"
|
|
assert endpoint.tls is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exposed_port_any_port_without_predeclaration(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
"""Blaxel previews can be created for any port on demand."""
|
|
session = _make_session(fake_sandbox)
|
|
# Call the public resolve_exposed_port (which checks _assert_exposed_port_configured).
|
|
# No exposed_ports were declared, but it should still work.
|
|
endpoint = await session.resolve_exposed_port(9999)
|
|
assert isinstance(endpoint, ExposedPortEndpoint)
|
|
assert endpoint.host == "preview.example.com"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exposed_port_error(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
fake_sandbox.previews.error = RuntimeError("backend down")
|
|
with pytest.raises(ExposedPortUnavailableError):
|
|
await session._resolve_exposed_port(3000)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exposed_port_public_preview(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
"""Public preview should not include a token query string."""
|
|
session = _make_session(fake_sandbox)
|
|
endpoint = await session._resolve_exposed_port(8080)
|
|
assert endpoint.query == ""
|
|
# Verify the preview was created with public=True.
|
|
assert fake_sandbox.previews.calls[-1]["spec"]["public"] is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exposed_port_private_preview(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
"""Private preview should create a token and set the query string."""
|
|
state = _make_state()
|
|
object.__setattr__(state, "exposed_port_public", False)
|
|
session = _make_session(fake_sandbox, state=state)
|
|
preview = _FakePreview(url="https://preview.example.com:443/")
|
|
preview.tokens.next_token = _FakePreviewToken(value="my-secret-token")
|
|
fake_sandbox.previews.next_preview = preview
|
|
endpoint = await session._resolve_exposed_port(8080)
|
|
# Verify the preview was created with public=False.
|
|
assert fake_sandbox.previews.calls[-1]["spec"]["public"] is False
|
|
# Verify token was created and attached as query.
|
|
assert len(preview.tokens.create_calls) == 1
|
|
assert endpoint.query == "bl_preview_token=my-secret-token"
|
|
assert "bl_preview_token=my-secret-token" in endpoint.url_for("http")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exposed_port_private_token_error(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
"""Token creation failure should raise ExposedPortUnavailableError."""
|
|
state = _make_state()
|
|
object.__setattr__(state, "exposed_port_public", False)
|
|
session = _make_session(fake_sandbox, state=state)
|
|
preview = _FakePreview(url="https://preview.example.com:443/")
|
|
preview.tokens.error = RuntimeError("token service down")
|
|
fake_sandbox.previews.next_preview = preview
|
|
with pytest.raises(ExposedPortUnavailableError):
|
|
await session._resolve_exposed_port(8080)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_supports_pty_with_url_and_token(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
session = _make_session(fake_sandbox, token="tok")
|
|
# Depends on aiohttp availability in test env.
|
|
try:
|
|
import aiohttp # noqa: F401
|
|
|
|
assert session.supports_pty() is True
|
|
except ImportError:
|
|
assert session.supports_pty() is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_supports_pty_without_token(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox, token=None)
|
|
assert session.supports_pty() is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_supports_pty_without_url(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
state = _make_state(sandbox_url=None)
|
|
session = _make_session(fake_sandbox, state=state, token="tok")
|
|
assert session.supports_pty() is False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Client tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestBlaxelSandboxClient:
|
|
@pytest.mark.asyncio
|
|
async def test_create(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
client = mod.BlaxelSandboxClient(token="test-token")
|
|
options = mod.BlaxelSandboxClientOptions(name="my-sandbox")
|
|
session = await client.create(options=options)
|
|
assert session is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_with_image(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
client = mod.BlaxelSandboxClient(token="test-token")
|
|
options = mod.BlaxelSandboxClientOptions(
|
|
name="img-sandbox",
|
|
image="blaxel/py-app:latest",
|
|
memory=4096,
|
|
region="us-pdx-1",
|
|
)
|
|
session = await client.create(options=options)
|
|
assert session is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
client = mod.BlaxelSandboxClient(token="test-token")
|
|
options = mod.BlaxelSandboxClientOptions(name="del-sandbox")
|
|
session = await client.create(options=options)
|
|
result = await client.delete(session)
|
|
assert result is session
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resume_reconnects(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
# Pre-populate the instance so get() finds it.
|
|
existing = _FakeSandboxInstance(name="resume-sandbox")
|
|
_FakeSandboxInstance._instances["resume-sandbox"] = existing
|
|
|
|
client = mod.BlaxelSandboxClient(token="test-token")
|
|
state = _make_state(sandbox_name="resume-sandbox", pause_on_exit=True)
|
|
session = await client.resume(state)
|
|
assert session is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resume_creates_new(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
client = mod.BlaxelSandboxClient(token="test-token")
|
|
state = _make_state(sandbox_name="new-sandbox", pause_on_exit=False)
|
|
session = await client.resume(state)
|
|
assert session is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deserialize_session_state(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
client = mod.BlaxelSandboxClient(token="test-token")
|
|
payload: dict[str, object] = {
|
|
"session_id": str(uuid.uuid4()),
|
|
"manifest": {"root": "/workspace"},
|
|
"snapshot": {"type": "noop", "id": "test-snap"},
|
|
"sandbox_name": "test",
|
|
}
|
|
state = client.deserialize_session_state(payload)
|
|
assert isinstance(state, mod.BlaxelSandboxSessionState)
|
|
assert state.sandbox_name == "test"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_context_manager(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
async with mod.BlaxelSandboxClient(token="test-token") as client:
|
|
assert client is not None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helper tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestHelpers:
|
|
def test_build_create_config_minimal(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _build_create_config
|
|
|
|
config = _build_create_config(name="test")
|
|
assert config["name"] == "test"
|
|
|
|
def test_build_create_config_full(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _build_create_config
|
|
|
|
config = _build_create_config(
|
|
name="full",
|
|
image=DEFAULT_PYTHON_SANDBOX_IMAGE,
|
|
memory=4096,
|
|
region="us-west",
|
|
env_vars={"KEY": "VAL"},
|
|
labels={"env": "test"},
|
|
ttl="24h",
|
|
)
|
|
assert config["image"] == DEFAULT_PYTHON_SANDBOX_IMAGE
|
|
assert config["memory"] == 4096
|
|
assert config["region"] == "us-west"
|
|
assert config["labels"] == {"env": "test"}
|
|
assert config["ttl"] == "24h"
|
|
assert "ports" not in config
|
|
assert config["envs"] == [{"name": "KEY", "value": "VAL"}]
|
|
|
|
def test_get_sandbox_url(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _get_sandbox_url
|
|
|
|
fake = _FakeSandboxInstance(url="https://sandbox.bl.run")
|
|
assert _get_sandbox_url(fake) == "https://sandbox.bl.run"
|
|
|
|
def test_get_sandbox_url_missing(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _get_sandbox_url
|
|
|
|
class _Bare:
|
|
pass
|
|
|
|
assert _get_sandbox_url(_Bare()) is None
|
|
|
|
def test_build_ws_url(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _build_ws_url
|
|
|
|
url = _build_ws_url(
|
|
sandbox_url="https://test.bl.run",
|
|
token="tok123",
|
|
session_id="sess-1",
|
|
cwd="/workspace",
|
|
)
|
|
assert url.startswith("wss://test.bl.run/terminal/ws?")
|
|
assert "token=tok123" in url
|
|
assert "sessionId=sess-1" in url
|
|
assert "workingDir=/workspace" in url
|
|
|
|
def test_extract_preview_url(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _extract_preview_url
|
|
|
|
assert _extract_preview_url(_FakePreview("https://p.bl.run")) == "https://p.bl.run"
|
|
|
|
def test_extract_preview_url_nested(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _extract_preview_url
|
|
|
|
class _Nested:
|
|
url = None
|
|
|
|
class status:
|
|
url = "https://nested.bl.run"
|
|
|
|
assert _extract_preview_url(_Nested()) == "https://nested.bl.run"
|
|
|
|
def test_extract_preview_url_direct_endpoint(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _extract_preview_url
|
|
|
|
class _Direct:
|
|
url = None
|
|
spec = None
|
|
status = None
|
|
endpoint = "https://direct.bl.run"
|
|
|
|
assert _extract_preview_url(_Direct()) == "https://direct.bl.run"
|
|
|
|
def test_extract_preview_url_inner_preview(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _extract_preview_url
|
|
|
|
class _Inner:
|
|
url = "https://inner.bl.run"
|
|
|
|
class _Outer:
|
|
url = None
|
|
spec = None
|
|
status = None
|
|
endpoint = None
|
|
preview = _Inner()
|
|
|
|
assert _extract_preview_url(_Outer()) == "https://inner.bl.run"
|
|
|
|
def test_extract_preview_url_returns_none(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _extract_preview_url
|
|
|
|
class _Empty:
|
|
pass
|
|
|
|
assert _extract_preview_url(_Empty()) is None
|
|
|
|
def test_get_sandbox_url_direct_url(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _get_sandbox_url
|
|
|
|
class _DirectUrl:
|
|
sandbox = None
|
|
url = "https://direct.bl.run"
|
|
|
|
assert _get_sandbox_url(_DirectUrl()) == "https://direct.bl.run"
|
|
|
|
def test_get_sandbox_url_empty_string(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _get_sandbox_url
|
|
|
|
class _EmptyUrl:
|
|
sandbox = None
|
|
url = ""
|
|
|
|
assert _get_sandbox_url(_EmptyUrl()) is None
|
|
|
|
def test_build_ws_url_http_scheme(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _build_ws_url
|
|
|
|
url = _build_ws_url(
|
|
sandbox_url="http://test.bl.run",
|
|
token="tok",
|
|
session_id="s1",
|
|
cwd="/w",
|
|
)
|
|
assert url.startswith("ws://test.bl.run/")
|
|
|
|
def test_build_create_config_with_ports(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _build_create_config
|
|
|
|
config = _build_create_config(
|
|
name="test",
|
|
ports=({"target": 3000, "protocol": "HTTP"},),
|
|
)
|
|
assert len(config["ports"]) == 1
|
|
assert config["ports"][0]["target"] == 3000
|
|
|
|
def test_build_create_config_region_from_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _build_create_config
|
|
|
|
monkeypatch.setenv("BL_REGION", "eu-ams-1")
|
|
config = _build_create_config(name="test")
|
|
assert config["region"] == "eu-ams-1"
|
|
|
|
def test_build_create_config_default_region(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _build_create_config
|
|
|
|
monkeypatch.delenv("BL_REGION", raising=False)
|
|
config = _build_create_config(name="test")
|
|
assert config["region"] == "us-pdx-1"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Import guard tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestImportGuards:
|
|
def test_import_blaxel_sdk_missing(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
def _fail() -> None:
|
|
raise ImportError("no blaxel")
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", _fail)
|
|
with pytest.raises(ImportError, match="no blaxel"):
|
|
mod._import_blaxel_sdk()
|
|
|
|
def test_import_aiohttp_missing(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _import_aiohttp
|
|
|
|
with patch.dict("sys.modules", {"aiohttp": None}):
|
|
with pytest.raises(ImportError, match="aiohttp"):
|
|
_import_aiohttp()
|
|
|
|
def test_has_aiohttp_false(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _has_aiohttp
|
|
|
|
with patch.dict("sys.modules", {"aiohttp": None}):
|
|
assert _has_aiohttp() is False
|
|
|
|
def test_has_aiohttp_true(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _has_aiohttp
|
|
|
|
# aiohttp should be available in the test environment.
|
|
try:
|
|
import aiohttp # noqa: F401
|
|
|
|
assert _has_aiohttp() is True
|
|
except ImportError:
|
|
pytest.skip("aiohttp not available")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tar validation tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_tar(members: dict[str, bytes | None] | None = None) -> bytes:
|
|
"""Build a tar archive in memory. Pass None as value for directories."""
|
|
buf = io.BytesIO()
|
|
with tarfile.open(fileobj=buf, mode="w") as tar:
|
|
for name, content in (members or {}).items():
|
|
if content is None:
|
|
info = tarfile.TarInfo(name=name)
|
|
info.type = tarfile.DIRTYPE
|
|
tar.addfile(info)
|
|
else:
|
|
info = tarfile.TarInfo(name=name)
|
|
info.size = len(content)
|
|
tar.addfile(info, io.BytesIO(content))
|
|
return buf.getvalue()
|
|
|
|
|
|
def _make_tar_with_symlink_and_file(*, symlink_name: str, target: str, file_name: str) -> bytes:
|
|
buf = io.BytesIO()
|
|
with tarfile.open(fileobj=buf, mode="w") as tar:
|
|
link = tarfile.TarInfo(name=symlink_name)
|
|
link.type = tarfile.SYMTYPE
|
|
link.linkname = target
|
|
tar.addfile(link)
|
|
|
|
contents = b"nested"
|
|
file_info = tarfile.TarInfo(name=file_name)
|
|
file_info.size = len(contents)
|
|
tar.addfile(file_info, io.BytesIO(contents))
|
|
return buf.getvalue()
|
|
|
|
|
|
class TestValidateTarBytes:
|
|
def _validate(self, raw: bytes) -> None:
|
|
validate_tar_bytes(raw)
|
|
|
|
def test_valid_tar(self) -> None:
|
|
raw = _make_tar({"hello.txt": b"content", "subdir/": None})
|
|
self._validate(raw)
|
|
|
|
def test_absolute_path_rejected(self) -> None:
|
|
buf = io.BytesIO()
|
|
with tarfile.open(fileobj=buf, mode="w") as tar:
|
|
info = tarfile.TarInfo(name="/etc/passwd")
|
|
info.size = 4
|
|
tar.addfile(info, io.BytesIO(b"root"))
|
|
with pytest.raises(ValueError, match="absolute path"):
|
|
self._validate(buf.getvalue())
|
|
|
|
def test_parent_traversal_rejected(self) -> None:
|
|
buf = io.BytesIO()
|
|
with tarfile.open(fileobj=buf, mode="w") as tar:
|
|
info = tarfile.TarInfo(name="../escape.txt")
|
|
info.size = 4
|
|
tar.addfile(info, io.BytesIO(b"data"))
|
|
with pytest.raises(ValueError, match="parent traversal"):
|
|
self._validate(buf.getvalue())
|
|
|
|
def test_tar_member_under_archive_symlink_rejected(self) -> None:
|
|
raw = _make_tar_with_symlink_and_file(
|
|
symlink_name="link.txt",
|
|
target="/etc/passwd",
|
|
file_name="link.txt/nested.txt",
|
|
)
|
|
with pytest.raises(ValueError, match="descends through symlink"):
|
|
self._validate(raw)
|
|
|
|
def test_corrupt_tar_rejected(self) -> None:
|
|
with pytest.raises(ValueError, match="invalid tar"):
|
|
self._validate(b"not a tar file at all")
|
|
|
|
def test_dot_entries_skipped(self) -> None:
|
|
raw = _make_tar({"./": None, "file.txt": b"ok"})
|
|
self._validate(raw)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Workspace persistence tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestWorkspacePersistence:
|
|
@pytest.mark.asyncio
|
|
async def test_persist_workspace(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
# Queue up results: mkdir for start, tar command success.
|
|
tar_data = _make_tar({"file.txt": b"hello"})
|
|
fake_sandbox.process._results_queue = [
|
|
_FakeExecResult(exit_code=0, output=""), # tar command
|
|
_FakeExecResult(exit_code=0, output=""), # rm cleanup
|
|
]
|
|
# Pre-populate the tar file so read_binary finds it.
|
|
tar_path = f"/tmp/bl-persist-{session.state.session_id.hex}.tar"
|
|
fake_sandbox.fs.files[tar_path] = tar_data
|
|
result = await session.persist_workspace()
|
|
assert result.read() == tar_data
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persist_workspace_tar_fails(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
fake_sandbox.process._results_queue = [
|
|
_FakeExecResult(exit_code=1, output="tar: error"), # tar command fails
|
|
_FakeExecResult(exit_code=0, output=""), # rm cleanup
|
|
]
|
|
with pytest.raises(WorkspaceArchiveReadError) as exc_info:
|
|
await session.persist_workspace()
|
|
assert exc_info.value.context["reason"] == "tar_failed"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persist_workspace_read_fails(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
fake_sandbox.process._results_queue = [
|
|
_FakeExecResult(exit_code=0, output=""), # tar succeeds
|
|
_FakeExecResult(exit_code=0, output=""), # rm cleanup
|
|
]
|
|
# No tar file in fs, so read_binary will raise FileNotFoundError.
|
|
with pytest.raises(WorkspaceArchiveReadError):
|
|
await session.persist_workspace()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persist_workspace_read_returns_str(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
tar_data = _make_tar({"a.txt": b"data"})
|
|
fake_sandbox.process._results_queue = [
|
|
_FakeExecResult(exit_code=0, output=""),
|
|
_FakeExecResult(exit_code=0, output=""),
|
|
]
|
|
tar_path = f"/tmp/bl-persist-{session.state.session_id.hex}.tar"
|
|
fake_sandbox.fs.files[tar_path] = tar_data
|
|
fake_sandbox.fs.return_str = True
|
|
# This will encode the string back to bytes.
|
|
result = await session.persist_workspace()
|
|
assert len(result.read()) > 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Workspace hydration tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestWorkspaceHydration:
|
|
@pytest.mark.asyncio
|
|
async def test_hydrate_workspace(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
tar_data = _make_tar({"file.txt": b"hello"})
|
|
fake_sandbox.process._results_queue = [
|
|
_FakeExecResult(exit_code=0, output=""), # tar extract
|
|
_FakeExecResult(exit_code=0, output=""), # rm cleanup
|
|
]
|
|
await session.hydrate_workspace(io.BytesIO(tar_data))
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_hydrate_invalid_tar(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
with pytest.raises(WorkspaceArchiveWriteError) as exc_info:
|
|
await session.hydrate_workspace(io.BytesIO(b"not a tar"))
|
|
assert exc_info.value.context["reason"] == "unsafe_or_invalid_tar"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_hydrate_tar_with_symlink(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
raw = _make_tar_with_symlink_and_file(
|
|
symlink_name="link.txt",
|
|
target="/etc/shadow",
|
|
file_name="link.txt/nested.txt",
|
|
)
|
|
with pytest.raises(WorkspaceArchiveWriteError) as exc_info:
|
|
await session.hydrate_workspace(io.BytesIO(raw))
|
|
assert "unsafe_or_invalid_tar" in str(exc_info.value.context)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_hydrate_extract_fails(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
tar_data = _make_tar({"file.txt": b"hello"})
|
|
fake_sandbox.process._results_queue = [
|
|
_FakeExecResult(exit_code=1, output="tar: extract error"), # extract fails
|
|
_FakeExecResult(exit_code=0, output=""), # rm cleanup
|
|
]
|
|
with pytest.raises(WorkspaceArchiveWriteError) as exc_info:
|
|
await session.hydrate_workspace(io.BytesIO(tar_data))
|
|
assert exc_info.value.context["reason"] == "tar_extract_failed"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_hydrate_str_payload_encoded(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
# A str payload gets encoded to bytes, then fails tar validation.
|
|
session = _make_session(fake_sandbox)
|
|
|
|
class _StrIO(io.IOBase):
|
|
def read(self) -> str:
|
|
return "not a valid tar"
|
|
|
|
with pytest.raises(WorkspaceArchiveWriteError) as exc_info:
|
|
await session.hydrate_workspace(_StrIO())
|
|
assert exc_info.value.context["reason"] == "unsafe_or_invalid_tar"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_hydrate_invalid_payload_type(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
|
|
class _IntIO(io.IOBase):
|
|
def read(self) -> int:
|
|
return 42
|
|
|
|
with pytest.raises(WorkspaceWriteTypeError):
|
|
await session.hydrate_workspace(_IntIO())
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_hydrate_write_binary_fails(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
tar_data = _make_tar({"file.txt": b"hello"})
|
|
fake_sandbox.fs.write_error = ConnectionError("upload failed")
|
|
with pytest.raises(WorkspaceArchiveWriteError):
|
|
await session.hydrate_workspace(io.BytesIO(tar_data))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Additional client tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestBlaxelSandboxClientExtra:
|
|
@pytest.mark.asyncio
|
|
async def test_delete_wrong_type(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
client = mod.BlaxelSandboxClient(token="test-token")
|
|
options = mod.BlaxelSandboxClientOptions(name="test")
|
|
session = await client.create(options=options)
|
|
# Replace the inner session with a non-Blaxel type.
|
|
session._inner = "not a BlaxelSandboxSession" # type: ignore[assignment]
|
|
with pytest.raises(TypeError, match="BlaxelSandboxClient.delete"):
|
|
await client.delete(session)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resume_wrong_state_type(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
from tests.utils.factories import TestSessionState
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
client = mod.BlaxelSandboxClient(token="test-token")
|
|
# Pass a non-Blaxel SandboxSessionState subclass.
|
|
state = TestSessionState(
|
|
session_id=uuid.uuid4(),
|
|
manifest=Manifest(root="/workspace"),
|
|
snapshot=NoopSnapshot(id="test"),
|
|
)
|
|
with pytest.raises(TypeError, match="BlaxelSandboxClient.resume"):
|
|
await client.resume(state)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resume_pause_on_exit_get_fails_falls_back(
|
|
self, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
# No instances exist, so get() will fail and fall back to create.
|
|
client = mod.BlaxelSandboxClient(token="test-token")
|
|
state = _make_state(sandbox_name="missing-sandbox", pause_on_exit=True)
|
|
session = await client.resume(state)
|
|
assert session is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_with_timeouts_dict(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
client = mod.BlaxelSandboxClient(token="test-token")
|
|
options = mod.BlaxelSandboxClientOptions(
|
|
name="timeout-test",
|
|
timeouts={"exec_timeout_s": 60, "cleanup_s": 10},
|
|
)
|
|
session = await client.create(options=options)
|
|
assert session is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_without_manifest(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
client = mod.BlaxelSandboxClient(token="test-token")
|
|
options = mod.BlaxelSandboxClientOptions(name="no-manifest")
|
|
session = await client.create(manifest=None, options=options)
|
|
assert session is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_with_all_options(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
client = mod.BlaxelSandboxClient(token="test-token")
|
|
options = mod.BlaxelSandboxClientOptions(
|
|
name="full-opts",
|
|
image=DEFAULT_PYTHON_SANDBOX_IMAGE,
|
|
memory=8192,
|
|
region="eu-ams-1",
|
|
ports=({"target": 3000, "protocol": "HTTP"},),
|
|
env_vars={"FOO": "bar"},
|
|
labels={"team": "test"},
|
|
ttl="1h",
|
|
pause_on_exit=True,
|
|
timeouts=mod.BlaxelTimeouts(exec_timeout_s=120),
|
|
)
|
|
session = await client.create(options=options)
|
|
assert session is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_client_token_from_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
monkeypatch.setenv("BL_API_KEY", "env-token")
|
|
|
|
client = mod.BlaxelSandboxClient()
|
|
assert client._token == "env-token"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_close_is_noop(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
client = mod.BlaxelSandboxClient(token="test-token")
|
|
await client.close() # Should not raise.
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Timeouts model tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestBlaxelTimeouts:
|
|
def test_defaults(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import BlaxelTimeouts
|
|
|
|
t = BlaxelTimeouts()
|
|
assert t.exec_timeout_s == 300.0
|
|
assert t.cleanup_s == 30.0
|
|
assert t.file_upload_s == 1800.0
|
|
assert t.file_download_s == 1800.0
|
|
assert t.workspace_tar_s == 300.0
|
|
assert t.fast_op_s == 30.0
|
|
|
|
def test_custom_values(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import BlaxelTimeouts
|
|
|
|
t = BlaxelTimeouts(exec_timeout_s=60, cleanup_s=10, fast_op_s=5)
|
|
assert t.exec_timeout_s == 60
|
|
assert t.cleanup_s == 10
|
|
assert t.fast_op_s == 5
|
|
|
|
def test_frozen(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import BlaxelTimeouts
|
|
|
|
t = BlaxelTimeouts()
|
|
with pytest.raises(ValidationError):
|
|
t.exec_timeout_s = 999
|
|
|
|
def test_validation_ge_1(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import BlaxelTimeouts
|
|
|
|
with pytest.raises(ValidationError):
|
|
BlaxelTimeouts(exec_timeout_s=0)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Session state tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestBlaxelSandboxSessionState:
|
|
def test_defaults(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import BlaxelSandboxSessionState
|
|
|
|
state = BlaxelSandboxSessionState(
|
|
session_id=uuid.uuid4(),
|
|
manifest=Manifest(root="/workspace"),
|
|
snapshot=NoopSnapshot(id="test"),
|
|
sandbox_name="test",
|
|
)
|
|
assert state.image is None
|
|
assert state.memory is None
|
|
assert state.region is None
|
|
assert state.base_env_vars == {}
|
|
assert state.labels == {}
|
|
assert state.ttl is None
|
|
assert state.pause_on_exit is False
|
|
assert state.sandbox_url is None
|
|
|
|
def test_serialization_roundtrip(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import (
|
|
BlaxelSandboxSessionState,
|
|
BlaxelTimeouts,
|
|
)
|
|
|
|
state = BlaxelSandboxSessionState(
|
|
session_id=uuid.uuid4(),
|
|
manifest=Manifest(root="/workspace"),
|
|
snapshot=NoopSnapshot(id="test"),
|
|
sandbox_name="test-rt",
|
|
image=DEFAULT_PYTHON_SANDBOX_IMAGE,
|
|
memory=4096,
|
|
region="us-pdx-1",
|
|
base_env_vars={"K": "V"},
|
|
labels={"env": "test"},
|
|
ttl="24h",
|
|
pause_on_exit=True,
|
|
timeouts=BlaxelTimeouts(exec_timeout_s=60),
|
|
sandbox_url="https://test.bl.run",
|
|
)
|
|
payload = state.model_dump()
|
|
restored = BlaxelSandboxSessionState.model_validate(payload)
|
|
assert restored.sandbox_name == "test-rt"
|
|
assert restored.image == DEFAULT_PYTHON_SANDBOX_IMAGE
|
|
assert restored.memory == 4096
|
|
assert restored.timeouts.exec_timeout_s == 60
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Client options tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestBlaxelSandboxClientOptions:
|
|
def test_defaults(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import BlaxelSandboxClientOptions
|
|
|
|
opts = BlaxelSandboxClientOptions()
|
|
assert opts.image is None
|
|
assert opts.memory is None
|
|
assert opts.region is None
|
|
assert opts.ports is None
|
|
assert opts.env_vars is None
|
|
assert opts.labels is None
|
|
assert opts.ttl is None
|
|
assert opts.name is None
|
|
assert opts.pause_on_exit is False
|
|
assert opts.timeouts is None
|
|
assert opts.exposed_port_public is True
|
|
|
|
def test_frozen(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import BlaxelSandboxClientOptions
|
|
|
|
opts = BlaxelSandboxClientOptions(name="test")
|
|
with pytest.raises(FrozenInstanceError):
|
|
opts.name = "changed" # type: ignore[misc]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tar exclude args tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestTarExcludeArgs:
|
|
@pytest.mark.asyncio
|
|
async def test_exclude_args_empty(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
args = session._tar_exclude_args()
|
|
# With default manifest (no skip paths), should be empty.
|
|
assert isinstance(args, list)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolved_envs(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
state = _make_state()
|
|
state.base_env_vars = {"BASE_KEY": "base_val"}
|
|
session = _make_session(fake_sandbox, state=state)
|
|
envs = await session._resolved_envs()
|
|
assert envs["BASE_KEY"] == "base_val"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Start lifecycle test
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestStartLifecycle:
|
|
@pytest.mark.asyncio
|
|
async def test_start_mkdir_failure_suppressed(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
|
|
async def _raise(*args: object, **kw: object) -> None:
|
|
raise ConnectionError("mkdir failed")
|
|
|
|
fake_sandbox.process.exec = _raise # type: ignore[assignment]
|
|
# start() should suppress the mkdir error and call super().start().
|
|
# super().start() will try to materialize the manifest, which may
|
|
# also call process.exec. We just verify it does not raise from the
|
|
# initial mkdir.
|
|
try:
|
|
await session.start()
|
|
except Exception:
|
|
# May fail in super().start() but not from the mkdir.
|
|
pass
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# PTY fake helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _FakeWSMessage:
|
|
def __init__(self, msg_type: Any, data: str | bytes) -> None:
|
|
self.type = msg_type
|
|
self.data = data
|
|
|
|
|
|
class _FakeWS:
|
|
"""Fake WebSocket that yields predefined messages then closes."""
|
|
|
|
def __init__(self, messages: list[_FakeWSMessage] | None = None) -> None:
|
|
self._messages = messages or []
|
|
self._sent: list[str] = []
|
|
self._closed = False
|
|
|
|
async def send_str(self, data: str) -> None:
|
|
self._sent.append(data)
|
|
|
|
async def close(self) -> None:
|
|
self._closed = True
|
|
|
|
def __aiter__(self) -> _FakeWS:
|
|
self._iter_index = 0
|
|
return self
|
|
|
|
async def __anext__(self) -> _FakeWSMessage:
|
|
if self._iter_index >= len(self._messages):
|
|
await asyncio.sleep(3600)
|
|
raise StopAsyncIteration
|
|
msg = self._messages[self._iter_index]
|
|
self._iter_index += 1
|
|
return msg
|
|
|
|
|
|
class _FakeHTTPSession:
|
|
def __init__(self, ws: _FakeWS | None = None) -> None:
|
|
self._ws = ws or _FakeWS()
|
|
self._closed = False
|
|
|
|
async def ws_connect(self, url: str) -> _FakeWS:
|
|
return self._ws
|
|
|
|
async def close(self) -> None:
|
|
self._closed = True
|
|
|
|
|
|
class _FakeAiohttp:
|
|
"""Minimal aiohttp mock module."""
|
|
|
|
class WSMsgType:
|
|
TEXT = 1
|
|
BINARY = 2
|
|
ERROR = 256
|
|
CLOSE = 257
|
|
CLOSING = 258
|
|
|
|
def __init__(self, ws: _FakeWS | None = None) -> None:
|
|
self._ws = ws
|
|
|
|
def ClientSession(self) -> _FakeHTTPSession:
|
|
return _FakeHTTPSession(self._ws)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# PTY tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestPtyExec:
|
|
@pytest.mark.asyncio
|
|
async def test_pty_exec_start_success(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
output_msg = json.dumps({"type": "output", "data": "hello from pty"})
|
|
ws = _FakeWS(messages=[_FakeWSMessage(_FakeAiohttp.WSMsgType.TEXT, output_msg)])
|
|
fake_aiohttp = _FakeAiohttp(ws=ws)
|
|
|
|
session = _make_session(fake_sandbox)
|
|
|
|
with patch.object(mod, "_import_aiohttp", return_value=fake_aiohttp):
|
|
update = await session.pty_exec_start("echo", "hello", yield_time_s=0.5)
|
|
assert update.output is not None
|
|
assert b"hello from pty" in update.output
|
|
# process_id may be None if the reader finishes before finalize (entry.done=True).
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_exec_start_timeout(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
session = _make_session(fake_sandbox)
|
|
|
|
class _SlowAiohttp:
|
|
WSMsgType = _FakeAiohttp.WSMsgType
|
|
|
|
def ClientSession(self) -> Any:
|
|
class _SlowSession:
|
|
async def ws_connect(self, url: str) -> None:
|
|
await asyncio.sleep(100)
|
|
|
|
async def close(self) -> None:
|
|
pass
|
|
|
|
return _SlowSession()
|
|
|
|
with patch.object(mod, "_import_aiohttp", return_value=_SlowAiohttp()):
|
|
with pytest.raises(ExecTimeoutError):
|
|
await session.pty_exec_start("echo", "hello", timeout=0.01)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_exec_start_timeout_reports_default_timeout(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
from agents.extensions.sandbox.blaxel.sandbox import BlaxelTimeouts
|
|
|
|
state = _make_state()
|
|
state.timeouts = BlaxelTimeouts(exec_timeout_s=1)
|
|
session = _make_session(fake_sandbox, state=state)
|
|
|
|
class _SlowAiohttp:
|
|
WSMsgType = _FakeAiohttp.WSMsgType
|
|
|
|
def ClientSession(self) -> Any:
|
|
class _SlowSession:
|
|
async def ws_connect(self, url: str) -> None:
|
|
await asyncio.sleep(100)
|
|
|
|
async def close(self) -> None:
|
|
pass
|
|
|
|
return _SlowSession()
|
|
|
|
with patch.object(mod, "_import_aiohttp", return_value=_SlowAiohttp()):
|
|
with pytest.raises(ExecTimeoutError) as exc_info:
|
|
await session.pty_exec_start("echo", "hello")
|
|
|
|
assert exc_info.value.timeout_s == 1.0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_exec_start_connection_error(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
session = _make_session(fake_sandbox)
|
|
|
|
class _ErrorAiohttp:
|
|
WSMsgType = _FakeAiohttp.WSMsgType
|
|
|
|
def ClientSession(self) -> Any:
|
|
class _ErrorSession:
|
|
async def ws_connect(self, url: str) -> None:
|
|
raise ConnectionError("ws connect failed")
|
|
|
|
async def close(self) -> None:
|
|
pass
|
|
|
|
return _ErrorSession()
|
|
|
|
with patch.object(mod, "_import_aiohttp", return_value=_ErrorAiohttp()):
|
|
with pytest.raises(ExecTransportError):
|
|
await session.pty_exec_start("echo", "hello")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_write_stdin(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
|
|
session = _make_session(fake_sandbox)
|
|
ws = _FakeWS()
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id="write-test",
|
|
ws=ws,
|
|
http_session=_FakeHTTPSession(ws),
|
|
)
|
|
session._pty_sessions[1] = entry
|
|
session._reserved_pty_process_ids.add(1)
|
|
|
|
with patch.object(mod, "_import_aiohttp", return_value=_FakeAiohttp()):
|
|
update = await session.pty_write_stdin(session_id=1, chars="input\n", yield_time_s=0.2)
|
|
assert update.output is not None
|
|
assert len(ws._sent) == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_write_stdin_empty_chars(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
|
|
session = _make_session(fake_sandbox)
|
|
ws = _FakeWS()
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id="empty-write",
|
|
ws=ws,
|
|
http_session=_FakeHTTPSession(ws),
|
|
)
|
|
session._pty_sessions[1] = entry
|
|
session._reserved_pty_process_ids.add(1)
|
|
|
|
with patch.object(mod, "_import_aiohttp", return_value=_FakeAiohttp()):
|
|
update = await session.pty_write_stdin(session_id=1, chars="", yield_time_s=0.2)
|
|
assert update.output is not None
|
|
# Empty chars should not send anything.
|
|
assert len(ws._sent) == 0
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_terminate_all(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
|
|
session = _make_session(fake_sandbox)
|
|
ws = _FakeWS()
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id="term-all",
|
|
ws=ws,
|
|
http_session=_FakeHTTPSession(ws),
|
|
)
|
|
session._pty_sessions[1] = entry
|
|
session._reserved_pty_process_ids.add(1)
|
|
|
|
await session.pty_terminate_all()
|
|
assert len(session._pty_sessions) == 0
|
|
assert len(session._reserved_pty_process_ids) == 0
|
|
assert ws._closed
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_ws_reader_error_message(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
error_msg = json.dumps({"type": "error", "data": "something failed"})
|
|
ws = _FakeWS(messages=[_FakeWSMessage(_FakeAiohttp.WSMsgType.TEXT, error_msg)])
|
|
fake_aiohttp = _FakeAiohttp(ws=ws)
|
|
session = _make_session(fake_sandbox)
|
|
|
|
with patch.object(mod, "_import_aiohttp", return_value=fake_aiohttp):
|
|
update = await session.pty_exec_start("bad_cmd", yield_time_s=0.5)
|
|
assert update.output is not None
|
|
assert b"something failed" in update.output
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_ws_reader_binary_message(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
output_msg = json.dumps({"type": "output", "data": "binary-data"}).encode()
|
|
ws = _FakeWS(messages=[_FakeWSMessage(_FakeAiohttp.WSMsgType.BINARY, output_msg)])
|
|
fake_aiohttp = _FakeAiohttp(ws=ws)
|
|
session = _make_session(fake_sandbox)
|
|
|
|
with patch.object(mod, "_import_aiohttp", return_value=fake_aiohttp):
|
|
update = await session.pty_exec_start("echo", "test", yield_time_s=0.5)
|
|
assert b"binary-data" in update.output
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_ws_reader_close_message(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
ws = _FakeWS(
|
|
messages=[
|
|
_FakeWSMessage(
|
|
_FakeAiohttp.WSMsgType.TEXT, json.dumps({"type": "output", "data": "hi"})
|
|
),
|
|
_FakeWSMessage(_FakeAiohttp.WSMsgType.CLOSE, ""),
|
|
]
|
|
)
|
|
fake_aiohttp = _FakeAiohttp(ws=ws)
|
|
session = _make_session(fake_sandbox)
|
|
|
|
with patch.object(mod, "_import_aiohttp", return_value=fake_aiohttp):
|
|
update = await session.pty_exec_start("echo", "test", yield_time_s=0.5)
|
|
assert b"hi" in update.output
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_ws_reader_invalid_json(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
ws = _FakeWS(
|
|
messages=[
|
|
_FakeWSMessage(_FakeAiohttp.WSMsgType.TEXT, "not json"),
|
|
_FakeWSMessage(
|
|
_FakeAiohttp.WSMsgType.TEXT,
|
|
json.dumps({"type": "output", "data": "valid"}),
|
|
),
|
|
]
|
|
)
|
|
fake_aiohttp = _FakeAiohttp(ws=ws)
|
|
session = _make_session(fake_sandbox)
|
|
|
|
with patch.object(mod, "_import_aiohttp", return_value=fake_aiohttp):
|
|
update = await session.pty_exec_start("echo", "test", yield_time_s=0.5)
|
|
# Invalid JSON should be silently ignored; valid output should appear.
|
|
assert b"valid" in update.output
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_ws_reader_error_type_message(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
ws = _FakeWS(
|
|
messages=[
|
|
_FakeWSMessage(_FakeAiohttp.WSMsgType.ERROR, "ws error"),
|
|
]
|
|
)
|
|
fake_aiohttp = _FakeAiohttp(ws=ws)
|
|
session = _make_session(fake_sandbox)
|
|
|
|
with patch.object(mod, "_import_aiohttp", return_value=fake_aiohttp):
|
|
update = await session.pty_exec_start("echo", "test", yield_time_s=0.3)
|
|
# Error WS message should break the reader loop.
|
|
assert update.output is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_finalize_done_session(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
|
|
session = _make_session(fake_sandbox)
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id="test-done",
|
|
ws=None,
|
|
http_session=None,
|
|
done=True,
|
|
exit_code=0,
|
|
)
|
|
# Manually register the entry.
|
|
session._pty_sessions[1] = entry
|
|
session._reserved_pty_process_ids.add(1)
|
|
|
|
result = await session._finalize_pty_update(
|
|
process_id=1,
|
|
entry=entry,
|
|
output=b"done output",
|
|
original_token_count=None,
|
|
)
|
|
assert result.process_id is None
|
|
assert result.exit_code == 0
|
|
assert 1 not in session._pty_sessions
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_prune_sessions(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
from agents.sandbox.session.pty_types import PTY_PROCESSES_MAX
|
|
|
|
session = _make_session(fake_sandbox)
|
|
# Fill to max capacity with done entries.
|
|
for i in range(PTY_PROCESSES_MAX):
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id=f"test-{i}",
|
|
ws=None,
|
|
http_session=None,
|
|
done=True,
|
|
exit_code=0,
|
|
)
|
|
entry.last_used = time.monotonic() - (PTY_PROCESSES_MAX - i)
|
|
session._pty_sessions[i] = entry
|
|
session._reserved_pty_process_ids.add(i)
|
|
|
|
pruned = session._prune_pty_sessions_if_needed()
|
|
assert pruned is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_prune_below_max(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
# Below max, no pruning.
|
|
pruned = session._prune_pty_sessions_if_needed()
|
|
assert pruned is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_terminate_pty_entry_with_reader_task(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
|
|
session = _make_session(fake_sandbox)
|
|
ws = _FakeWS()
|
|
http = _FakeHTTPSession(ws)
|
|
|
|
async def _reader() -> None:
|
|
await asyncio.sleep(100)
|
|
|
|
task = asyncio.create_task(_reader())
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id="term-test",
|
|
ws=ws,
|
|
http_session=http,
|
|
reader_task=task,
|
|
)
|
|
await session._terminate_pty_entry(entry)
|
|
assert task.cancelled() or task.done()
|
|
assert ws._closed
|
|
assert http._closed
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_terminate_pty_entry_all_none(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
|
|
session = _make_session(fake_sandbox)
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id="null-test",
|
|
ws=None,
|
|
http_session=None,
|
|
reader_task=None,
|
|
)
|
|
# Should not raise.
|
|
await session._terminate_pty_entry(entry)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_exec_default_yield_time(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
ws = _FakeWS(
|
|
messages=[
|
|
_FakeWSMessage(
|
|
_FakeAiohttp.WSMsgType.TEXT,
|
|
json.dumps({"type": "output", "data": "quick"}),
|
|
),
|
|
]
|
|
)
|
|
fake_aiohttp = _FakeAiohttp(ws=ws)
|
|
session = _make_session(fake_sandbox)
|
|
|
|
with patch.object(mod, "_import_aiohttp", return_value=fake_aiohttp):
|
|
# Pass yield_time_s=None to test default (10s), but with a short timeout.
|
|
# We use a small timeout to not wait 10 seconds.
|
|
update = await session.pty_exec_start("echo", "test", yield_time_s=0.1)
|
|
assert b"quick" in update.output
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_ws_reader_capital_type_keys(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
# Test the alternative capitalized key paths (Type/Data).
|
|
output_msg = json.dumps({"Type": "output", "Data": "cap-data"})
|
|
ws = _FakeWS(messages=[_FakeWSMessage(_FakeAiohttp.WSMsgType.TEXT, output_msg)])
|
|
fake_aiohttp = _FakeAiohttp(ws=ws)
|
|
session = _make_session(fake_sandbox)
|
|
|
|
with patch.object(mod, "_import_aiohttp", return_value=fake_aiohttp):
|
|
update = await session.pty_exec_start("echo", "test", yield_time_s=0.5)
|
|
assert b"cap-data" in update.output
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_max_output_tokens(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
long_output = "x" * 10000
|
|
output_msg = json.dumps({"type": "output", "data": long_output})
|
|
ws = _FakeWS(messages=[_FakeWSMessage(_FakeAiohttp.WSMsgType.TEXT, output_msg)])
|
|
fake_aiohttp = _FakeAiohttp(ws=ws)
|
|
session = _make_session(fake_sandbox)
|
|
|
|
with patch.object(mod, "_import_aiohttp", return_value=fake_aiohttp):
|
|
update = await session.pty_exec_start(
|
|
"echo", "test", yield_time_s=0.5, max_output_tokens=10
|
|
)
|
|
# Output should be truncated.
|
|
assert len(update.output) < len(long_output.encode())
|
|
assert update.original_token_count is not None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Persist workspace with mount handling
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestPersistWithMounts:
|
|
@pytest.mark.asyncio
|
|
async def test_persist_unmount_error(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
|
|
mock_strategy = MagicMock()
|
|
mock_strategy.teardown_for_snapshot = AsyncMock(side_effect=RuntimeError("unmount fail"))
|
|
|
|
mock_mount = MagicMock()
|
|
mock_mount.mount_strategy = mock_strategy
|
|
mount_path = Path("/workspace/mount")
|
|
|
|
orig_manifest = session.state.manifest
|
|
mock_manifest = MagicMock(wraps=orig_manifest)
|
|
mock_manifest.root = orig_manifest.root
|
|
mock_manifest.environment = orig_manifest.environment
|
|
mock_manifest.ephemeral_mount_targets = MagicMock(return_value=[(mock_mount, mount_path)])
|
|
session.state.manifest = mock_manifest
|
|
|
|
with pytest.raises(WorkspaceArchiveReadError):
|
|
await session.persist_workspace()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persist_remount_error(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
tar_data = _make_tar({"file.txt": b"data"})
|
|
|
|
mock_strategy = MagicMock()
|
|
mock_strategy.teardown_for_snapshot = AsyncMock()
|
|
mock_strategy.restore_after_snapshot = AsyncMock(side_effect=RuntimeError("remount fail"))
|
|
|
|
mock_mount = MagicMock()
|
|
mock_mount.mount_strategy = mock_strategy
|
|
mount_path = Path("/workspace/mount")
|
|
|
|
fake_sandbox.process._results_queue = [
|
|
_FakeExecResult(exit_code=0, output=""),
|
|
_FakeExecResult(exit_code=0, output=""),
|
|
]
|
|
tar_path = f"/tmp/bl-persist-{session.state.session_id.hex}.tar"
|
|
fake_sandbox.fs.files[tar_path] = tar_data
|
|
|
|
orig_manifest = session.state.manifest
|
|
mock_manifest = MagicMock(wraps=orig_manifest)
|
|
mock_manifest.root = orig_manifest.root
|
|
mock_manifest.environment = orig_manifest.environment
|
|
mock_manifest.ephemeral_mount_targets = MagicMock(return_value=[(mock_mount, mount_path)])
|
|
session.state.manifest = mock_manifest
|
|
|
|
with pytest.raises(WorkspaceArchiveReadError):
|
|
await session.persist_workspace()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persist_snapshot_error_still_remounts(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
|
|
mock_strategy = MagicMock()
|
|
mock_strategy.teardown_for_snapshot = AsyncMock()
|
|
mock_strategy.restore_after_snapshot = AsyncMock()
|
|
|
|
mock_mount = MagicMock()
|
|
mock_mount.mount_strategy = mock_strategy
|
|
mount_path = Path("/workspace/mount")
|
|
|
|
fake_sandbox.process._results_queue = [
|
|
_FakeExecResult(exit_code=1, output="tar fail"),
|
|
_FakeExecResult(exit_code=0, output=""),
|
|
]
|
|
|
|
orig_manifest = session.state.manifest
|
|
mock_manifest = MagicMock(wraps=orig_manifest)
|
|
mock_manifest.root = orig_manifest.root
|
|
mock_manifest.environment = orig_manifest.environment
|
|
mock_manifest.ephemeral_mount_targets = MagicMock(return_value=[(mock_mount, mount_path)])
|
|
session.state.manifest = mock_manifest
|
|
|
|
with pytest.raises(WorkspaceArchiveReadError):
|
|
await session.persist_workspace()
|
|
|
|
mock_strategy.restore_after_snapshot.assert_called_once()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _import_blaxel_sdk actual error path
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestImportBlaxelSdkActual:
|
|
def test_actual_import_error(self) -> None:
|
|
# Force the actual function (not mocked) to fail by hiding the module.
|
|
from agents.extensions.sandbox.blaxel.sandbox import _import_blaxel_sdk
|
|
|
|
with patch.dict(
|
|
"sys.modules", {"blaxel": None, "blaxel.core": None, "blaxel.core.sandbox": None}
|
|
):
|
|
with pytest.raises(ImportError, match="BlaxelSandboxClient requires"):
|
|
_import_blaxel_sdk()
|
|
|
|
def test_actual_import_aiohttp_error(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _import_aiohttp
|
|
|
|
with patch.dict("sys.modules", {"aiohttp": None}):
|
|
with pytest.raises(ImportError, match="aiohttp"):
|
|
_import_aiohttp()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# shared tar validation: unsupported member type (for example, device or fifo)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestValidateTarBytesExtra:
|
|
def test_unsupported_member_type(self) -> None:
|
|
buf = io.BytesIO()
|
|
with tarfile.open(fileobj=buf, mode="w") as tar:
|
|
info = tarfile.TarInfo(name="device")
|
|
info.type = tarfile.CHRTYPE # Character device, not dir or reg.
|
|
tar.addfile(info)
|
|
|
|
with pytest.raises(ValueError, match="unsupported member type"):
|
|
validate_tar_bytes(buf.getvalue())
|
|
|
|
def test_hardlink_rejected(self) -> None:
|
|
buf = io.BytesIO()
|
|
with tarfile.open(fileobj=buf, mode="w") as tar:
|
|
info = tarfile.TarInfo(name="hardlink")
|
|
info.type = tarfile.LNKTYPE
|
|
info.linkname = "target"
|
|
tar.addfile(info)
|
|
|
|
with pytest.raises(ValueError, match="hardlink"):
|
|
validate_tar_bytes(buf.getvalue())
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Additional coverage: tar_exclude_args with skip paths
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestTarExcludeArgsWithSkipPaths:
|
|
@pytest.mark.asyncio
|
|
async def test_exclude_args_with_skip_paths(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
session._runtime_persist_workspace_skip_relpaths = {
|
|
Path("node_modules"),
|
|
Path(".git"),
|
|
}
|
|
args = session._tar_exclude_args()
|
|
assert len(args) > 0
|
|
assert any("node_modules" in a for a in args)
|
|
assert any(".git" in a for a in args)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exclude_args_skips_empty_and_dot(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
session._runtime_persist_workspace_skip_relpaths = {
|
|
Path("."),
|
|
Path("keep_me"),
|
|
}
|
|
args = session._tar_exclude_args()
|
|
# "." should be skipped, "keep_me" should be included.
|
|
assert any("keep_me" in a for a in args)
|
|
assert not any(a == "--exclude='.'" for a in args)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Additional coverage: terminate entry with close errors
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestTerminatePtyEntryErrors:
|
|
@pytest.mark.asyncio
|
|
async def test_terminate_ws_close_error(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
|
|
session = _make_session(fake_sandbox)
|
|
|
|
class _ErrorWS:
|
|
async def close(self) -> None:
|
|
raise ConnectionError("ws close failed")
|
|
|
|
class _ErrorHTTP:
|
|
async def close(self) -> None:
|
|
raise ConnectionError("http close failed")
|
|
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id="err-close",
|
|
ws=_ErrorWS(),
|
|
http_session=_ErrorHTTP(),
|
|
reader_task=None,
|
|
)
|
|
# Should not raise.
|
|
await session._terminate_pty_entry(entry)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_terminate_reader_already_done(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
|
|
session = _make_session(fake_sandbox)
|
|
|
|
async def _done_task() -> None:
|
|
pass
|
|
|
|
task = asyncio.create_task(_done_task())
|
|
await task # Let it complete.
|
|
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id="done-reader",
|
|
ws=_FakeWS(),
|
|
http_session=_FakeHTTPSession(),
|
|
reader_task=task,
|
|
)
|
|
await session._terminate_pty_entry(entry)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Additional coverage: _collect_pty_output with entry already done at start
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCollectPtyOutputEdgeCases:
|
|
@pytest.mark.asyncio
|
|
async def test_collect_output_entry_done_immediately(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
|
|
session = _make_session(fake_sandbox)
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id="done-imm",
|
|
ws=None,
|
|
http_session=None,
|
|
done=True,
|
|
)
|
|
entry.output_chunks.append(b"final output")
|
|
output, token_count = await session._collect_pty_output(
|
|
entry=entry, yield_time_ms=100, max_output_tokens=None
|
|
)
|
|
assert b"final output" in output
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_collect_output_timeout_path(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
|
|
session = _make_session(fake_sandbox)
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id="timeout-collect",
|
|
ws=None,
|
|
http_session=None,
|
|
)
|
|
# Very short yield time, no output, not done.
|
|
output, token_count = await session._collect_pty_output(
|
|
entry=entry, yield_time_ms=1, max_output_tokens=None
|
|
)
|
|
assert output == b""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Additional coverage: actual import success paths
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestActualImportSuccess:
|
|
def test_import_blaxel_sdk_success(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _import_blaxel_sdk
|
|
|
|
try:
|
|
result = _import_blaxel_sdk()
|
|
assert result is not None
|
|
except ImportError:
|
|
pytest.skip("blaxel not available")
|
|
|
|
def test_import_aiohttp_success(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _import_aiohttp
|
|
|
|
try:
|
|
result = _import_aiohttp()
|
|
assert result is not None
|
|
except ImportError:
|
|
pytest.skip("aiohttp not available")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Additional coverage: hydrate cleanup and persist cleanup rm paths
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCleanupPaths:
|
|
@pytest.mark.asyncio
|
|
async def test_persist_cleanup_rm_failure_suppressed(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
tar_data = _make_tar({"file.txt": b"hello"})
|
|
|
|
call_count = 0
|
|
|
|
async def _counting_exec(config: dict[str, Any], **kw: object) -> _FakeExecResult:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count == 1:
|
|
# tar command succeeds.
|
|
return _FakeExecResult(exit_code=0, output="")
|
|
# rm cleanup fails.
|
|
raise ConnectionError("rm failed")
|
|
|
|
fake_sandbox.process.exec = _counting_exec # type: ignore[method-assign]
|
|
tar_path = f"/tmp/bl-persist-{session.state.session_id.hex}.tar"
|
|
fake_sandbox.fs.files[tar_path] = tar_data
|
|
|
|
# Should succeed despite rm failure.
|
|
result = await session.persist_workspace()
|
|
assert result.read() == tar_data
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_hydrate_cleanup_rm_failure_suppressed(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
tar_data = _make_tar({"file.txt": b"hello"})
|
|
|
|
call_count = 0
|
|
|
|
async def _counting_exec(config: dict[str, Any], **kw: object) -> _FakeExecResult:
|
|
nonlocal call_count
|
|
call_count += 1
|
|
command = str(config.get("command", ""))
|
|
helper_result = _fake_helper_exec_result(
|
|
command, symlinks=fake_sandbox.process.symlinks
|
|
)
|
|
if helper_result is not None:
|
|
return helper_result
|
|
if "tar" in command:
|
|
if "xf" in command:
|
|
# tar extract succeeds.
|
|
return _FakeExecResult(exit_code=0, output="")
|
|
if "rm" in command:
|
|
raise ConnectionError("rm failed")
|
|
return _FakeExecResult(exit_code=0, output="")
|
|
|
|
fake_sandbox.process.exec = _counting_exec # type: ignore[method-assign]
|
|
|
|
# Should succeed despite rm failure.
|
|
await session.hydrate_workspace(io.BytesIO(tar_data))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Additional coverage: client branch partials
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestClientBranchCoverage:
|
|
@pytest.mark.asyncio
|
|
async def test_create_no_name_generates_one(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
client = mod.BlaxelSandboxClient(token="test-token")
|
|
options = mod.BlaxelSandboxClientOptions() # No name.
|
|
session = await client.create(options=options)
|
|
assert session is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resume_reconnects_no_new_url(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
# Create an instance with no URL.
|
|
class _NoUrlSandbox(_FakeSandboxInstance):
|
|
def __init__(self, name: str = "no-url") -> None:
|
|
super().__init__(name=name)
|
|
self.sandbox = _FakeSandboxModel(name=name, url="")
|
|
|
|
_FakeSandboxInstance._instances["no-url-sandbox"] = _NoUrlSandbox("no-url-sandbox")
|
|
|
|
client = mod.BlaxelSandboxClient(token="test-token")
|
|
state = _make_state(sandbox_name="no-url-sandbox", pause_on_exit=True)
|
|
session = await client.resume(state)
|
|
assert session is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_shutdown_raises(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
monkeypatch.setattr(mod, "_import_blaxel_sdk", lambda: _FakeSandboxInstance)
|
|
|
|
client = mod.BlaxelSandboxClient(token="test-token")
|
|
options = mod.BlaxelSandboxClientOptions(name="del-err")
|
|
session = await client.create(options=options)
|
|
|
|
# Make shutdown raise.
|
|
async def _raise() -> None:
|
|
raise RuntimeError("shutdown error")
|
|
|
|
session._inner.shutdown = _raise # type: ignore[method-assign]
|
|
# delete should suppress the error.
|
|
result = await client.delete(session)
|
|
assert result is session
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Final coverage gap tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFinalCoverageGaps:
|
|
@pytest.mark.asyncio
|
|
async def test_exec_reraises_exec_timeout_error(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
"""Cover line 401: except (ExecTimeoutError, ExecTransportError): raise."""
|
|
session = _make_session(fake_sandbox)
|
|
|
|
async def _timeout_exec(*args: object, **kw: object) -> None:
|
|
raise ExecTimeoutError(command=("test",), timeout_s=1.0, cause=None)
|
|
|
|
fake_sandbox.process.exec = _timeout_exec # type: ignore[assignment]
|
|
with pytest.raises(ExecTimeoutError):
|
|
await session._exec_internal("test")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persist_rm_exception_suppressed(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
"""Cover lines 493-494: except Exception: pass in persist cleanup."""
|
|
session = _make_session(fake_sandbox)
|
|
tar_data = _make_tar({"file.txt": b"hello"})
|
|
|
|
async def _exec_with_rm_fail(config: dict[str, Any], **kw: object) -> _FakeExecResult:
|
|
command = str(config.get("command", ""))
|
|
helper_result = _fake_helper_exec_result(
|
|
command, symlinks=fake_sandbox.process.symlinks
|
|
)
|
|
if helper_result is not None:
|
|
return helper_result
|
|
if "rm" in command:
|
|
raise OSError("rm failed")
|
|
return _FakeExecResult(exit_code=0, output="")
|
|
|
|
fake_sandbox.process.exec = _exec_with_rm_fail # type: ignore[method-assign]
|
|
tar_path = f"/tmp/bl-persist-{session.state.session_id.hex}.tar"
|
|
fake_sandbox.fs.files[tar_path] = tar_data
|
|
|
|
result = await session.persist_workspace()
|
|
assert result.read() == tar_data
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_hydrate_rm_exception_suppressed(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
"""Cover lines 560-561: except Exception: pass in hydrate cleanup."""
|
|
session = _make_session(fake_sandbox)
|
|
tar_data = _make_tar({"file.txt": b"hello"})
|
|
|
|
async def _exec_with_rm_fail(config: dict[str, Any], **kw: object) -> _FakeExecResult:
|
|
command = str(config.get("command", ""))
|
|
helper_result = _fake_helper_exec_result(
|
|
command, symlinks=fake_sandbox.process.symlinks
|
|
)
|
|
if helper_result is not None:
|
|
return helper_result
|
|
if "rm" in command:
|
|
raise OSError("rm failed")
|
|
return _FakeExecResult(exit_code=0, output="")
|
|
|
|
fake_sandbox.process.exec = _exec_with_rm_fail # type: ignore[method-assign]
|
|
|
|
await session.hydrate_workspace(io.BytesIO(tar_data))
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_exec_with_pruning(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
"""Cover line 638: pruned entry termination in pty_exec_start."""
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
from agents.sandbox.session.pty_types import PTY_PROCESSES_MAX
|
|
|
|
session = _make_session(fake_sandbox)
|
|
|
|
# Fill sessions to capacity with done entries.
|
|
for i in range(PTY_PROCESSES_MAX):
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id=f"fill-{i}",
|
|
ws=None,
|
|
http_session=None,
|
|
done=True,
|
|
exit_code=0,
|
|
)
|
|
entry.last_used = time.monotonic() - (PTY_PROCESSES_MAX - i)
|
|
session._pty_sessions[i + 100] = entry
|
|
session._reserved_pty_process_ids.add(i + 100)
|
|
|
|
ws = _FakeWS(
|
|
messages=[
|
|
_FakeWSMessage(
|
|
_FakeAiohttp.WSMsgType.TEXT,
|
|
json.dumps({"type": "output", "data": "pruned-test"}),
|
|
),
|
|
]
|
|
)
|
|
fake_aiohttp = _FakeAiohttp(ws=ws)
|
|
|
|
with patch.object(mod, "_import_aiohttp", return_value=fake_aiohttp):
|
|
update = await session.pty_exec_start("echo", "test", yield_time_s=0.3)
|
|
assert b"pruned-test" in update.output
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_warning_threshold(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
"""Cover line 641: warning log for high PTY count."""
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
from agents.sandbox.session.pty_types import PTY_PROCESSES_WARNING
|
|
|
|
session = _make_session(fake_sandbox)
|
|
|
|
# Fill up to just below warning threshold.
|
|
for i in range(PTY_PROCESSES_WARNING - 1):
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id=f"warn-{i}",
|
|
ws=None,
|
|
http_session=None,
|
|
)
|
|
session._pty_sessions[i + 200] = entry
|
|
session._reserved_pty_process_ids.add(i + 200)
|
|
|
|
ws = _FakeWS(
|
|
messages=[
|
|
_FakeWSMessage(
|
|
_FakeAiohttp.WSMsgType.TEXT,
|
|
json.dumps({"type": "output", "data": "warn-test"}),
|
|
),
|
|
]
|
|
)
|
|
fake_aiohttp = _FakeAiohttp(ws=ws)
|
|
|
|
with patch.object(mod, "_import_aiohttp", return_value=fake_aiohttp):
|
|
update = await session.pty_exec_start("echo", "test", yield_time_s=0.3)
|
|
assert update.output is not None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pty_ws_reader_exception_in_iter(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
"""Cover line 744: except Exception: pass in _pty_ws_reader."""
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
|
|
session = _make_session(fake_sandbox)
|
|
|
|
class _ErrorWS:
|
|
_sent: list[str] = []
|
|
_closed = False
|
|
|
|
async def send_str(self, data: str) -> None:
|
|
self._sent.append(data)
|
|
|
|
async def close(self) -> None:
|
|
self._closed = True
|
|
|
|
def __aiter__(self) -> _ErrorWS:
|
|
return self
|
|
|
|
async def __anext__(self) -> None:
|
|
raise RuntimeError("WS iteration error")
|
|
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id="err-iter",
|
|
ws=_ErrorWS(),
|
|
http_session=_FakeHTTPSession(),
|
|
)
|
|
|
|
# Run the reader directly.
|
|
await session._pty_ws_reader(entry)
|
|
assert entry.done is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_terminate_pty_outer_exception(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
"""Cover lines 841-842: outer except Exception: pass in _terminate_pty_entry."""
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
|
|
session = _make_session(fake_sandbox)
|
|
|
|
class _BadReaderTask:
|
|
"""Fake task whose done() raises."""
|
|
|
|
def done(self) -> bool:
|
|
raise RuntimeError("task check failed")
|
|
|
|
def cancel(self) -> None:
|
|
pass
|
|
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id="outer-err",
|
|
ws=None,
|
|
http_session=None,
|
|
reader_task=_BadReaderTask(), # type: ignore[arg-type]
|
|
)
|
|
# Should not raise.
|
|
await session._terminate_pty_entry(entry)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prune_returns_none_when_no_pid(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
"""Cover line 819: prune returns None when process_id_to_prune_from_meta returns None."""
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
from agents.sandbox.session.pty_types import PTY_PROCESSES_MAX
|
|
|
|
session = _make_session(fake_sandbox)
|
|
|
|
# Fill to max with entries, then patch process_id_to_prune_from_meta to return None.
|
|
for i in range(PTY_PROCESSES_MAX):
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id=f"no-prune-{i}",
|
|
ws=None,
|
|
http_session=None,
|
|
)
|
|
session._pty_sessions[i + 300] = entry
|
|
session._reserved_pty_process_ids.add(i + 300)
|
|
|
|
with patch(
|
|
"agents.extensions.sandbox.blaxel.sandbox.process_id_to_prune_from_meta",
|
|
return_value=None,
|
|
):
|
|
result = session._prune_pty_sessions_if_needed()
|
|
assert result is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_collect_output_deadline_break(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
"""Cover lines 765, 774: deadline and remaining_s break paths."""
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
|
|
session = _make_session(fake_sandbox)
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id="deadline-test",
|
|
ws=None,
|
|
http_session=None,
|
|
)
|
|
entry.output_chunks.append(b"some data")
|
|
|
|
# yield_time_ms=1 means very short deadline, should hit deadline break.
|
|
output, _ = await session._collect_pty_output(
|
|
entry=entry, yield_time_ms=1, max_output_tokens=None
|
|
)
|
|
assert b"some data" in output
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_collect_output_done_with_remaining_chunks(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
"""Cover line 769: collecting remaining chunks when entry is done."""
|
|
from agents.extensions.sandbox.blaxel.sandbox import _BlaxelPtySessionEntry
|
|
|
|
session = _make_session(fake_sandbox)
|
|
entry = _BlaxelPtySessionEntry(
|
|
ws_session_id="done-chunks",
|
|
ws=None,
|
|
http_session=None,
|
|
done=True,
|
|
)
|
|
# Add chunks after marking done, to test the inner drain loop.
|
|
entry.output_chunks.append(b"chunk1")
|
|
entry.output_chunks.append(b"chunk2")
|
|
|
|
output, _ = await session._collect_pty_output(
|
|
entry=entry, yield_time_ms=5000, max_output_tokens=None
|
|
)
|
|
assert b"chunk1" in output
|
|
assert b"chunk2" in output
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Mounts tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _FakeExecResultForMount:
|
|
def __init__(self, exit_code: int = 0, stdout: bytes = b"", stderr: bytes = b"") -> None:
|
|
self.exit_code = exit_code
|
|
self.stdout = stdout
|
|
self.stderr = stderr
|
|
|
|
|
|
class _FakeMountSession:
|
|
"""Minimal BaseSandboxSession stand-in for mount tests."""
|
|
|
|
__name__ = "BlaxelSandboxSession"
|
|
|
|
def __init__(self) -> None:
|
|
self.exec_calls: list[tuple[tuple[str, ...], dict[str, float]]] = []
|
|
self._next_results: list[_FakeExecResultForMount] = []
|
|
self._default_result = _FakeExecResultForMount()
|
|
|
|
async def exec(self, *cmd: str, timeout: float = 120) -> _FakeExecResultForMount:
|
|
self.exec_calls.append((cmd, {"timeout": timeout}))
|
|
if self._next_results:
|
|
return self._next_results.pop(0)
|
|
return self._default_result
|
|
|
|
class __class__:
|
|
__name__ = "BlaxelSandboxSession"
|
|
|
|
|
|
# Override type name for _assert_blaxel_session check.
|
|
_FakeMountSession.__name__ = "BlaxelSandboxSession"
|
|
|
|
|
|
def _bl_strategy() -> Any:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelCloudBucketMountStrategy
|
|
|
|
return BlaxelCloudBucketMountStrategy()
|
|
|
|
|
|
class TestMountsModule:
|
|
def test_build_mount_config_s3(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _build_mount_config
|
|
from agents.sandbox.entries import S3Mount
|
|
|
|
mount = S3Mount(
|
|
bucket="my-bucket",
|
|
mount_strategy=_bl_strategy(),
|
|
access_key_id="AKID",
|
|
secret_access_key="SECRET",
|
|
region="us-east-1",
|
|
prefix="data/",
|
|
read_only=True,
|
|
)
|
|
config = _build_mount_config(mount, mount_path="/mnt/s3")
|
|
assert config.provider == "s3"
|
|
assert config.bucket == "my-bucket"
|
|
assert config.mount_path == "/mnt/s3"
|
|
assert config.access_key_id == "AKID"
|
|
assert config.region == "us-east-1"
|
|
assert config.prefix == "data/"
|
|
|
|
def test_build_mount_config_r2(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _build_mount_config
|
|
from agents.sandbox.entries import R2Mount
|
|
|
|
mount = R2Mount(
|
|
bucket="r2-bucket",
|
|
mount_strategy=_bl_strategy(),
|
|
account_id="acc123",
|
|
access_key_id="R2KEY",
|
|
secret_access_key="R2SECRET",
|
|
)
|
|
config = _build_mount_config(mount, mount_path="/mnt/r2")
|
|
assert config.provider == "r2"
|
|
assert "r2.cloudflarestorage.com" in (config.endpoint_url or "")
|
|
|
|
def test_build_mount_config_r2_custom_domain(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _build_mount_config
|
|
from agents.sandbox.entries import R2Mount
|
|
|
|
mount = R2Mount(
|
|
bucket="r2-bucket",
|
|
account_id="acc123",
|
|
mount_strategy=_bl_strategy(),
|
|
access_key_id="R2KEY",
|
|
secret_access_key="R2SECRET",
|
|
custom_domain="https://custom.example.com",
|
|
)
|
|
config = _build_mount_config(mount, mount_path="/mnt/r2")
|
|
assert config.endpoint_url == "https://custom.example.com"
|
|
|
|
def test_build_mount_config_gcs(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _build_mount_config
|
|
from agents.sandbox.entries import GCSMount
|
|
|
|
mount = GCSMount(
|
|
bucket="gcs-bucket",
|
|
mount_strategy=_bl_strategy(),
|
|
service_account_credentials='{"type":"service_account"}',
|
|
prefix="prefix/",
|
|
)
|
|
config = _build_mount_config(mount, mount_path="/mnt/gcs")
|
|
assert config.provider == "gcs"
|
|
assert config.service_account_key is not None
|
|
|
|
def test_build_mount_config_gcs_hmac(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _build_mount_config
|
|
from agents.sandbox.entries import GCSMount
|
|
|
|
mount = GCSMount(
|
|
bucket="gcs-bucket",
|
|
mount_strategy=_bl_strategy(),
|
|
access_id="GOOG1",
|
|
secret_access_key="SECRET",
|
|
endpoint_url="https://storage.googleapis.com",
|
|
prefix="prefix/",
|
|
)
|
|
config = _build_mount_config(mount, mount_path="/mnt/gcs")
|
|
assert config.provider == "s3"
|
|
assert config.access_key_id == "GOOG1"
|
|
assert config.secret_access_key == "SECRET"
|
|
assert config.endpoint_url == "https://storage.googleapis.com"
|
|
assert config.prefix == "prefix/"
|
|
|
|
def test_build_mount_config_unsupported(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _build_mount_config
|
|
from agents.sandbox.errors import MountConfigError
|
|
|
|
# Use a MagicMock with a type attribute to simulate an unsupported mount.
|
|
mount = MagicMock()
|
|
mount.type = "unsupported_mount"
|
|
with pytest.raises(MountConfigError, match="only support"):
|
|
_build_mount_config(mount, mount_path="/mnt/x")
|
|
|
|
def test_assert_blaxel_session_wrong_type(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _assert_blaxel_session
|
|
from agents.sandbox.errors import MountConfigError
|
|
|
|
class _WrongSession:
|
|
pass
|
|
|
|
with pytest.raises(MountConfigError, match="BlaxelSandboxSession"):
|
|
_assert_blaxel_session(_WrongSession()) # type: ignore[arg-type]
|
|
|
|
def test_validate_mount(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelCloudBucketMountStrategy
|
|
from agents.sandbox.entries import S3Mount
|
|
|
|
strategy = BlaxelCloudBucketMountStrategy()
|
|
mount = S3Mount(bucket="test-bucket", mount_strategy=_bl_strategy())
|
|
strategy.validate_mount(mount)
|
|
|
|
def test_build_docker_volume_driver_config_returns_none(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelCloudBucketMountStrategy
|
|
from agents.sandbox.entries import S3Mount
|
|
|
|
strategy = BlaxelCloudBucketMountStrategy()
|
|
mount = S3Mount(bucket="test", mount_strategy=_bl_strategy())
|
|
assert strategy.build_docker_volume_driver_config(mount) is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mount_s3_with_credentials(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelCloudBucketMountConfig, _mount_s3
|
|
|
|
session = _FakeMountSession()
|
|
# Simulate: which s3fs succeeds.
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0, stdout=b"/usr/bin/s3fs"), # which s3fs
|
|
_FakeExecResultForMount(exit_code=0), # write cred file
|
|
_FakeExecResultForMount(exit_code=0), # mkdir
|
|
_FakeExecResultForMount(exit_code=0), # s3fs mount
|
|
_FakeExecResultForMount(exit_code=0), # rm cred file
|
|
]
|
|
|
|
config = BlaxelCloudBucketMountConfig(
|
|
provider="s3",
|
|
bucket="my-bucket",
|
|
mount_path="/mnt/s3",
|
|
access_key_id="AKID",
|
|
secret_access_key="SECRET",
|
|
region="us-east-1",
|
|
prefix="data/",
|
|
read_only=True,
|
|
)
|
|
await _mount_s3(session, config) # type: ignore[arg-type]
|
|
assert len(session.exec_calls) == 5
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mount_s3_public_bucket(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelCloudBucketMountConfig, _mount_s3
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0), # which s3fs
|
|
_FakeExecResultForMount(exit_code=0), # mkdir
|
|
_FakeExecResultForMount(exit_code=0), # s3fs mount (no cred cleanup)
|
|
]
|
|
|
|
config = BlaxelCloudBucketMountConfig(
|
|
provider="s3",
|
|
bucket="public-bucket",
|
|
mount_path="/mnt/pub",
|
|
read_only=True,
|
|
)
|
|
await _mount_s3(session, config) # type: ignore[arg-type]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mount_s3_with_endpoint(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelCloudBucketMountConfig, _mount_s3
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0), # which s3fs
|
|
_FakeExecResultForMount(exit_code=0), # mkdir
|
|
_FakeExecResultForMount(exit_code=0), # s3fs mount
|
|
]
|
|
|
|
config = BlaxelCloudBucketMountConfig(
|
|
provider="s3",
|
|
bucket="endpoint-bucket",
|
|
mount_path="/mnt/ep",
|
|
endpoint_url="https://custom-s3.example.com",
|
|
)
|
|
await _mount_s3(session, config) # type: ignore[arg-type]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mount_s3_r2_sigv4(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelCloudBucketMountConfig, _mount_s3
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0), # which s3fs
|
|
_FakeExecResultForMount(exit_code=0), # write cred
|
|
_FakeExecResultForMount(exit_code=0), # mkdir
|
|
_FakeExecResultForMount(exit_code=0), # s3fs mount
|
|
_FakeExecResultForMount(exit_code=0), # rm cred
|
|
]
|
|
|
|
config = BlaxelCloudBucketMountConfig(
|
|
provider="r2",
|
|
bucket="r2-bucket",
|
|
mount_path="/mnt/r2",
|
|
access_key_id="KEY",
|
|
secret_access_key="SECRET",
|
|
endpoint_url="https://acc.r2.cloudflarestorage.com",
|
|
)
|
|
await _mount_s3(session, config) # type: ignore[arg-type]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mount_s3_fails(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelCloudBucketMountConfig, _mount_s3
|
|
from agents.sandbox.errors import MountConfigError
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0), # which s3fs
|
|
_FakeExecResultForMount(exit_code=0), # mkdir
|
|
_FakeExecResultForMount(exit_code=1, stderr=b"mount error"), # s3fs fails
|
|
]
|
|
|
|
config = BlaxelCloudBucketMountConfig(
|
|
provider="s3",
|
|
bucket="fail-bucket",
|
|
mount_path="/mnt/fail",
|
|
)
|
|
with pytest.raises(MountConfigError, match="s3fs mount failed"):
|
|
await _mount_s3(session, config) # type: ignore[arg-type]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mount_gcs_with_key(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelCloudBucketMountConfig, _mount_gcs
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0), # which gcsfuse
|
|
_FakeExecResultForMount(exit_code=0), # write key
|
|
_FakeExecResultForMount(exit_code=0), # mkdir
|
|
_FakeExecResultForMount(exit_code=0), # gcsfuse mount
|
|
_FakeExecResultForMount(exit_code=0), # rm key
|
|
]
|
|
|
|
config = BlaxelCloudBucketMountConfig(
|
|
provider="gcs",
|
|
bucket="gcs-bucket",
|
|
mount_path="/mnt/gcs",
|
|
service_account_key='{"type":"service_account"}',
|
|
read_only=True,
|
|
prefix="data/",
|
|
)
|
|
await _mount_gcs(session, config) # type: ignore[arg-type]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mount_gcs_anonymous(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelCloudBucketMountConfig, _mount_gcs
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0), # which gcsfuse
|
|
_FakeExecResultForMount(exit_code=0), # mkdir
|
|
_FakeExecResultForMount(exit_code=0), # gcsfuse mount
|
|
]
|
|
|
|
config = BlaxelCloudBucketMountConfig(
|
|
provider="gcs",
|
|
bucket="pub-gcs",
|
|
mount_path="/mnt/pub-gcs",
|
|
)
|
|
await _mount_gcs(session, config) # type: ignore[arg-type]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mount_gcs_fails(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelCloudBucketMountConfig, _mount_gcs
|
|
from agents.sandbox.errors import MountConfigError
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0), # which gcsfuse
|
|
_FakeExecResultForMount(exit_code=0), # mkdir
|
|
_FakeExecResultForMount(exit_code=1, stderr=b"gcs error"), # fails
|
|
]
|
|
|
|
config = BlaxelCloudBucketMountConfig(
|
|
provider="gcs",
|
|
bucket="fail-gcs",
|
|
mount_path="/mnt/fail-gcs",
|
|
)
|
|
with pytest.raises(MountConfigError, match="gcsfuse mount failed"):
|
|
await _mount_gcs(session, config) # type: ignore[arg-type]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mount_bucket_dispatch_s3(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import (
|
|
BlaxelCloudBucketMountConfig,
|
|
_mount_bucket,
|
|
)
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0), # which s3fs
|
|
_FakeExecResultForMount(exit_code=0), # mkdir
|
|
_FakeExecResultForMount(exit_code=0), # s3fs mount
|
|
]
|
|
config = BlaxelCloudBucketMountConfig(provider="s3", bucket="b", mount_path="/m")
|
|
await _mount_bucket(session, config) # type: ignore[arg-type]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mount_bucket_dispatch_gcs(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import (
|
|
BlaxelCloudBucketMountConfig,
|
|
_mount_bucket,
|
|
)
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0),
|
|
_FakeExecResultForMount(exit_code=0),
|
|
_FakeExecResultForMount(exit_code=0),
|
|
]
|
|
config = BlaxelCloudBucketMountConfig(provider="gcs", bucket="b", mount_path="/m")
|
|
await _mount_bucket(session, config) # type: ignore[arg-type]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unmount_bucket_fusermount(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _unmount_bucket
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0), # fusermount succeeds
|
|
]
|
|
await _unmount_bucket(session, "/mnt/test") # type: ignore[arg-type]
|
|
assert len(session.exec_calls) == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unmount_bucket_umount_fallback(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _unmount_bucket
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=1), # fusermount fails
|
|
_FakeExecResultForMount(exit_code=0), # umount succeeds
|
|
]
|
|
await _unmount_bucket(session, "/mnt/test") # type: ignore[arg-type]
|
|
assert len(session.exec_calls) == 2
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unmount_bucket_lazy(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _unmount_bucket
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=1), # fusermount fails
|
|
_FakeExecResultForMount(exit_code=1), # umount fails
|
|
_FakeExecResultForMount(exit_code=0), # umount -l
|
|
]
|
|
await _unmount_bucket(session, "/mnt/test") # type: ignore[arg-type]
|
|
assert len(session.exec_calls) == 3
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_install_tool_with_apk(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _install_tool
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0, stdout=b"apk"), # detect pkg mgr
|
|
_FakeExecResultForMount(exit_code=0), # apk add succeeds
|
|
]
|
|
await _install_tool(session, "s3fs") # type: ignore[arg-type]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_install_tool_with_apt(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _install_tool
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0, stdout=b"apt"), # detect pkg mgr
|
|
_FakeExecResultForMount(exit_code=0), # apt-get install succeeds
|
|
]
|
|
await _install_tool(session, "gcsfuse") # type: ignore[arg-type]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_install_tool_fails_after_retries(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _install_tool
|
|
from agents.sandbox.errors import MountConfigError
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0, stdout=b"apt"), # detect
|
|
_FakeExecResultForMount(exit_code=1), # attempt 1
|
|
_FakeExecResultForMount(exit_code=1), # attempt 2
|
|
_FakeExecResultForMount(exit_code=1), # attempt 3
|
|
]
|
|
with pytest.raises(MountConfigError, match="failed to install"):
|
|
await _install_tool(session, "s3fs") # type: ignore[arg-type]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_tool_already_installed(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _ensure_tool
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0), # which s3fs succeeds
|
|
]
|
|
await _ensure_tool(session, "s3fs") # type: ignore[arg-type]
|
|
assert len(session.exec_calls) == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ensure_tool_needs_install(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _ensure_tool
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=1), # which fails
|
|
_FakeExecResultForMount(exit_code=0, stdout=b"apt"), # detect
|
|
_FakeExecResultForMount(exit_code=0), # install
|
|
]
|
|
await _ensure_tool(session, "s3fs") # type: ignore[arg-type]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_activate(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelCloudBucketMountStrategy
|
|
from agents.sandbox.entries import S3Mount
|
|
|
|
strategy = BlaxelCloudBucketMountStrategy()
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0), # which
|
|
_FakeExecResultForMount(exit_code=0), # mkdir
|
|
_FakeExecResultForMount(exit_code=0), # mount
|
|
]
|
|
mount = S3Mount(bucket="test", mount_strategy=_bl_strategy(), mount_path=Path("/mnt/s3"))
|
|
# activate needs a real mount path resolution, mock it.
|
|
mount._resolve_mount_path = lambda s, d: Path("/workspace/mnt/s3") # type: ignore[assignment]
|
|
result = await strategy.activate(
|
|
mount,
|
|
session, # type: ignore[arg-type]
|
|
Path("/workspace/mnt/s3"),
|
|
Path("/workspace"),
|
|
)
|
|
assert result == []
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deactivate(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelCloudBucketMountStrategy
|
|
from agents.sandbox.entries import S3Mount
|
|
|
|
strategy = BlaxelCloudBucketMountStrategy()
|
|
session = _FakeMountSession()
|
|
session._next_results = [_FakeExecResultForMount(exit_code=0)]
|
|
mount = S3Mount(bucket="test", mount_strategy=_bl_strategy(), mount_path=Path("/mnt/s3"))
|
|
mount._resolve_mount_path = lambda s, d: Path("/workspace/mnt/s3") # type: ignore[assignment]
|
|
await strategy.deactivate(
|
|
mount,
|
|
session, # type: ignore[arg-type]
|
|
Path("/workspace/mnt/s3"),
|
|
Path("/workspace"),
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_teardown_for_snapshot(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelCloudBucketMountStrategy
|
|
from agents.sandbox.entries import S3Mount
|
|
|
|
strategy = BlaxelCloudBucketMountStrategy()
|
|
session = _FakeMountSession()
|
|
session._next_results = [_FakeExecResultForMount(exit_code=0)]
|
|
mount = S3Mount(bucket="test", mount_strategy=_bl_strategy())
|
|
await strategy.teardown_for_snapshot(
|
|
mount,
|
|
session, # type: ignore[arg-type]
|
|
Path("/workspace/mnt/s3"),
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restore_after_snapshot(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelCloudBucketMountStrategy
|
|
from agents.sandbox.entries import S3Mount
|
|
|
|
strategy = BlaxelCloudBucketMountStrategy()
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=0), # which
|
|
_FakeExecResultForMount(exit_code=0), # mkdir
|
|
_FakeExecResultForMount(exit_code=0), # mount
|
|
]
|
|
mount = S3Mount(bucket="test", mount_strategy=_bl_strategy())
|
|
await strategy.restore_after_snapshot(
|
|
mount,
|
|
session, # type: ignore[arg-type]
|
|
Path("/workspace/mnt/s3"),
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# SDK exception mapping tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSdkExceptionMapping:
|
|
def test_import_sandbox_api_error(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _import_sandbox_api_error
|
|
|
|
cls = _import_sandbox_api_error()
|
|
if cls is None:
|
|
pytest.skip("blaxel not available")
|
|
assert issubclass(cls, BaseException)
|
|
|
|
def test_import_sandbox_api_error_missing_sdk(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.sandbox import _import_sandbox_api_error
|
|
|
|
with patch.dict(
|
|
"sys.modules",
|
|
{"blaxel": None, "blaxel.core": None, "blaxel.core.sandbox": None},
|
|
):
|
|
assert _import_sandbox_api_error() is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exec_maps_sdk_api_error_408_to_timeout(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
"""SandboxAPIError with status_code=408 should map to ExecTimeoutError."""
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
session = _make_session(fake_sandbox)
|
|
|
|
# Create a fake SandboxAPIError with status_code.
|
|
class FakeApiError(Exception):
|
|
def __init__(self, msg: str, status_code: int) -> None:
|
|
super().__init__(msg)
|
|
self.status_code = status_code
|
|
|
|
async def _raise_timeout(*args: object, **kw: object) -> None:
|
|
raise FakeApiError("request timeout", status_code=408)
|
|
|
|
fake_sandbox.process.exec = _raise_timeout # type: ignore[assignment]
|
|
|
|
with patch.object(mod, "_import_sandbox_api_error", return_value=FakeApiError):
|
|
with pytest.raises(ExecTimeoutError):
|
|
await session._exec_internal("sleep", "100")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exec_maps_sdk_api_error_504_to_timeout(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
"""SandboxAPIError with status_code=504 should map to ExecTimeoutError."""
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
session = _make_session(fake_sandbox)
|
|
|
|
class FakeApiError(Exception):
|
|
def __init__(self, msg: str, status_code: int) -> None:
|
|
super().__init__(msg)
|
|
self.status_code = status_code
|
|
|
|
async def _raise_504(*args: object, **kw: object) -> None:
|
|
raise FakeApiError("gateway timeout", status_code=504)
|
|
|
|
fake_sandbox.process.exec = _raise_504 # type: ignore[assignment]
|
|
|
|
with patch.object(mod, "_import_sandbox_api_error", return_value=FakeApiError):
|
|
with pytest.raises(ExecTimeoutError):
|
|
await session._exec_internal("sleep", "100")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exec_non_timeout_api_error_becomes_transport(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
"""SandboxAPIError with status_code=500 should map to ExecTransportError."""
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
session = _make_session(fake_sandbox)
|
|
|
|
class FakeApiError(Exception):
|
|
def __init__(self, msg: str, status_code: int) -> None:
|
|
super().__init__(msg)
|
|
self.status_code = status_code
|
|
|
|
async def _raise_500(*args: object, **kw: object) -> None:
|
|
raise FakeApiError("internal error", status_code=500)
|
|
|
|
fake_sandbox.process.exec = _raise_500 # type: ignore[assignment]
|
|
|
|
with patch.object(mod, "_import_sandbox_api_error", return_value=FakeApiError):
|
|
with pytest.raises(ExecTransportError) as exc_info:
|
|
await session._exec_internal("echo", "hello")
|
|
assert str(exc_info.value) == "Blaxel exec failed: HTTP 500: internal error"
|
|
assert exc_info.value.context["backend"] == "blaxel"
|
|
assert exc_info.value.context["http_status"] == 500
|
|
assert exc_info.value.context["provider_error"] == "HTTP 500: internal error"
|
|
assert exc_info.value.retryable is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exec_uses_structured_blaxel_non_retryable_error_code(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
session = _make_session(fake_sandbox)
|
|
|
|
class FakeApiError(Exception):
|
|
def __init__(self) -> None:
|
|
super().__init__("route not found")
|
|
self.status_code = 404
|
|
self.body = {
|
|
"error": {
|
|
"code": "ROUTE_NOT_FOUND",
|
|
"message": "Preview not found: sandbox",
|
|
"retryable": False,
|
|
"status": 404,
|
|
}
|
|
}
|
|
|
|
async def _raise_route_not_found(*args: object, **kw: object) -> None:
|
|
raise FakeApiError()
|
|
|
|
fake_sandbox.process.exec = _raise_route_not_found # type: ignore[assignment]
|
|
|
|
with patch.object(mod, "_import_sandbox_api_error", return_value=FakeApiError):
|
|
with pytest.raises(ExecTransportError) as exc_info:
|
|
await session._exec_internal("echo", "hello")
|
|
|
|
assert str(exc_info.value) == "Blaxel exec failed: HTTP 404: route not found"
|
|
assert exc_info.value.context["backend"] == "blaxel"
|
|
assert exc_info.value.context["http_status"] == 404
|
|
assert exc_info.value.context["provider_error"] == "HTTP 404: route not found"
|
|
assert exc_info.value.context["provider_error_code"] == "ROUTE_NOT_FOUND"
|
|
assert exc_info.value.retryable is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exec_uses_structured_blaxel_retryable_error_code(
|
|
self, fake_sandbox: _FakeSandboxInstance
|
|
) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
session = _make_session(fake_sandbox)
|
|
|
|
class FakeApiError(Exception):
|
|
def __init__(self) -> None:
|
|
super().__init__("workload unavailable")
|
|
self.status_code = 404
|
|
self.body = {
|
|
"error": {
|
|
"code": "WORKLOAD_UNAVAILABLE",
|
|
"message": "No healthy replica is serving workload",
|
|
"retryable": True,
|
|
"status": 404,
|
|
}
|
|
}
|
|
|
|
async def _raise_workload_unavailable(*args: object, **kw: object) -> None:
|
|
raise FakeApiError()
|
|
|
|
fake_sandbox.process.exec = _raise_workload_unavailable # type: ignore[assignment]
|
|
|
|
with patch.object(mod, "_import_sandbox_api_error", return_value=FakeApiError):
|
|
with pytest.raises(ExecTransportError) as exc_info:
|
|
await session._exec_internal("echo", "hello")
|
|
|
|
assert str(exc_info.value) == "Blaxel exec failed: HTTP 404: workload unavailable"
|
|
assert exc_info.value.context["backend"] == "blaxel"
|
|
assert exc_info.value.context["http_status"] == 404
|
|
assert exc_info.value.context["provider_error"] == "HTTP 404: workload unavailable"
|
|
assert exc_info.value.context["provider_error_code"] == "WORKLOAD_UNAVAILABLE"
|
|
assert exc_info.value.retryable is True
|
|
|
|
@pytest.mark.parametrize(
|
|
("code", "expected_retryable"),
|
|
[
|
|
("ROUTE_NOT_FOUND", False),
|
|
("WORKLOAD_NOT_FOUND", False),
|
|
("WORKSPACE_NOT_FOUND", False),
|
|
("WORKLOAD_UNAVAILABLE", True),
|
|
("AUTHENTICATION_REQUIRED", False),
|
|
("AUTHENTICATION_FAILED", False),
|
|
("FORBIDDEN", False),
|
|
("BAD_REQUEST", False),
|
|
("USAGE_LIMIT_EXCEEDED", False),
|
|
("POLICY_VIOLATION", False),
|
|
],
|
|
)
|
|
def test_blaxel_retryability_error_code_table(
|
|
self,
|
|
code: str,
|
|
expected_retryable: bool,
|
|
) -> None:
|
|
from agents.extensions.sandbox.blaxel import sandbox as mod
|
|
|
|
class FakeApiError(Exception):
|
|
def __init__(self) -> None:
|
|
super().__init__(code)
|
|
self.body = {"error": {"code": code, "message": code}}
|
|
|
|
retryable, provider_error_code = mod._blaxel_provider_retryability(FakeApiError())
|
|
|
|
assert retryable is expected_retryable
|
|
assert provider_error_code == code
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Timeout coercion tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCoerceExecTimeout:
|
|
def test_none_returns_default(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
result = session._coerce_exec_timeout(None)
|
|
assert result == 300.0 # Default from BlaxelTimeouts.
|
|
|
|
def test_positive_value_passthrough(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
assert session._coerce_exec_timeout(42.5) == 42.5
|
|
|
|
def test_zero_returns_small_positive(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
assert session._coerce_exec_timeout(0) == 0.001
|
|
|
|
def test_negative_returns_small_positive(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
session = _make_session(fake_sandbox)
|
|
assert session._coerce_exec_timeout(-5) == 0.001
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Drive mount tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestDriveMounts:
|
|
@pytest.mark.asyncio
|
|
async def test_attach_drive_success(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelDriveMountConfig, _attach_drive
|
|
|
|
sandbox = _FakeSandboxInstance()
|
|
config = BlaxelDriveMountConfig(
|
|
drive_name="test-drive", mount_path="/mnt/data", drive_path="/"
|
|
)
|
|
await _attach_drive(sandbox, config)
|
|
assert sandbox.drives.mount_calls == [("test-drive", "/mnt/data", "/")]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_attach_drive_error(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelDriveMountConfig, _attach_drive
|
|
from agents.sandbox.errors import MountConfigError
|
|
|
|
sandbox = _FakeSandboxInstance()
|
|
sandbox.drives.mount_error = RuntimeError("mount api error")
|
|
config = BlaxelDriveMountConfig(
|
|
drive_name="test-drive", mount_path="/mnt/data", drive_path="/"
|
|
)
|
|
with pytest.raises(MountConfigError, match="drive mount failed"):
|
|
await _attach_drive(sandbox, config)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_attach_drive_no_drives_api(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelDriveMountConfig, _attach_drive
|
|
from agents.sandbox.errors import MountConfigError
|
|
|
|
class _NoDrives:
|
|
pass
|
|
|
|
config = BlaxelDriveMountConfig(
|
|
drive_name="test-drive", mount_path="/mnt/data", drive_path="/"
|
|
)
|
|
with pytest.raises(MountConfigError, match="does not expose a drives API"):
|
|
await _attach_drive(_NoDrives(), config)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detach_drive_success(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _detach_drive
|
|
|
|
sandbox = _FakeSandboxInstance()
|
|
await _detach_drive(sandbox, "/mnt/data")
|
|
assert sandbox.drives.unmount_calls == ["/mnt/data"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detach_drive_error_logged_not_raised(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _detach_drive
|
|
|
|
sandbox = _FakeSandboxInstance()
|
|
sandbox.drives.unmount_error = RuntimeError("unmount failed")
|
|
# Should not raise; error is logged.
|
|
await _detach_drive(sandbox, "/mnt/data")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detach_drive_no_drives_api(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _detach_drive
|
|
|
|
class _NoDrives:
|
|
pass
|
|
|
|
# Should not raise when drives API is missing.
|
|
await _detach_drive(_NoDrives(), "/mnt/data")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_drive_strategy_validate_wrong_mount_type(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelDriveMountStrategy
|
|
from agents.sandbox.errors import MountConfigError
|
|
|
|
strategy = BlaxelDriveMountStrategy()
|
|
mount = MagicMock()
|
|
mount.type = "blaxel_drive"
|
|
with pytest.raises(MountConfigError, match="BlaxelDriveMount"):
|
|
strategy.validate_mount(mount)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_drive_strategy_validate_non_drive_mount(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelDriveMountStrategy
|
|
from agents.sandbox.errors import MountConfigError
|
|
|
|
strategy = BlaxelDriveMountStrategy()
|
|
mount = MagicMock()
|
|
mount.type = "s3_mount"
|
|
with pytest.raises(MountConfigError, match="BlaxelDriveMount"):
|
|
strategy.validate_mount(mount)
|
|
|
|
def test_drive_strategy_build_docker_volume_returns_none(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import BlaxelDriveMountStrategy
|
|
|
|
strategy = BlaxelDriveMountStrategy()
|
|
mount = MagicMock()
|
|
assert strategy.build_docker_volume_driver_config(mount) is None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Unmount bucket stderr logging tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestUnmountBucketLogging:
|
|
@pytest.mark.asyncio
|
|
async def test_unmount_all_attempts_fail_logs_warning(self) -> None:
|
|
from agents.extensions.sandbox.blaxel.mounts import _unmount_bucket
|
|
|
|
session = _FakeMountSession()
|
|
session._next_results = [
|
|
_FakeExecResultForMount(exit_code=1), # fusermount fails
|
|
_FakeExecResultForMount(exit_code=1), # umount fails
|
|
_FakeExecResultForMount(exit_code=1), # umount -l fails
|
|
]
|
|
# Should not raise, just log warning.
|
|
await _unmount_bucket(session, "/mnt/test") # type: ignore[arg-type]
|
|
assert len(session.exec_calls) == 3
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# FakeFs.ls improvement tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFakeFs:
|
|
@pytest.mark.asyncio
|
|
async def test_ls_returns_matching_paths(self) -> None:
|
|
fs = _FakeFs()
|
|
fs.files["/workspace/a.txt"] = b"a"
|
|
fs.files["/workspace/b.txt"] = b"b"
|
|
fs.files["/other/c.txt"] = b"c"
|
|
result = await fs.ls("/workspace")
|
|
assert "/workspace/a.txt" in result
|
|
assert "/workspace/b.txt" in result
|
|
assert "/other/c.txt" not in result
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ls_empty_returns_path(self) -> None:
|
|
fs = _FakeFs()
|
|
result = await fs.ls("/empty")
|
|
assert result == ["/empty"]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Shutdown logging tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestShutdownLogging:
|
|
@pytest.mark.asyncio
|
|
async def test_shutdown_delete_logs_warning(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
"""shutdown() should log a warning when delete fails, not silently suppress."""
|
|
session = _make_session(fake_sandbox)
|
|
|
|
async def _raise() -> None:
|
|
raise RuntimeError("delete failed")
|
|
|
|
fake_sandbox.delete = _raise # type: ignore[method-assign]
|
|
# Should not raise.
|
|
await session.shutdown()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_running_false_logs_debug(self, fake_sandbox: _FakeSandboxInstance) -> None:
|
|
"""running() should log at debug level when health check fails."""
|
|
session = _make_session(fake_sandbox)
|
|
|
|
async def _raise(*args: object, **kw: object) -> None:
|
|
raise ConnectionError("offline")
|
|
|
|
fake_sandbox.fs.ls = _raise # type: ignore[assignment]
|
|
assert await session.running() is False
|