mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
318 lines
9.9 KiB
Python
318 lines
9.9 KiB
Python
"""
|
|
Tests for the e2b code execution sandbox primitive.
|
|
|
|
Unit tests inject a fake async HTTP client (dependency injection, no
|
|
monkeypatching) and assert request shapes and result mapping. Real-network
|
|
integration tests live in tests/integration/sandbox/test_e2b_sandbox.py.
|
|
"""
|
|
|
|
import json
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.llms.base_llm.sandbox.transformation import ContainerHandle
|
|
from litellm.llms.e2b.sandbox.transformation import (
|
|
MAX_OUTPUT_BYTES,
|
|
E2BSandboxConfig,
|
|
)
|
|
|
|
|
|
class FakeResponse:
|
|
def __init__(self, *, json_data=None, lines=None, status_code=200):
|
|
self._json = json_data
|
|
self._lines = lines or []
|
|
self.status_code = status_code
|
|
|
|
def json(self):
|
|
return self._json
|
|
|
|
async def aiter_lines(self):
|
|
for line in self._lines:
|
|
yield line
|
|
|
|
|
|
class FakeHTTPClient:
|
|
"""Records outbound requests and returns canned responses keyed by URL."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
create_json=None,
|
|
execute_lines=None,
|
|
delete_status=204,
|
|
execute_raises=None,
|
|
):
|
|
self.create_json = create_json or {
|
|
"sandboxID": "sbx_123",
|
|
"domain": "e2b.app",
|
|
"envdAccessToken": "tok_abc",
|
|
}
|
|
self.execute_lines = execute_lines or []
|
|
self.delete_status = delete_status
|
|
self.execute_raises = execute_raises
|
|
self.calls = []
|
|
|
|
async def post(self, url, headers=None, json=None, stream=False, **kwargs):
|
|
self.calls.append(("POST", url, headers, json))
|
|
if url.endswith("/sandboxes"):
|
|
return FakeResponse(json_data=self.create_json)
|
|
if url.endswith("/execute"):
|
|
if self.execute_raises is not None:
|
|
raise self.execute_raises
|
|
return FakeResponse(lines=self.execute_lines)
|
|
raise AssertionError(f"unexpected POST {url}")
|
|
|
|
async def delete(self, url, headers=None, **kwargs):
|
|
self.calls.append(("DELETE", url, headers, None))
|
|
if not (200 <= self.delete_status < 300):
|
|
raise httpx.HTTPStatusError(
|
|
f"status {self.delete_status}",
|
|
request=httpx.Request("DELETE", url),
|
|
response=httpx.Response(self.delete_status),
|
|
)
|
|
return FakeResponse(status_code=self.delete_status)
|
|
|
|
|
|
# ---------- pure parser ----------
|
|
|
|
|
|
def test_parse_lines_stdout_and_count():
|
|
lines = [
|
|
json.dumps({"type": "stdout", "text": "6\n", "timestamp": 1}),
|
|
json.dumps({"type": "number_of_executions", "execution_count": 1}),
|
|
]
|
|
result = E2BSandboxConfig._parse_lines(lines)
|
|
assert result.stdout == "6\n"
|
|
assert result.execution_count == 1
|
|
assert result.error is None
|
|
|
|
|
|
def test_parse_lines_error_surfaces_name_and_traceback():
|
|
lines = [
|
|
json.dumps(
|
|
{
|
|
"type": "error",
|
|
"name": "ZeroDivisionError",
|
|
"value": "division by zero",
|
|
"traceback": "Traceback (most recent call last): ...",
|
|
}
|
|
)
|
|
]
|
|
result = E2BSandboxConfig._parse_lines(lines)
|
|
assert result.error["name"] == "ZeroDivisionError"
|
|
assert "Traceback" in result.error["traceback"]
|
|
|
|
|
|
def test_parse_lines_result_carries_png():
|
|
lines = [
|
|
json.dumps({"type": "result", "png": "BASE64DATA", "is_main_result": True})
|
|
]
|
|
result = E2BSandboxConfig._parse_lines(lines)
|
|
assert result.results and result.results[0]["png"] == "BASE64DATA"
|
|
assert "type" not in result.results[0]
|
|
|
|
|
|
# ---------- request shapes ----------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_template_flows_into_create_request_as_templateID():
|
|
client = FakeHTTPClient()
|
|
cfg = E2BSandboxConfig()
|
|
handle = await cfg.acreate_sandbox(
|
|
template="my-custom-template", api_key="e2b_key", client=client
|
|
)
|
|
|
|
method, url, headers, body = client.calls[0]
|
|
assert method == "POST"
|
|
assert url.endswith("/sandboxes")
|
|
assert body["templateID"] == "my-custom-template" # not "template"
|
|
assert body["secure"] is True
|
|
assert headers["X-API-Key"] == "e2b_key"
|
|
assert handle.id == "sbx_123"
|
|
assert handle._hidden_params["envd_access_token"] == "tok_abc"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_defaults_template_when_omitted():
|
|
client = FakeHTTPClient()
|
|
await E2BSandboxConfig().acreate_sandbox(api_key="e2b_key", client=client)
|
|
_, _, _, body = client.calls[0]
|
|
assert body["templateID"] == "code-interpreter-v1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_code_targets_jupyter_host_with_access_token():
|
|
client = FakeHTTPClient(
|
|
execute_lines=[json.dumps({"type": "stdout", "text": "42\n", "timestamp": 1})]
|
|
)
|
|
handle = ContainerHandle(id="sbx_xyz", provider="e2b", domain="e2b.app")
|
|
handle._hidden_params = {"envd_access_token": "tok_run"}
|
|
|
|
result = await E2BSandboxConfig().arun_code(
|
|
container=handle, code="print(6*7)", client=client
|
|
)
|
|
|
|
method, url, headers, body = client.calls[0]
|
|
assert url == "https://49999-sbx_xyz.e2b.app/execute"
|
|
assert headers["X-Access-Token"] == "tok_run"
|
|
assert body["code"] == "print(6*7)"
|
|
assert result.stdout.strip() == "42"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_issues_delete_to_sandbox_id():
|
|
client = FakeHTTPClient(delete_status=204)
|
|
handle = ContainerHandle(id="sbx_del", provider="e2b", domain="e2b.app")
|
|
handle._hidden_params = {"api_key": "e2b_key"}
|
|
|
|
ok = await E2BSandboxConfig().adelete_sandbox(container=handle, client=client)
|
|
|
|
method, url, headers, _ = client.calls[0]
|
|
assert method == "DELETE"
|
|
assert url.endswith("/sandboxes/sbx_del")
|
|
assert ok is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_returns_false_on_404():
|
|
client = FakeHTTPClient(delete_status=404)
|
|
handle = ContainerHandle(id="sbx_gone", provider="e2b", domain="e2b.app")
|
|
handle._hidden_params = {"api_key": "e2b_key"}
|
|
ok = await E2BSandboxConfig().adelete_sandbox(container=handle, client=client)
|
|
assert ok is False
|
|
|
|
|
|
# ---------- ephemeral teardown ----------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_code_interpreter_tool_deletes_even_when_run_raises():
|
|
client = FakeHTTPClient(execute_raises=RuntimeError("boom"))
|
|
|
|
with pytest.raises(RuntimeError, match="boom"):
|
|
await litellm.acode_interpreter_tool(
|
|
provider="e2b", code="1/0", api_key="e2b_key", client=client
|
|
)
|
|
|
|
methods = [c[0] for c in client.calls]
|
|
urls = [c[1] for c in client.calls]
|
|
assert methods == ["POST", "POST", "DELETE"] # create, run(raises), delete
|
|
assert urls[0].endswith("/sandboxes")
|
|
assert urls[1].endswith("/execute")
|
|
assert urls[2].endswith("/sandboxes/sbx_123")
|
|
|
|
|
|
# ---------- correctness guards ----------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_reraises_non_404_http_error():
|
|
client = FakeHTTPClient(delete_status=500)
|
|
handle = ContainerHandle(id="sbx_err", provider="e2b", domain="e2b.app")
|
|
handle._hidden_params = {"api_key": "e2b_key"}
|
|
with pytest.raises(httpx.HTTPStatusError):
|
|
await E2BSandboxConfig().adelete_sandbox(container=handle, client=client)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_preserves_explicit_zero_timeout():
|
|
client = FakeHTTPClient()
|
|
await E2BSandboxConfig().acreate_sandbox(
|
|
timeout=0, api_key="e2b_key", client=client
|
|
)
|
|
_, _, _, body = client.calls[0]
|
|
assert body["timeout"] == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_code_rejects_bare_id_without_access_token():
|
|
client = FakeHTTPClient()
|
|
with pytest.raises(ValueError, match="access token"):
|
|
await E2BSandboxConfig().arun_code(
|
|
container="sbx_no_token", code="print(1)", client=client
|
|
)
|
|
assert client.calls == [] # never reached the network
|
|
|
|
|
|
def test_parse_lines_skips_non_json_lines():
|
|
lines = [
|
|
"not-json-heartbeat",
|
|
json.dumps({"type": "stdout", "text": "ok\n"}),
|
|
"",
|
|
"{partial",
|
|
]
|
|
result = E2BSandboxConfig._parse_lines(lines)
|
|
assert result.stdout == "ok\n"
|
|
assert result.error is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_run_code_aborts_on_output_over_cap():
|
|
big_line = "x" * (MAX_OUTPUT_BYTES + 1)
|
|
client = FakeHTTPClient(execute_lines=[big_line])
|
|
handle = ContainerHandle(id="sbx_big", provider="e2b", domain="e2b.app")
|
|
handle._hidden_params = {"envd_access_token": "tok"}
|
|
with pytest.raises(ValueError, match="exceeded"):
|
|
await E2BSandboxConfig().arun_code(
|
|
container=handle, code="print('x'*999)", client=client
|
|
)
|
|
|
|
|
|
# ---------- public entrypoints ----------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_public_lifecycle_create_run_delete():
|
|
client = FakeHTTPClient(
|
|
execute_lines=[json.dumps({"type": "stdout", "text": "42\n"})]
|
|
)
|
|
container = await litellm.acreate_sandbox(
|
|
provider="e2b", api_key="e2b_key", client=client
|
|
)
|
|
assert container.id == "sbx_123"
|
|
|
|
result = await litellm.arun_code(
|
|
provider="e2b",
|
|
container=container,
|
|
api_key="e2b_key",
|
|
code="print(6*7)",
|
|
client=client,
|
|
)
|
|
assert result.stdout.strip() == "42"
|
|
|
|
assert (
|
|
await litellm.adelete_sandbox(
|
|
provider="e2b", container=container, api_key="e2b_key", client=client
|
|
)
|
|
is True
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_unsupported_provider_raises():
|
|
with pytest.raises(ValueError, match="not-a-provider' is not a valid SandboxProviders"):
|
|
await litellm.acreate_sandbox(provider="not-a-provider")
|
|
|
|
|
|
# ---------- api_base override ----------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_uses_api_base_override():
|
|
client = FakeHTTPClient()
|
|
await E2BSandboxConfig().acreate_sandbox(
|
|
api_base="http://my-sandbox:8080", api_key="k", client=client
|
|
)
|
|
_, url, _, _ = client.calls[0]
|
|
assert url == "http://my-sandbox:8080/sandboxes"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_defaults_to_e2b_api_base():
|
|
client = FakeHTTPClient()
|
|
await E2BSandboxConfig().acreate_sandbox(api_key="k", client=client)
|
|
_, url, _, _ = client.calls[0]
|
|
assert url == "https://api.e2b.app/sandboxes"
|