wehub-resource-sync
841 lines
29 KiB
Python
841 lines
29 KiB
Python
"""Unit tests for ``E2BSandboxProvider`` and its companion ``E2BSandbox``."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
import threading
|
|
from collections import OrderedDict
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from deerflow.config.paths import Paths
|
|
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
# Fakes for the e2b SDK
|
|
# ──────────────────────────────────────────────────────────────────────────────
|
|
|
|
|
|
class FakeCommandsAPI:
|
|
"""Stand-in for ``client.commands``."""
|
|
|
|
GONE = "__GONE__"
|
|
NOT_FOUND_MSG = "The sandbox was not found: This error is likely due to sandbox timeout."
|
|
|
|
def __init__(self, responses: list[Any] | None = None) -> None:
|
|
self.calls: list[str] = []
|
|
self._responses = list(responses or [])
|
|
|
|
def _next(self) -> Any:
|
|
if not self._responses:
|
|
return SimpleNamespace(stdout="BOOTSTRAP_OK", stderr="", exit_code=0)
|
|
head = self._responses.pop(0)
|
|
return head
|
|
|
|
def run(self, cmd: str, envs: dict[str, str] | None = None, **kwargs) -> SimpleNamespace:
|
|
self.calls.append(cmd)
|
|
self.envs = getattr(self, "envs", [])
|
|
self.envs.append(envs)
|
|
head = self._next()
|
|
if head == self.GONE:
|
|
raise RuntimeError(self.NOT_FOUND_MSG)
|
|
if callable(head):
|
|
return head(cmd)
|
|
if isinstance(head, SimpleNamespace):
|
|
return head
|
|
return SimpleNamespace(stdout=str(head), stderr="", exit_code=0)
|
|
|
|
|
|
class _FakeFileStream:
|
|
"""Minimal stand-in for ``e2b.FileStreamReader``.
|
|
|
|
Yields fixed-size chunks and tracks whether ``close()`` was invoked so
|
|
tests can assert we release the connection on both success and abort.
|
|
"""
|
|
|
|
def __init__(self, data: bytes, *, chunk_size: int = 4096) -> None:
|
|
self._data = bytes(data)
|
|
self._chunk_size = max(1, int(chunk_size))
|
|
self._offset = 0
|
|
self.closed = False
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def __next__(self) -> bytes:
|
|
if self.closed or self._offset >= len(self._data):
|
|
raise StopIteration
|
|
end = min(self._offset + self._chunk_size, len(self._data))
|
|
chunk = self._data[self._offset : end]
|
|
self._offset = end
|
|
return chunk
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
def __enter__(self) -> _FakeFileStream:
|
|
return self
|
|
|
|
def __exit__(self, *exc_info) -> None:
|
|
self.close()
|
|
|
|
|
|
class FakeFilesAPI:
|
|
def __init__(
|
|
self,
|
|
store: dict[str, bytes] | None = None,
|
|
*,
|
|
stream_chunk_size: int = 4096,
|
|
) -> None:
|
|
self.store = dict(store or {})
|
|
self.read_calls: list[tuple[str, str | None]] = []
|
|
self.write_calls: list[tuple[str, bytes]] = []
|
|
self.streams: list[_FakeFileStream] = []
|
|
self._stream_chunk_size = stream_chunk_size
|
|
|
|
def read(self, path: str, *, format: str | None = None):
|
|
self.read_calls.append((path, format))
|
|
if path not in self.store:
|
|
raise FileNotFoundError(path)
|
|
data = self.store[path]
|
|
if format == "bytes":
|
|
return data
|
|
if format == "stream":
|
|
stream = _FakeFileStream(data, chunk_size=self._stream_chunk_size)
|
|
self.streams.append(stream)
|
|
return stream
|
|
try:
|
|
return data.decode("utf-8")
|
|
except UnicodeDecodeError:
|
|
return data
|
|
|
|
def write(self, path: str, content: bytes) -> None:
|
|
self.write_calls.append((path, content))
|
|
self.store[path] = content
|
|
|
|
|
|
class FakeClient:
|
|
"""Lightweight ``e2b.Sandbox`` substitute used by the provider tests."""
|
|
|
|
def __init__(
|
|
self,
|
|
sandbox_id: str = "fake-sb-1",
|
|
*,
|
|
commands: FakeCommandsAPI | None = None,
|
|
files: FakeFilesAPI | None = None,
|
|
) -> None:
|
|
self.sandbox_id = sandbox_id
|
|
self.commands = commands or FakeCommandsAPI()
|
|
self.files = files or FakeFilesAPI()
|
|
self.timeouts_set: list[int] = []
|
|
self.killed = False
|
|
self.closed = False
|
|
|
|
def set_timeout(self, seconds: int) -> None:
|
|
self.timeouts_set.append(int(seconds))
|
|
|
|
def kill(self) -> None:
|
|
self.killed = True
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class FakeSandboxClass:
|
|
"""Stand-in for ``e2b_code_interpreter.Sandbox`` (the class itself)."""
|
|
|
|
def __init__(self) -> None:
|
|
self.create_calls: list[dict[str, Any]] = []
|
|
self.connect_calls: list[tuple[str, dict[str, Any]]] = []
|
|
self.list_calls: list[dict[str, Any]] = []
|
|
self.create_factory = lambda **kw: FakeClient(sandbox_id=f"created-{len(self.create_calls)}")
|
|
self.connect_factory = lambda sid, **kw: FakeClient(sandbox_id=sid)
|
|
self.list_return: Any = []
|
|
|
|
def create(self, **kwargs: Any) -> FakeClient:
|
|
self.create_calls.append(kwargs)
|
|
return self.create_factory(**kwargs)
|
|
|
|
def connect(self, sandbox_id: str, **kwargs: Any) -> FakeClient:
|
|
self.connect_calls.append((sandbox_id, kwargs))
|
|
return self.connect_factory(sandbox_id, **kwargs)
|
|
|
|
def list(self, **kwargs: Any) -> Any:
|
|
self.list_calls.append(kwargs)
|
|
return self.list_return
|
|
|
|
|
|
def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800) -> Any:
|
|
"""Build a ``E2BSandboxProvider`` instance bypassing ``__init__``."""
|
|
mod = importlib.import_module("deerflow.community.e2b_sandbox.e2b_sandbox_provider")
|
|
provider = mod.E2BSandboxProvider.__new__(mod.E2BSandboxProvider)
|
|
provider._lock = threading.Lock()
|
|
provider._sandboxes = {}
|
|
provider._thread_sandboxes = {}
|
|
provider._thread_locks = {}
|
|
provider._warm_pool = OrderedDict()
|
|
provider._shutdown_called = False
|
|
provider._config = {
|
|
"api_key": "test-key",
|
|
"template": "code-interpreter-v1",
|
|
"domain": None,
|
|
"home_dir": "/home/user",
|
|
"idle_timeout": idle_timeout,
|
|
"replicas": replicas,
|
|
"mounts": [],
|
|
"environment": {},
|
|
}
|
|
return provider
|
|
|
|
|
|
def _install_fake_sdk(monkeypatch, provider) -> FakeSandboxClass:
|
|
fake_cls = FakeSandboxClass()
|
|
monkeypatch.setattr(provider, "_get_sandbox_cls", lambda: fake_cls)
|
|
return fake_cls
|
|
|
|
|
|
def _make_sandbox(client: FakeClient, *, sandbox_id: str | None = None) -> Any:
|
|
mod = importlib.import_module("deerflow.community.e2b_sandbox.e2b_sandbox")
|
|
return mod.E2BSandbox(
|
|
id=sandbox_id or client.sandbox_id,
|
|
client=client,
|
|
home_dir="/home/user",
|
|
)
|
|
|
|
|
|
def test_thread_key_returns_user_thread_tuple():
|
|
p = _make_provider()
|
|
assert p._thread_key("t1", "u1") == ("u1", "t1")
|
|
|
|
|
|
def test_stable_seed_is_deterministic_and_user_scoped():
|
|
p = _make_provider()
|
|
s_a = p._stable_seed("t1", "u1")
|
|
s_b = p._stable_seed("t1", "u1")
|
|
s_other_user = p._stable_seed("t1", "u2")
|
|
s_other_thread = p._stable_seed("t2", "u1")
|
|
assert s_a == s_b
|
|
assert s_a != s_other_user
|
|
assert s_a != s_other_thread
|
|
|
|
|
|
def test_is_sandbox_gone_error_matches_known_signatures():
|
|
mod = importlib.import_module("deerflow.community.e2b_sandbox.e2b_sandbox")
|
|
f = mod._is_sandbox_gone_error
|
|
assert f(RuntimeError("Paused sandbox abcdef not found"))
|
|
assert f(Exception("The sandbox was not found: due to timeout"))
|
|
assert f(Exception("sandbox not found"))
|
|
# Unrelated errors must not flip the dead flag.
|
|
assert not f(Exception("Connection reset by peer"))
|
|
assert not f(ValueError("invalid path"))
|
|
|
|
|
|
def test_execute_command_marks_dead_on_sandbox_gone_error():
|
|
client = FakeClient(commands=FakeCommandsAPI([FakeCommandsAPI.GONE]))
|
|
sb = _make_sandbox(client)
|
|
out = sb.execute_command("echo hi")
|
|
assert "Error: " in out and "sandbox was not found" in out
|
|
assert sb.is_dead is True
|
|
out2 = sb.execute_command("echo again")
|
|
assert "reaped" in out2.lower()
|
|
assert client.commands.calls == ["echo hi"]
|
|
|
|
|
|
def test_execute_command_returns_stdout_on_success():
|
|
client = FakeClient(commands=FakeCommandsAPI([SimpleNamespace(stdout="hello\n", stderr="", exit_code=0)]))
|
|
sb = _make_sandbox(client)
|
|
assert sb.execute_command("printf hello").rstrip() == "hello"
|
|
assert sb.is_dead is False
|
|
|
|
|
|
def test_execute_command_does_not_mark_dead_on_unrelated_error():
|
|
|
|
def boom(_cmd: str, **kwargs) -> Any:
|
|
raise RuntimeError("Connection reset by peer")
|
|
|
|
client = FakeClient(commands=FakeCommandsAPI([boom]))
|
|
sb = _make_sandbox(client)
|
|
out = sb.execute_command("echo hi")
|
|
assert "Error" in out
|
|
assert sb.is_dead is False
|
|
|
|
|
|
def test_execute_command_forwards_env_and_timeout_to_commands_run():
|
|
"""execute_command(env=..., timeout=...) routes env as ``envs`` and the
|
|
timeout through to ``commands.run`` so request-scoped secrets (#3861) reach
|
|
the e2b subprocess without entering the command string. Regression for the
|
|
signature mismatch that broke bash for every e2b user."""
|
|
commands = MagicMock()
|
|
commands.run.return_value = SimpleNamespace(stdout="ok\n", stderr="", exit_code=0)
|
|
client = FakeClient(commands=commands)
|
|
sb = _make_sandbox(client)
|
|
|
|
out = sb.execute_command("echo $TOK", env={"TOK": "secret-v"}, timeout=120)
|
|
|
|
assert out.rstrip() == "ok"
|
|
args, kwargs = commands.run.call_args
|
|
assert args == ("echo $TOK",)
|
|
assert kwargs["envs"] == {"TOK": "secret-v"}
|
|
assert kwargs["timeout"] == 120
|
|
# The secret must not be smuggled into the command string.
|
|
assert "secret-v" not in args[0]
|
|
|
|
|
|
def test_execute_command_env_none_passes_no_envs_kwarg():
|
|
"""env=None is fully backward-compatible — ``commands.run`` is called with no
|
|
``envs``/``timeout`` kwargs, so existing (non-secret) callers are unaffected."""
|
|
commands = MagicMock()
|
|
commands.run.return_value = SimpleNamespace(stdout="ok\n", stderr="", exit_code=0)
|
|
client = FakeClient(commands=commands)
|
|
sb = _make_sandbox(client)
|
|
|
|
sb.execute_command("echo hi")
|
|
|
|
_, kwargs = commands.run.call_args
|
|
assert "envs" not in kwargs
|
|
assert "timeout" not in kwargs
|
|
|
|
|
|
def test_execute_command_forwards_env_as_envs():
|
|
"""Per-call ``env`` reaches the e2b SDK as ``envs`` so secrets like
|
|
``GITHUB_TOKEN`` are scoped to a single command without mutating shared
|
|
state. Mirrors the local/AIO sandboxes' overlay contract.
|
|
"""
|
|
client = FakeClient(commands=FakeCommandsAPI([SimpleNamespace(stdout="ok", stderr="", exit_code=0)]))
|
|
sb = _make_sandbox(client)
|
|
sb.execute_command("gh pr create", env={"GH_TOKEN": "tok-123"})
|
|
assert client.commands.envs == [{"GH_TOKEN": "tok-123"}]
|
|
|
|
|
|
def test_execute_command_rejects_invalid_env_key():
|
|
client = FakeClient(commands=FakeCommandsAPI([]))
|
|
sb = _make_sandbox(client)
|
|
with pytest.raises(ValueError, match="extra_env key"):
|
|
sb.execute_command("echo hi", env={"X;rm -rf /;Y": "v"})
|
|
# The SDK was never reached — validation happens before commands.run.
|
|
assert client.commands.calls == []
|
|
|
|
|
|
def test_ping_returns_false_when_sandbox_gone():
|
|
client = FakeClient(commands=FakeCommandsAPI([FakeCommandsAPI.GONE]))
|
|
sb = _make_sandbox(client)
|
|
assert sb.ping() is False
|
|
assert sb.is_dead is True
|
|
|
|
|
|
def test_ping_returns_true_on_unknown_error():
|
|
def boom(_cmd: str) -> Any:
|
|
raise RuntimeError("upstream timeout")
|
|
|
|
client = FakeClient(commands=FakeCommandsAPI([boom]))
|
|
sb = _make_sandbox(client)
|
|
assert sb.ping() is True
|
|
assert sb.is_dead is False
|
|
|
|
|
|
def test_client_alive_true_for_healthy_client():
|
|
p = _make_provider()
|
|
client = FakeClient()
|
|
assert p._client_alive(client) is True
|
|
assert client.commands.calls == ["true"]
|
|
|
|
|
|
def test_client_alive_false_when_sandbox_gone():
|
|
p = _make_provider()
|
|
client = FakeClient(commands=FakeCommandsAPI([FakeCommandsAPI.GONE]))
|
|
assert p._client_alive(client) is False
|
|
|
|
|
|
def test_client_alive_treats_unknown_errors_as_alive():
|
|
p = _make_provider()
|
|
|
|
def boom(_cmd: str) -> Any:
|
|
raise RuntimeError("flaky network")
|
|
|
|
client = FakeClient(commands=FakeCommandsAPI([boom]))
|
|
assert p._client_alive(client) is True
|
|
|
|
|
|
def test_safe_close_client_swallows_close_failures():
|
|
p = _make_provider()
|
|
|
|
class BadCloseClient:
|
|
def close(self) -> None:
|
|
raise RuntimeError("boom")
|
|
|
|
p._safe_close_client(BadCloseClient())
|
|
p._safe_close_client(None)
|
|
|
|
|
|
def test_kill_and_close_invokes_kill_and_close_in_order():
|
|
p = _make_provider()
|
|
client = FakeClient()
|
|
sb = _make_sandbox(client)
|
|
p._kill_and_close(sb)
|
|
assert client.killed is True
|
|
|
|
|
|
def test_kill_and_close_swallows_kill_exceptions():
|
|
p = _make_provider()
|
|
client = FakeClient()
|
|
sb = _make_sandbox(client)
|
|
|
|
def explode():
|
|
raise RuntimeError("already gone")
|
|
|
|
client.kill = explode
|
|
p._kill_and_close(sb)
|
|
|
|
|
|
def test_reuse_in_process_sandbox_returns_cached_id_on_healthy_reuse():
|
|
p = _make_provider()
|
|
client = FakeClient()
|
|
sb = _make_sandbox(client, sandbox_id="sb-1")
|
|
p._sandboxes["sb-1"] = sb
|
|
p._thread_sandboxes[("u1", "t1")] = "sb-1"
|
|
|
|
sid = p._reuse_in_process_sandbox("t1", user_id="u1")
|
|
assert sid == "sb-1"
|
|
assert client.timeouts_set, "expected set_timeout to be called on reuse"
|
|
|
|
|
|
def test_reuse_in_process_sandbox_evicts_dead_sandbox():
|
|
p = _make_provider()
|
|
client = FakeClient(commands=FakeCommandsAPI([FakeCommandsAPI.GONE]))
|
|
sb = _make_sandbox(client, sandbox_id="sb-dead")
|
|
sb._dead = True
|
|
p._sandboxes["sb-dead"] = sb
|
|
p._thread_sandboxes[("u1", "t1")] = "sb-dead"
|
|
|
|
sid = p._reuse_in_process_sandbox("t1", user_id="u1")
|
|
assert sid is None
|
|
assert "sb-dead" not in p._sandboxes
|
|
assert ("u1", "t1") not in p._thread_sandboxes
|
|
|
|
|
|
def test_reuse_in_process_sandbox_evicts_when_ping_fails():
|
|
p = _make_provider()
|
|
client = FakeClient(commands=FakeCommandsAPI([FakeCommandsAPI.GONE]))
|
|
sb = _make_sandbox(client, sandbox_id="sb-stale")
|
|
p._sandboxes["sb-stale"] = sb
|
|
p._thread_sandboxes[("u1", "t1")] = "sb-stale"
|
|
|
|
sid = p._reuse_in_process_sandbox("t1", user_id="u1")
|
|
assert sid is None
|
|
assert sb.is_dead is True
|
|
assert "sb-stale" not in p._sandboxes
|
|
|
|
|
|
def test_reuse_in_process_sandbox_cleans_dangling_mapping():
|
|
p = _make_provider()
|
|
p._thread_sandboxes[("u1", "t1")] = "ghost"
|
|
sid = p._reuse_in_process_sandbox("t1", user_id="u1")
|
|
assert sid is None
|
|
assert ("u1", "t1") not in p._thread_sandboxes
|
|
|
|
|
|
def test_reuse_in_process_sandbox_returns_none_when_no_mapping():
|
|
p = _make_provider()
|
|
assert p._reuse_in_process_sandbox("t-x", user_id="u-x") is None
|
|
|
|
|
|
def test_reclaim_warm_pool_sandbox_happy_path(monkeypatch):
|
|
p = _make_provider()
|
|
fake_cls = _install_fake_sdk(monkeypatch, p)
|
|
seed = p._stable_seed("t1", "u1")
|
|
p._warm_pool["sb-warm"] = (seed, 12345.0)
|
|
|
|
sid = p._reclaim_warm_pool_sandbox("t1", user_id="u1")
|
|
assert sid == "sb-warm"
|
|
assert "sb-warm" in p._sandboxes
|
|
assert p._thread_sandboxes[("u1", "t1")] == "sb-warm"
|
|
assert "sb-warm" not in p._warm_pool
|
|
assert [c[0] for c in fake_cls.connect_calls] == ["sb-warm"]
|
|
|
|
|
|
def test_reclaim_warm_pool_sandbox_drops_dead_entry(monkeypatch):
|
|
p = _make_provider()
|
|
fake_cls = _install_fake_sdk(monkeypatch, p)
|
|
fake_cls.connect_factory = lambda sid, **kw: FakeClient(sandbox_id=sid, commands=FakeCommandsAPI([FakeCommandsAPI.GONE]))
|
|
seed = p._stable_seed("t1", "u1")
|
|
p._warm_pool["sb-zombie"] = (seed, 12345.0)
|
|
|
|
sid = p._reclaim_warm_pool_sandbox("t1", user_id="u1")
|
|
assert sid is None
|
|
assert "sb-zombie" not in p._sandboxes
|
|
assert "sb-zombie" not in p._warm_pool
|
|
|
|
|
|
def test_reclaim_warm_pool_sandbox_handles_reconnect_exception(monkeypatch):
|
|
p = _make_provider()
|
|
fake_cls = _install_fake_sdk(monkeypatch, p)
|
|
|
|
def boom(sid, **kw):
|
|
raise RuntimeError("404 Not Found")
|
|
|
|
fake_cls.connect_factory = boom
|
|
seed = p._stable_seed("t1", "u1")
|
|
p._warm_pool["sb-broken"] = (seed, 12345.0)
|
|
|
|
sid = p._reclaim_warm_pool_sandbox("t1", user_id="u1")
|
|
assert sid is None
|
|
assert "sb-broken" not in p._warm_pool
|
|
|
|
|
|
def test_reclaim_warm_pool_sandbox_returns_none_on_seed_mismatch(monkeypatch):
|
|
p = _make_provider()
|
|
_install_fake_sdk(monkeypatch, p)
|
|
p._warm_pool["sb-other"] = ("some-other-seed", 12345.0)
|
|
assert p._reclaim_warm_pool_sandbox("t1", user_id="u1") is None
|
|
# The unrelated entry must remain untouched.
|
|
assert "sb-other" in p._warm_pool
|
|
|
|
|
|
class _FakePaginator:
|
|
"""Mirror of e2b SDK's ``SandboxPaginator``: items via ``next_items``."""
|
|
|
|
def __init__(self, pages: list[list[Any]]) -> None:
|
|
self._pages = list(pages)
|
|
self.has_next = bool(self._pages)
|
|
self.calls = 0
|
|
|
|
def next_items(self) -> list[Any]:
|
|
self.calls += 1
|
|
if not self._pages:
|
|
self.has_next = False
|
|
return []
|
|
page = self._pages.pop(0)
|
|
self.has_next = bool(self._pages)
|
|
return page
|
|
|
|
|
|
def _info(sandbox_id: str, user_id: str, thread_id: str):
|
|
return SimpleNamespace(
|
|
sandbox_id=sandbox_id,
|
|
metadata={
|
|
"deer_flow_provider": "e2b_sandbox_provider",
|
|
"deer_flow_user": user_id,
|
|
"deer_flow_thread": thread_id,
|
|
},
|
|
)
|
|
|
|
|
|
def test_discover_remote_sandbox_walks_paginator(monkeypatch):
|
|
p = _make_provider()
|
|
fake_cls = _install_fake_sdk(monkeypatch, p)
|
|
fake_cls.list_return = _FakePaginator(
|
|
[
|
|
[_info("sb-other", "u-x", "t-x")],
|
|
[_info("sb-match", "u1", "t1")],
|
|
]
|
|
)
|
|
|
|
sid = p._discover_remote_sandbox("t1", user_id="u1")
|
|
assert sid == "sb-match"
|
|
assert p._thread_sandboxes[("u1", "t1")] == "sb-match"
|
|
|
|
|
|
def test_discover_remote_sandbox_accepts_legacy_list(monkeypatch):
|
|
p = _make_provider()
|
|
fake_cls = _install_fake_sdk(monkeypatch, p)
|
|
fake_cls.list_return = [_info("sb-legacy", "u1", "t1")]
|
|
|
|
sid = p._discover_remote_sandbox("t1", user_id="u1")
|
|
assert sid == "sb-legacy"
|
|
|
|
|
|
def test_discover_remote_sandbox_skips_dead_candidate(monkeypatch):
|
|
p = _make_provider()
|
|
fake_cls = _install_fake_sdk(monkeypatch, p)
|
|
fake_cls.list_return = [_info("sb-dead", "u1", "t1")]
|
|
fake_cls.connect_factory = lambda sid, **kw: FakeClient(sandbox_id=sid, commands=FakeCommandsAPI([FakeCommandsAPI.GONE]))
|
|
|
|
assert p._discover_remote_sandbox("t1", user_id="u1") is None
|
|
assert ("u1", "t1") not in p._thread_sandboxes
|
|
|
|
|
|
def test_discover_remote_sandbox_returns_none_when_list_raises(monkeypatch):
|
|
p = _make_provider()
|
|
fake_cls = _install_fake_sdk(monkeypatch, p)
|
|
|
|
def boom(**kw):
|
|
raise RuntimeError("API unreachable")
|
|
|
|
fake_cls.list = boom
|
|
assert p._discover_remote_sandbox("t1", user_id="u1") is None
|
|
|
|
|
|
def test_bootstrap_sandbox_paths_emits_expected_script():
|
|
p = _make_provider()
|
|
client = FakeClient()
|
|
p._bootstrap_sandbox_paths(client)
|
|
assert len(client.commands.calls) == 1
|
|
script = client.commands.calls[0]
|
|
assert "ln -sfn" in script
|
|
assert "/mnt/user-data" in script
|
|
assert "/mnt/acp-workspace" in script
|
|
assert "BOOTSTRAP_OK" in script
|
|
for sub in ("workspace", "uploads", "outputs", "acp-workspace"):
|
|
assert f"/home/user/{sub}" in script
|
|
|
|
|
|
def test_bootstrap_sandbox_paths_swallows_command_failure():
|
|
p = _make_provider()
|
|
|
|
def boom(_cmd: str) -> Any:
|
|
raise RuntimeError("sudo not allowed")
|
|
|
|
client = FakeClient(commands=FakeCommandsAPI([boom]))
|
|
p._bootstrap_sandbox_paths(client)
|
|
|
|
|
|
def test_release_unknown_sandbox_id_is_noop():
|
|
p = _make_provider()
|
|
p.release("nonexistent")
|
|
assert p._warm_pool == OrderedDict()
|
|
|
|
|
|
def test_release_dead_sandbox_skips_warm_pool(monkeypatch):
|
|
p = _make_provider()
|
|
client = FakeClient()
|
|
sb = _make_sandbox(client, sandbox_id="sb-dead")
|
|
sb._dead = True
|
|
p._sandboxes["sb-dead"] = sb
|
|
p._thread_sandboxes[("u1", "t1")] = "sb-dead"
|
|
|
|
p.release("sb-dead")
|
|
|
|
assert "sb-dead" not in p._warm_pool, "dead sandbox must not be parked"
|
|
assert "sb-dead" not in p._sandboxes
|
|
assert ("u1", "t1") not in p._thread_sandboxes
|
|
assert client.killed is True, "release of dead sandbox must kill the remote VM"
|
|
|
|
|
|
def test_release_healthy_sandbox_parks_in_warm_pool(monkeypatch, tmp_path):
|
|
p = _make_provider()
|
|
_setup_paths(monkeypatch, tmp_path)
|
|
cmds = FakeCommandsAPI([SimpleNamespace(stdout="", stderr="", exit_code=0)])
|
|
client = FakeClient(commands=cmds)
|
|
sb = _make_sandbox(client, sandbox_id="sb-warm-1")
|
|
p._sandboxes["sb-warm-1"] = sb
|
|
p._thread_sandboxes[("u1", "t1")] = "sb-warm-1"
|
|
|
|
p.release("sb-warm-1")
|
|
|
|
assert "sb-warm-1" in p._warm_pool
|
|
seed_in_pool, _ts = p._warm_pool["sb-warm-1"]
|
|
assert seed_in_pool == p._stable_seed("t1", "u1")
|
|
assert client.killed is False
|
|
assert client.timeouts_set
|
|
|
|
|
|
def test_release_skips_warm_pool_when_sync_reveals_dead_vm(monkeypatch, tmp_path):
|
|
p = _make_provider()
|
|
_setup_paths(monkeypatch, tmp_path)
|
|
client = FakeClient(commands=FakeCommandsAPI([FakeCommandsAPI.GONE]))
|
|
sb = _make_sandbox(client, sandbox_id="sb-died-during-sync")
|
|
p._sandboxes["sb-died-during-sync"] = sb
|
|
p._thread_sandboxes[("u1", "t1")] = "sb-died-during-sync"
|
|
|
|
p.release("sb-died-during-sync")
|
|
|
|
assert sb.is_dead is True
|
|
assert "sb-died-during-sync" not in p._warm_pool
|
|
assert client.killed is True
|
|
|
|
|
|
def _setup_paths(monkeypatch, tmp_path):
|
|
paths_mod = importlib.import_module("deerflow.config.paths")
|
|
monkeypatch.setattr(paths_mod, "get_paths", lambda: Paths(base_dir=tmp_path), raising=False)
|
|
|
|
|
|
def test_sync_outputs_to_host_writes_new_files(monkeypatch, tmp_path):
|
|
p = _make_provider()
|
|
_setup_paths(monkeypatch, tmp_path)
|
|
listing = "13\t/home/user/outputs/random.pdf\x00"
|
|
files = FakeFilesAPI(store={"/home/user/outputs/random.pdf": b"%PDF-1.4hello"})
|
|
cmds = FakeCommandsAPI([SimpleNamespace(stdout=listing, stderr="", exit_code=0)])
|
|
client = FakeClient(commands=cmds, files=files)
|
|
sb = _make_sandbox(client, sandbox_id="sb-sync-1")
|
|
|
|
p._sync_outputs_to_host(sb, thread_id="t1", user_id="u1")
|
|
|
|
expected = Paths(base_dir=tmp_path).thread_dir("t1", user_id="u1") / "user-data" / "outputs" / "random.pdf"
|
|
assert expected.exists()
|
|
assert expected.read_bytes() == b"%PDF-1.4hello"
|
|
|
|
|
|
def test_sync_outputs_to_host_skips_unchanged_files(monkeypatch, tmp_path):
|
|
p = _make_provider()
|
|
_setup_paths(monkeypatch, tmp_path)
|
|
out_dir = Paths(base_dir=tmp_path).thread_dir("t1", user_id="u1") / "user-data" / "outputs"
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
target = out_dir / "random.pdf"
|
|
target.write_bytes(b"%PDF-1.4hello")
|
|
|
|
listing = "13\t/home/user/outputs/random.pdf\x00"
|
|
files = FakeFilesAPI(store={"/home/user/outputs/random.pdf": b"DIFFERENT-SAME-LEN"})
|
|
cmds = FakeCommandsAPI([SimpleNamespace(stdout=listing, stderr="", exit_code=0)])
|
|
client = FakeClient(commands=cmds, files=files)
|
|
sb = _make_sandbox(client, sandbox_id="sb-sync-2")
|
|
|
|
p._sync_outputs_to_host(sb, thread_id="t1", user_id="u1")
|
|
|
|
assert files.read_calls == [], "size match should skip the download round-trip"
|
|
assert target.read_bytes() == b"%PDF-1.4hello"
|
|
|
|
|
|
def test_sync_outputs_to_host_marks_dead_on_sandbox_gone(monkeypatch, tmp_path):
|
|
p = _make_provider()
|
|
_setup_paths(monkeypatch, tmp_path)
|
|
cmds = FakeCommandsAPI([FakeCommandsAPI.GONE])
|
|
client = FakeClient(commands=cmds)
|
|
sb = _make_sandbox(client, sandbox_id="sb-sync-dead")
|
|
|
|
p._sync_outputs_to_host(sb, thread_id="t1", user_id="u1")
|
|
|
|
assert sb.is_dead is True
|
|
|
|
|
|
def test_sync_outputs_to_host_uses_virtual_path_for_download(monkeypatch, tmp_path):
|
|
"""`download_file` requires paths under ``/mnt/user-data``; the sync
|
|
helper must translate the physical /home/user/... back to the virtual
|
|
prefix before calling it."""
|
|
p = _make_provider()
|
|
_setup_paths(monkeypatch, tmp_path)
|
|
|
|
listing = "5\t/home/user/outputs/sub/x.txt\x00"
|
|
files = FakeFilesAPI(store={"/home/user/outputs/sub/x.txt": b"hello"})
|
|
cmds = FakeCommandsAPI([SimpleNamespace(stdout=listing, stderr="", exit_code=0)])
|
|
client = FakeClient(commands=cmds, files=files)
|
|
sb = _make_sandbox(client, sandbox_id="sb-sync-3")
|
|
|
|
p._sync_outputs_to_host(sb, thread_id="t1", user_id="u1")
|
|
|
|
read_paths = [r[0] for r in files.read_calls]
|
|
assert "/home/user/outputs/sub/x.txt" in read_paths
|
|
|
|
|
|
def test_sync_outputs_to_host_is_noop_when_client_closed():
|
|
p = _make_provider()
|
|
sb = _make_sandbox(FakeClient(), sandbox_id="sb-x")
|
|
sb.close()
|
|
p._sync_outputs_to_host(sb, thread_id="t1", user_id="u1")
|
|
|
|
|
|
def test_download_file_uses_streaming_read_and_returns_full_bytes():
|
|
payload = b"A" * (128 * 1024) # 128 KiB — well below the cap.
|
|
files = FakeFilesAPI(store={"/home/user/outputs/small.bin": payload})
|
|
client = FakeClient(files=files)
|
|
sb = _make_sandbox(client, sandbox_id="sb-stream-1")
|
|
|
|
data = sb.download_file("/mnt/user-data/outputs/small.bin")
|
|
|
|
assert data == payload
|
|
formats_used = [fmt for _p, fmt in files.read_calls]
|
|
assert "stream" in formats_used, f"expected download_file to invoke read(format='stream'), got {formats_used!r}"
|
|
assert files.streams, "download_file must actually consume a stream"
|
|
assert files.streams[-1].closed, "stream must be closed after successful read"
|
|
|
|
|
|
def test_download_file_streaming_raises_efbig_before_full_buffering():
|
|
import errno as _errno
|
|
|
|
from deerflow.community.e2b_sandbox import e2b_sandbox as e2b_sb_mod
|
|
|
|
cap = e2b_sb_mod._MAX_DOWNLOAD_SIZE
|
|
|
|
class _OversizeStream:
|
|
def __init__(self) -> None:
|
|
self.bytes_yielded = 0
|
|
self.closed = False
|
|
self._chunk = b"X" * (1024 * 1024) # 1 MiB per chunk
|
|
|
|
def __iter__(self):
|
|
return self
|
|
|
|
def __next__(self) -> bytes:
|
|
if self.closed:
|
|
raise StopIteration
|
|
# Yield up to ``cap + a bit`` — the caller must abort before
|
|
# actually buffering all of that in memory.
|
|
if self.bytes_yielded > cap + 4 * len(self._chunk):
|
|
raise StopIteration
|
|
self.bytes_yielded += len(self._chunk)
|
|
return self._chunk
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
stream = _OversizeStream()
|
|
|
|
class _StubFilesAPI:
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, str | None]] = []
|
|
|
|
def read(self, path: str, *, format: str | None = None):
|
|
self.calls.append((path, format))
|
|
assert format == "stream", "provider must request a streamed download"
|
|
return stream
|
|
|
|
files = _StubFilesAPI()
|
|
client = FakeClient(files=files) # type: ignore[arg-type]
|
|
sb = _make_sandbox(client, sandbox_id="sb-stream-oversize")
|
|
|
|
try:
|
|
sb.download_file("/mnt/user-data/outputs/huge.bin")
|
|
except OSError as exc:
|
|
assert exc.errno == _errno.EFBIG, f"expected EFBIG, got errno={exc.errno!r} ({exc})"
|
|
else: # pragma: no cover - defensive
|
|
raise AssertionError("download_file must raise OSError(EFBIG) on oversize stream")
|
|
|
|
assert stream.closed is True, "stream must be closed on abort so the pooled connection is released"
|
|
assert stream.bytes_yielded <= cap + 2 * 1024 * 1024, f"aborted too late: yielded={stream.bytes_yielded} vs cap={cap}"
|
|
|
|
|
|
def test_download_file_falls_back_to_buffered_read_for_legacy_sdk():
|
|
|
|
class _LegacyFilesAPI:
|
|
def __init__(self, data: bytes) -> None:
|
|
self._data = data
|
|
self.calls: list[tuple[str, str | None]] = []
|
|
|
|
def read(self, path: str, *, format: str | None = None):
|
|
self.calls.append((path, format))
|
|
if format == "stream":
|
|
raise TypeError("format='stream' unsupported")
|
|
if format == "bytes":
|
|
return self._data
|
|
return self._data.decode("utf-8", errors="replace")
|
|
|
|
files = _LegacyFilesAPI(b"legacy-payload")
|
|
client = FakeClient(files=files) # type: ignore[arg-type]
|
|
sb = _make_sandbox(client, sandbox_id="sb-legacy")
|
|
|
|
data = sb.download_file("/mnt/user-data/outputs/legacy.bin")
|
|
assert data == b"legacy-payload"
|
|
formats_used = [fmt for _p, fmt in files.calls]
|
|
assert formats_used == ["stream", "bytes"], f"expected stream then bytes fallback, got {formats_used!r}"
|
|
|
|
|
|
def test_sync_outputs_to_host_skips_oversize_files(monkeypatch, tmp_path):
|
|
from deerflow.community.e2b_sandbox import e2b_sandbox as e2b_sb_mod
|
|
|
|
p = _make_provider()
|
|
_setup_paths(monkeypatch, tmp_path)
|
|
|
|
oversize = e2b_sb_mod._MAX_DOWNLOAD_SIZE + 1
|
|
listing = f"{oversize}\t/home/user/outputs/huge.bin\x00"
|
|
files = FakeFilesAPI() # no store entry: any read attempt would raise
|
|
cmds = FakeCommandsAPI([SimpleNamespace(stdout=listing, stderr="", exit_code=0)])
|
|
client = FakeClient(commands=cmds, files=files)
|
|
sb = _make_sandbox(client, sandbox_id="sb-sync-oversize")
|
|
|
|
p._sync_outputs_to_host(sb, thread_id="t1", user_id="u1")
|
|
|
|
assert files.read_calls == [], "oversize files must be skipped without invoking download_file"
|
|
host_target = Paths(base_dir=tmp_path).thread_dir("t1", user_id="u1") / "user-data" / "outputs" / "huge.bin"
|
|
assert not host_target.exists(), "no oversize artefact must be written to host"
|