forked from wrenn/python-sdk
feat: implement client architecture and sandbox environment
Introduces the core Wrenn client and a dedicated sandbox execution environment. This includes automated model generation and a custom exception hierarchy to support robust integration. - Add `WrennClient` in `src/wrenn/client.py` for API interaction. - Implement `Sandbox` in `src/wrenn/sandbox.py` for isolated execution. - Add Pydantic/model support via `_generated.py`. - Define project-specific error types in `exceptions.py`. - Include AGENTS.md documentation for specialized logic. - Add comprehensive unit and integration tests. - Update build system (Makefile, uv.lock, pyproject.toml) and LICENSE.
This commit is contained in:
Tasnim Kabir Sadik
417
tests/test_client.py
Normal file
417
tests/test_client.py
Normal file
@ -0,0 +1,417 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
from wrenn.client import AsyncWrennClient, WrennClient
|
||||
from wrenn.exceptions import (
|
||||
WrennAgentError,
|
||||
WrennAuthenticationError,
|
||||
WrennConflictError,
|
||||
WrennForbiddenError,
|
||||
WrennHostHasSandboxesError,
|
||||
WrennInternalError,
|
||||
WrennNotFoundError,
|
||||
WrennValidationError,
|
||||
)
|
||||
from wrenn.models import (
|
||||
APIKeyResponse,
|
||||
AuthResponse,
|
||||
CreateHostResponse,
|
||||
Host,
|
||||
Sandbox,
|
||||
Status,
|
||||
Template,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client():
|
||||
with WrennClient(api_key="wrn_test1234567890abcdef12345678") as c:
|
||||
yield c
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def async_client():
|
||||
return AsyncWrennClient(api_key="wrn_test1234567890abcdef12345678")
|
||||
|
||||
|
||||
class TestAuth:
|
||||
@respx.mock
|
||||
def test_signup(self, client):
|
||||
respx.post("https://api.wrenn.dev/v1/auth/signup").respond(
|
||||
201,
|
||||
json={
|
||||
"token": "jwt-token",
|
||||
"user_id": "u-1",
|
||||
"team_id": "t-1",
|
||||
"email": "a@b.com",
|
||||
},
|
||||
)
|
||||
resp = client.auth.signup("a@b.com", "password123")
|
||||
assert isinstance(resp, AuthResponse)
|
||||
assert resp.token == "jwt-token"
|
||||
assert resp.user_id == "u-1"
|
||||
|
||||
@respx.mock
|
||||
def test_login(self, client):
|
||||
respx.post("https://api.wrenn.dev/v1/auth/login").respond(
|
||||
200,
|
||||
json={"token": "jwt-token", "email": "a@b.com"},
|
||||
)
|
||||
resp = client.auth.login("a@b.com", "password123")
|
||||
assert resp.token == "jwt-token"
|
||||
|
||||
|
||||
class TestAPIKeys:
|
||||
@respx.mock
|
||||
def test_create(self, client):
|
||||
respx.post("https://api.wrenn.dev/v1/api-keys").respond(
|
||||
201,
|
||||
json={
|
||||
"id": "key-1",
|
||||
"name": "my-key",
|
||||
"key_prefix": "wrn_ab12cd34",
|
||||
"key": "wrn_ab12cd34fullkey",
|
||||
},
|
||||
)
|
||||
resp = client.api_keys.create(name="my-key")
|
||||
assert isinstance(resp, APIKeyResponse)
|
||||
assert resp.name == "my-key"
|
||||
assert resp.key == "wrn_ab12cd34fullkey"
|
||||
|
||||
@respx.mock
|
||||
def test_list(self, client):
|
||||
respx.get("https://api.wrenn.dev/v1/api-keys").respond(
|
||||
200,
|
||||
json=[{"id": "key-1", "name": "k1"}, {"id": "key-2", "name": "k2"}],
|
||||
)
|
||||
keys = client.api_keys.list()
|
||||
assert len(keys) == 2
|
||||
assert keys[0].id == "key-1"
|
||||
|
||||
@respx.mock
|
||||
def test_delete(self, client):
|
||||
route = respx.delete("https://api.wrenn.dev/v1/api-keys/key-1").respond(204)
|
||||
client.api_keys.delete("key-1")
|
||||
assert route.called
|
||||
|
||||
|
||||
class TestSandboxes:
|
||||
@respx.mock
|
||||
def test_create(self, client):
|
||||
respx.post("https://api.wrenn.dev/v1/sandboxes").respond(
|
||||
201,
|
||||
json={
|
||||
"id": "sb-1",
|
||||
"status": "pending",
|
||||
"template": "base-python",
|
||||
"vcpus": 2,
|
||||
"memory_mb": 1024,
|
||||
},
|
||||
)
|
||||
resp = client.sandboxes.create(template="base-python", vcpus=2, memory_mb=1024)
|
||||
assert isinstance(resp, Sandbox)
|
||||
assert resp.id == "sb-1"
|
||||
assert resp.status == Status.pending
|
||||
|
||||
@respx.mock
|
||||
def test_create_defaults(self, client):
|
||||
respx.post("https://api.wrenn.dev/v1/sandboxes").respond(
|
||||
201, json={"id": "sb-2", "status": "pending"}
|
||||
)
|
||||
resp = client.sandboxes.create()
|
||||
assert resp.id == "sb-2"
|
||||
|
||||
@respx.mock
|
||||
def test_list(self, client):
|
||||
respx.get("https://api.wrenn.dev/v1/sandboxes").respond(
|
||||
200, json=[{"id": "sb-1", "status": "running"}]
|
||||
)
|
||||
boxes = client.sandboxes.list()
|
||||
assert len(boxes) == 1
|
||||
assert boxes[0].status == Status.running
|
||||
|
||||
@respx.mock
|
||||
def test_get(self, client):
|
||||
respx.get("https://api.wrenn.dev/v1/sandboxes/sb-1").respond(
|
||||
200, json={"id": "sb-1", "status": "running"}
|
||||
)
|
||||
resp = client.sandboxes.get("sb-1")
|
||||
assert resp.id == "sb-1"
|
||||
|
||||
@respx.mock
|
||||
def test_destroy(self, client):
|
||||
route = respx.delete("https://api.wrenn.dev/v1/sandboxes/sb-1").respond(204)
|
||||
client.sandboxes.destroy("sb-1")
|
||||
assert route.called
|
||||
|
||||
|
||||
class TestSnapshots:
|
||||
@respx.mock
|
||||
def test_create(self, client):
|
||||
respx.post("https://api.wrenn.dev/v1/snapshots").respond(
|
||||
201,
|
||||
json={"name": "snap-1", "type": "snapshot", "vcpus": 1},
|
||||
)
|
||||
resp = client.snapshots.create(sandbox_id="sb-1", name="snap-1")
|
||||
assert isinstance(resp, Template)
|
||||
assert resp.name == "snap-1"
|
||||
|
||||
@respx.mock
|
||||
def test_create_with_overwrite(self, client):
|
||||
route = respx.post("https://api.wrenn.dev/v1/snapshots").respond(
|
||||
201, json={"name": "snap-1", "type": "snapshot"}
|
||||
)
|
||||
client.snapshots.create(sandbox_id="sb-1", overwrite=True)
|
||||
req = route.calls[0].request
|
||||
assert "overwrite=true" in str(req.url)
|
||||
|
||||
@respx.mock
|
||||
def test_list(self, client):
|
||||
respx.get("https://api.wrenn.dev/v1/snapshots").respond(
|
||||
200, json=[{"name": "base-python", "type": "base"}]
|
||||
)
|
||||
snaps = client.snapshots.list()
|
||||
assert len(snaps) == 1
|
||||
|
||||
@respx.mock
|
||||
def test_list_with_filter(self, client):
|
||||
route = respx.get("https://api.wrenn.dev/v1/snapshots").respond(200, json=[])
|
||||
client.snapshots.list(type="snapshot")
|
||||
req = route.calls[0].request
|
||||
assert "type=snapshot" in str(req.url)
|
||||
|
||||
@respx.mock
|
||||
def test_delete(self, client):
|
||||
route = respx.delete("https://api.wrenn.dev/v1/snapshots/snap-1").respond(204)
|
||||
client.snapshots.delete("snap-1")
|
||||
assert route.called
|
||||
|
||||
|
||||
class TestHosts:
|
||||
@respx.mock
|
||||
def test_create(self, client):
|
||||
respx.post("https://api.wrenn.dev/v1/hosts").respond(
|
||||
201,
|
||||
json={
|
||||
"host": {"id": "h-1", "type": "regular", "status": "pending"},
|
||||
"registration_token": "reg-tok-123",
|
||||
},
|
||||
)
|
||||
resp = client.hosts.create(type="regular")
|
||||
assert isinstance(resp, CreateHostResponse)
|
||||
assert resp.registration_token == "reg-tok-123"
|
||||
|
||||
@respx.mock
|
||||
def test_list(self, client):
|
||||
respx.get("https://api.wrenn.dev/v1/hosts").respond(
|
||||
200, json=[{"id": "h-1", "status": "online"}]
|
||||
)
|
||||
hosts = client.hosts.list()
|
||||
assert len(hosts) == 1
|
||||
assert isinstance(hosts[0], Host)
|
||||
|
||||
@respx.mock
|
||||
def test_get(self, client):
|
||||
respx.get("https://api.wrenn.dev/v1/hosts/h-1").respond(
|
||||
200, json={"id": "h-1", "status": "online"}
|
||||
)
|
||||
resp = client.hosts.get("h-1")
|
||||
assert resp.id == "h-1"
|
||||
|
||||
@respx.mock
|
||||
def test_delete(self, client):
|
||||
route = respx.delete("https://api.wrenn.dev/v1/hosts/h-1").respond(204)
|
||||
client.hosts.delete("h-1")
|
||||
assert route.called
|
||||
|
||||
@respx.mock
|
||||
def test_regenerate_token(self, client):
|
||||
respx.post("https://api.wrenn.dev/v1/hosts/h-1/token").respond(
|
||||
201,
|
||||
json={
|
||||
"host": {"id": "h-1"},
|
||||
"registration_token": "new-tok",
|
||||
},
|
||||
)
|
||||
resp = client.hosts.regenerate_token("h-1")
|
||||
assert resp.registration_token == "new-tok"
|
||||
|
||||
@respx.mock
|
||||
def test_list_tags(self, client):
|
||||
respx.get("https://api.wrenn.dev/v1/hosts/h-1/tags").respond(
|
||||
200, json=["gpu", "high-mem"]
|
||||
)
|
||||
tags = client.hosts.list_tags("h-1")
|
||||
assert tags == ["gpu", "high-mem"]
|
||||
|
||||
@respx.mock
|
||||
def test_add_tag(self, client):
|
||||
route = respx.post("https://api.wrenn.dev/v1/hosts/h-1/tags").respond(204)
|
||||
client.hosts.add_tag("h-1", "gpu")
|
||||
assert route.called
|
||||
|
||||
@respx.mock
|
||||
def test_remove_tag(self, client):
|
||||
route = respx.delete("https://api.wrenn.dev/v1/hosts/h-1/tags/gpu").respond(204)
|
||||
client.hosts.remove_tag("h-1", "gpu")
|
||||
assert route.called
|
||||
|
||||
|
||||
class TestErrorHandling:
|
||||
@respx.mock
|
||||
def test_validation_error(self, client):
|
||||
respx.post("https://api.wrenn.dev/v1/sandboxes").respond(
|
||||
400,
|
||||
json={"error": {"code": "invalid_request", "message": "bad input"}},
|
||||
)
|
||||
with pytest.raises(WrennValidationError) as exc_info:
|
||||
client.sandboxes.create()
|
||||
assert exc_info.value.code == "invalid_request"
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
@respx.mock
|
||||
def test_auth_error(self, client):
|
||||
respx.get("https://api.wrenn.dev/v1/sandboxes").respond(
|
||||
401,
|
||||
json={"error": {"code": "unauthorized", "message": "bad key"}},
|
||||
)
|
||||
with pytest.raises(WrennAuthenticationError):
|
||||
client.sandboxes.list()
|
||||
|
||||
@respx.mock
|
||||
def test_forbidden_error(self, client):
|
||||
respx.post("https://api.wrenn.dev/v1/hosts").respond(
|
||||
403,
|
||||
json={"error": {"code": "forbidden", "message": "nope"}},
|
||||
)
|
||||
with pytest.raises(WrennForbiddenError):
|
||||
client.hosts.create(type="regular")
|
||||
|
||||
@respx.mock
|
||||
def test_not_found_error(self, client):
|
||||
respx.get("https://api.wrenn.dev/v1/sandboxes/nope").respond(
|
||||
404,
|
||||
json={"error": {"code": "not_found", "message": "sandbox not found"}},
|
||||
)
|
||||
with pytest.raises(WrennNotFoundError):
|
||||
client.sandboxes.get("nope")
|
||||
|
||||
@respx.mock
|
||||
def test_conflict_error(self, client):
|
||||
respx.get("https://api.wrenn.dev/v1/sandboxes/sb-1").respond(
|
||||
409,
|
||||
json={"error": {"code": "invalid_state", "message": "not running"}},
|
||||
)
|
||||
with pytest.raises(WrennConflictError):
|
||||
client.sandboxes.get("sb-1")
|
||||
|
||||
@respx.mock
|
||||
def test_host_has_sandboxes_error(self, client):
|
||||
respx.delete("https://api.wrenn.dev/v1/hosts/h-1").respond(
|
||||
409,
|
||||
json={
|
||||
"error": {
|
||||
"code": "host_has_sandboxes",
|
||||
"message": "host has running sandboxes",
|
||||
},
|
||||
"sandbox_ids": ["sb-1", "sb-2"],
|
||||
},
|
||||
)
|
||||
with pytest.raises(WrennHostHasSandboxesError) as exc_info:
|
||||
client.hosts.delete("h-1")
|
||||
assert exc_info.value.sandbox_ids == ["sb-1", "sb-2"]
|
||||
|
||||
@respx.mock
|
||||
def test_agent_error(self, client):
|
||||
respx.post("https://api.wrenn.dev/v1/sandboxes").respond(
|
||||
502,
|
||||
json={"error": {"code": "agent_error", "message": "host agent failed"}},
|
||||
)
|
||||
with pytest.raises(WrennAgentError):
|
||||
client.sandboxes.create()
|
||||
|
||||
@respx.mock
|
||||
def test_internal_error(self, client):
|
||||
respx.get("https://api.wrenn.dev/v1/sandboxes/sb-1").respond(
|
||||
500,
|
||||
json={"error": {"code": "internal_error", "message": "oops"}},
|
||||
)
|
||||
with pytest.raises(WrennInternalError):
|
||||
client.sandboxes.get("sb-1")
|
||||
|
||||
@respx.mock
|
||||
def test_unknown_error_code_falls_back(self, client):
|
||||
respx.get("https://api.wrenn.dev/v1/sandboxes/sb-1").respond(
|
||||
418,
|
||||
json={"error": {"code": "teapot", "message": "I'm a teapot"}},
|
||||
)
|
||||
from wrenn.exceptions import WrennError
|
||||
|
||||
with pytest.raises(WrennError) as exc_info:
|
||||
client.sandboxes.get("sb-1")
|
||||
assert exc_info.value.code == "teapot"
|
||||
|
||||
|
||||
class TestAuthModes:
|
||||
def test_api_key_header(self):
|
||||
with WrennClient(api_key="wrn_test1234567890abcdef12345678") as c:
|
||||
assert c._http.headers["X-API-Key"] == "wrn_test1234567890abcdef12345678"
|
||||
|
||||
def test_token_header(self):
|
||||
with WrennClient(token="jwt-token-abc") as c:
|
||||
assert c._http.headers["Authorization"] == "Bearer jwt-token-abc"
|
||||
|
||||
def test_no_auth_raises(self):
|
||||
with pytest.raises(ValueError, match="Either api_key or token"):
|
||||
WrennClient()
|
||||
|
||||
@respx.mock
|
||||
def test_jwt_auth_on_api_keys(self):
|
||||
route = respx.get("https://api.wrenn.dev/v1/api-keys").respond(200, json=[])
|
||||
with WrennClient(token="jwt-abc") as c:
|
||||
c.api_keys.list()
|
||||
req = route.calls[0].request
|
||||
assert req.headers["Authorization"] == "Bearer jwt-abc"
|
||||
|
||||
|
||||
class TestAsyncClient:
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_async_sandboxes_create(self, async_client):
|
||||
async with async_client:
|
||||
respx.post("https://api.wrenn.dev/v1/sandboxes").respond(
|
||||
201, json={"id": "sb-1", "status": "pending"}
|
||||
)
|
||||
resp = await async_client.sandboxes.create(template="base-python")
|
||||
assert resp.id == "sb-1"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_async_sandboxes_list(self, async_client):
|
||||
async with async_client:
|
||||
respx.get("https://api.wrenn.dev/v1/sandboxes").respond(
|
||||
200, json=[{"id": "sb-1"}]
|
||||
)
|
||||
boxes = await async_client.sandboxes.list()
|
||||
assert len(boxes) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_async_hosts_list(self, async_client):
|
||||
async with async_client:
|
||||
respx.get("https://api.wrenn.dev/v1/hosts").respond(200, json=[])
|
||||
hosts = await async_client.hosts.list()
|
||||
assert hosts == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@respx.mock
|
||||
async def test_async_error_handling(self, async_client):
|
||||
async with async_client:
|
||||
respx.get("https://api.wrenn.dev/v1/sandboxes/nope").respond(
|
||||
404,
|
||||
json={"error": {"code": "not_found", "message": "not found"}},
|
||||
)
|
||||
with pytest.raises(WrennNotFoundError):
|
||||
await async_client.sandboxes.get("nope")
|
||||
289
tests/test_integration.py
Normal file
289
tests/test_integration.py
Normal file
@ -0,0 +1,289 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Generator
|
||||
|
||||
import pytest
|
||||
|
||||
from wrenn.client import AsyncWrennClient, WrennClient
|
||||
from wrenn.exceptions import WrennNotFoundError, WrennValidationError
|
||||
|
||||
WRENN_API_KEY = os.environ.get("WRENN_API_KEY")
|
||||
WRENN_TOKEN = os.environ.get("WRENN_TOKEN")
|
||||
WRENN_BASE_URL = os.environ.get("WRENN_BASE_URL", "http://localhost:8080")
|
||||
WRENN_TEST_EMAIL = os.environ.get("WRENN_TEST_EMAIL")
|
||||
WRENN_TEST_PASSWORD = os.environ.get("WRENN_TEST_PASSWORD")
|
||||
|
||||
|
||||
def _has_auth() -> bool:
|
||||
return bool(WRENN_API_KEY or WRENN_TOKEN)
|
||||
|
||||
|
||||
requires_auth = pytest.mark.skipif(
|
||||
not _has_auth(),
|
||||
reason="Set WRENN_API_KEY or WRENN_TOKEN to run integration tests",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client() -> Generator[WrennClient, None, None]:
|
||||
with WrennClient(
|
||||
api_key=WRENN_API_KEY,
|
||||
token=WRENN_TOKEN,
|
||||
base_url=WRENN_BASE_URL,
|
||||
) as c:
|
||||
yield c
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def async_client() -> AsyncWrennClient:
|
||||
return AsyncWrennClient(
|
||||
api_key=WRENN_API_KEY,
|
||||
token=WRENN_TOKEN,
|
||||
base_url=WRENN_BASE_URL,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bearer_client() -> Generator[WrennClient, None, None]:
|
||||
if WRENN_TOKEN:
|
||||
with WrennClient(token=WRENN_TOKEN, base_url=WRENN_BASE_URL) as c:
|
||||
yield c
|
||||
elif WRENN_TEST_EMAIL and WRENN_TEST_PASSWORD:
|
||||
with WrennClient(
|
||||
api_key=WRENN_API_KEY, token=WRENN_TOKEN, base_url=WRENN_BASE_URL
|
||||
) as c:
|
||||
resp = c.auth.login(WRENN_TEST_EMAIL, WRENN_TEST_PASSWORD)
|
||||
with WrennClient(token=resp.token, base_url=WRENN_BASE_URL) as c:
|
||||
yield c
|
||||
else:
|
||||
pytest.skip(
|
||||
"Set WRENN_TOKEN or WRENN_TEST_EMAIL+WRENN_TEST_PASSWORD for bearer-auth tests"
|
||||
)
|
||||
|
||||
|
||||
@requires_auth
|
||||
class TestSandboxLifecycle:
|
||||
def test_create_exec_destroy(self, client):
|
||||
with client.sandboxes.create(template="minimal", timeout_sec=120) as sb:
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
result = sb.exec("echo", args=["hello"])
|
||||
assert result.exit_code == 0
|
||||
assert "hello" in result.stdout
|
||||
|
||||
def test_exec_with_args(self, client):
|
||||
with client.sandboxes.create(template="minimal", timeout_sec=120) as sb:
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
result = sb.exec("echo", args=["hello", "world"])
|
||||
assert result.exit_code == 0
|
||||
assert "hello world" in result.stdout
|
||||
|
||||
def test_exec_nonzero_exit(self, client):
|
||||
with client.sandboxes.create(template="minimal", timeout_sec=120) as sb:
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
result = sb.exec("sh", args=["-c", "exit 42"])
|
||||
assert result.exit_code == 42
|
||||
|
||||
def test_exec_stderr(self, client):
|
||||
with client.sandboxes.create(template="minimal", timeout_sec=120) as sb:
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
result = sb.exec("sh", args=["-c", "echo err>&2"])
|
||||
assert result.exit_code == 0
|
||||
assert "err" in result.stderr
|
||||
|
||||
def test_context_manager_cleanup(self, client):
|
||||
sb = client.sandboxes.create(template="minimal", timeout_sec=120)
|
||||
sb_id = sb.id
|
||||
|
||||
with sb:
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
|
||||
fetched = client.sandboxes.get(sb_id)
|
||||
assert fetched.status in ("stopped", "destroyed")
|
||||
|
||||
|
||||
@requires_auth
|
||||
class TestFileIO:
|
||||
def test_upload_and_download(self, client):
|
||||
with client.sandboxes.create(template="minimal", timeout_sec=120) as sb:
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
content = b"Hello from integration test!"
|
||||
sb.upload("/tmp/test_file.txt", content)
|
||||
downloaded = sb.download("/tmp/test_file.txt")
|
||||
assert downloaded == content
|
||||
|
||||
def test_download_nonexistent_file(self, client):
|
||||
with client.sandboxes.create(template="minimal", timeout_sec=120) as sb:
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
with pytest.raises(Exception):
|
||||
sb.download("/tmp/no_such_file_12345")
|
||||
|
||||
|
||||
@requires_auth
|
||||
class TestPauseResume:
|
||||
def test_pause_and_resume(self, client):
|
||||
with client.sandboxes.create(template="minimal", timeout_sec=120) as sb:
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
sb.pause()
|
||||
assert sb.status == "paused"
|
||||
|
||||
sb.resume()
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
|
||||
result = sb.exec("echo", args=["resumed"])
|
||||
assert result.exit_code == 0
|
||||
assert "resumed" in result.stdout
|
||||
|
||||
|
||||
@requires_auth
|
||||
class TestPing:
|
||||
def test_ping_resets_timer(self, client):
|
||||
with client.sandboxes.create(template="minimal", timeout_sec=120) as sb:
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
sb.ping()
|
||||
result = sb.exec("echo", args=["still_alive"])
|
||||
assert result.exit_code == 0
|
||||
assert "still_alive" in result.stdout
|
||||
|
||||
|
||||
@requires_auth
|
||||
class TestProxy:
|
||||
def test_get_url(self, client):
|
||||
with client.sandboxes.create(template="minimal", timeout_sec=120) as sb:
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
url = sb.get_url(8888)
|
||||
assert sb.id in url
|
||||
assert "8888" in url
|
||||
|
||||
|
||||
@requires_auth
|
||||
class TestListAndGet:
|
||||
def test_list_sandboxes(self, client):
|
||||
with client.sandboxes.create(template="minimal", timeout_sec=120) as sb:
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
boxes = client.sandboxes.list()
|
||||
ids = [b.id for b in boxes]
|
||||
assert sb.id in ids
|
||||
|
||||
def test_get_existing_sandbox(self, client):
|
||||
with client.sandboxes.create(template="minimal", timeout_sec=120) as sb:
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
fetched = client.sandboxes.get(sb.id)
|
||||
assert fetched.id == sb.id
|
||||
assert fetched.status == "running"
|
||||
|
||||
def test_get_nonexistent_sandbox(self, client):
|
||||
with pytest.raises((WrennNotFoundError, WrennValidationError)):
|
||||
client.sandboxes.get("cl-nonexistent00000000000000000")
|
||||
|
||||
|
||||
@requires_auth
|
||||
class TestSnapshots:
|
||||
def test_list_templates(self, client):
|
||||
templates = client.snapshots.list()
|
||||
assert isinstance(templates, list)
|
||||
|
||||
|
||||
@requires_auth
|
||||
class TestAPIKeys:
|
||||
def test_create_list_delete(self, bearer_client):
|
||||
key_resp = bearer_client.api_keys.create(name="integration-test-key")
|
||||
assert key_resp.name == "integration-test-key"
|
||||
assert key_resp.key is not None
|
||||
assert key_resp.id is not None
|
||||
|
||||
try:
|
||||
keys = bearer_client.api_keys.list()
|
||||
ids = [k.id for k in keys]
|
||||
assert key_resp.id in ids
|
||||
finally:
|
||||
bearer_client.api_keys.delete(key_resp.id)
|
||||
|
||||
|
||||
@requires_auth
|
||||
class TestRunCode:
|
||||
def test_basic_execution(self, client):
|
||||
with client.sandboxes.create(
|
||||
template="python-interpreter-v0-beta", timeout_sec=120
|
||||
) as sb:
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
|
||||
r = sb.run_code("x = 42")
|
||||
assert r.error is None
|
||||
|
||||
r = sb.run_code("x * 2")
|
||||
assert r.text == "84"
|
||||
|
||||
def test_state_persists(self, client):
|
||||
with client.sandboxes.create(
|
||||
template="python-interpreter-v0-beta", timeout_sec=120
|
||||
) as sb:
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
|
||||
sb.run_code("def greet(name): return f'hello {name}'")
|
||||
r = sb.run_code("greet('sandbox')")
|
||||
assert "hello sandbox" in (r.text or "")
|
||||
|
||||
def test_error_traceback(self, client):
|
||||
with client.sandboxes.create(
|
||||
template="python-interpreter-v0-beta", timeout_sec=120
|
||||
) as sb:
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
|
||||
r = sb.run_code("1/0")
|
||||
assert r.error is not None
|
||||
assert "ZeroDivisionError" in r.error
|
||||
|
||||
def test_stdout_capture(self, client):
|
||||
with client.sandboxes.create(
|
||||
template="python-interpreter-v0-beta", timeout_sec=120
|
||||
) as sb:
|
||||
sb.wait_ready(timeout=60, interval=1)
|
||||
|
||||
r = sb.run_code("print('hello from kernel')")
|
||||
assert "hello from kernel" in r.stdout
|
||||
|
||||
|
||||
@requires_auth
|
||||
class TestAsyncSandboxLifecycle:
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_create_exec_destroy(self, async_client):
|
||||
async with async_client:
|
||||
sb = await async_client.sandboxes.create(
|
||||
template="minimal", timeout_sec=120
|
||||
)
|
||||
try:
|
||||
await sb.async_wait_ready(timeout=60, interval=1)
|
||||
result = await sb.async_exec("echo", args=["async_hello"])
|
||||
assert result.exit_code == 0
|
||||
assert "async_hello" in result.stdout
|
||||
finally:
|
||||
await sb.async_destroy()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_upload_download(self, async_client):
|
||||
async with async_client:
|
||||
sb = await async_client.sandboxes.create(
|
||||
template="minimal", timeout_sec=120
|
||||
)
|
||||
try:
|
||||
await sb.async_wait_ready(timeout=60, interval=1)
|
||||
content = b"Async upload test"
|
||||
await sb.async_upload("/tmp/async_test.txt", content)
|
||||
downloaded = await sb.async_download("/tmp/async_test.txt")
|
||||
assert downloaded == content
|
||||
finally:
|
||||
await sb.async_destroy()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_run_code(self, async_client):
|
||||
async with async_client:
|
||||
sb = await async_client.sandboxes.create(
|
||||
template="python-interpreter-v0-beta", timeout_sec=120
|
||||
)
|
||||
try:
|
||||
await sb.async_wait_ready(timeout=60, interval=1)
|
||||
r = await sb.async_run_code("42 * 2")
|
||||
assert r.text == "84"
|
||||
finally:
|
||||
await sb.async_destroy()
|
||||
175
tests/test_sandbox_features.py
Normal file
175
tests/test_sandbox_features.py
Normal file
@ -0,0 +1,175 @@
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
import pytest
|
||||
import respx
|
||||
|
||||
from wrenn.client import WrennClient
|
||||
from wrenn.exceptions import WrennAuthenticationError
|
||||
from wrenn.sandbox import CodeResult, Sandbox, _build_proxy_url
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client():
|
||||
with WrennClient(api_key="wrn_test1234567890abcdef12345678") as c:
|
||||
yield c
|
||||
|
||||
|
||||
class TestBuildProxyUrl:
|
||||
def test_https_production(self):
|
||||
url = _build_proxy_url("https://api.wrenn.dev", "cl-abc123", 8888)
|
||||
assert url == "wss://8888-cl-abc123.api.wrenn.dev"
|
||||
|
||||
def test_http_localhost(self):
|
||||
url = _build_proxy_url("http://localhost:8080", "cl-abc123", 3000)
|
||||
assert url == "ws://3000-cl-abc123.localhost:8080"
|
||||
|
||||
def test_https_custom_port(self):
|
||||
url = _build_proxy_url("https://api.example.com:9443", "sb-1", 8080)
|
||||
assert url == "wss://8080-sb-1.api.example.com:9443"
|
||||
|
||||
def test_http_no_port(self):
|
||||
url = _build_proxy_url("http://192.168.1.1", "sb-2", 5000)
|
||||
assert url == "ws://5000-sb-2.192.168.1.1"
|
||||
|
||||
|
||||
class TestSandboxGetUrl:
|
||||
@respx.mock
|
||||
def test_get_url_returns_proxy_url(self, client):
|
||||
respx.post("https://api.wrenn.dev/v1/sandboxes").respond(
|
||||
201, json={"id": "cl-abc", "status": "pending"}
|
||||
)
|
||||
sb = client.sandboxes.create(template="minimal")
|
||||
url = sb.get_url(8888)
|
||||
assert url == "wss://8888-cl-abc.api.wrenn.dev"
|
||||
|
||||
@respx.mock
|
||||
def test_get_url_localhost(self):
|
||||
with WrennClient(
|
||||
api_key="wrn_test1234567890abcdef12345678",
|
||||
base_url="http://localhost:8080",
|
||||
) as c:
|
||||
respx.post("http://localhost:8080/v1/sandboxes").respond(
|
||||
201, json={"id": "cl-xyz", "status": "pending"}
|
||||
)
|
||||
sb = c.sandboxes.create()
|
||||
url = sb.get_url(3000)
|
||||
assert url == "ws://3000-cl-xyz.localhost:8080"
|
||||
|
||||
|
||||
class TestProxyAuthGuard:
|
||||
def test_jwt_only_get_url_raises(self):
|
||||
with WrennClient(token="jwt-abc") as c:
|
||||
sb = Sandbox(id="cl-abc")
|
||||
sb._bind(c._http, str(c._http.base_url), api_key=None, token="jwt-abc")
|
||||
with pytest.raises(WrennAuthenticationError):
|
||||
sb.get_url(8888)
|
||||
|
||||
def test_jwt_only_http_client_raises(self):
|
||||
with WrennClient(token="jwt-abc") as c:
|
||||
sb = Sandbox(id="cl-abc")
|
||||
sb._bind(c._http, str(c._http.base_url), api_key=None, token="jwt-abc")
|
||||
with pytest.raises(WrennAuthenticationError):
|
||||
_ = sb.http_client
|
||||
|
||||
|
||||
class TestSandboxHttpClient:
|
||||
@respx.mock
|
||||
def test_http_client_has_api_key_header(self, client):
|
||||
respx.post("https://api.wrenn.dev/v1/sandboxes").respond(
|
||||
201, json={"id": "cl-abc", "status": "pending"}
|
||||
)
|
||||
sb = client.sandboxes.create()
|
||||
hc = sb.http_client
|
||||
assert hc.headers["X-API-Key"] == "wrn_test1234567890abcdef12345678"
|
||||
|
||||
@respx.mock
|
||||
def test_http_client_sends_to_proxy(self, client):
|
||||
route = respx.get("https://8888-cl-abc.api.wrenn.dev/api/kernels").respond(
|
||||
200, json=[]
|
||||
)
|
||||
respx.post("https://api.wrenn.dev/v1/sandboxes").respond(
|
||||
201, json={"id": "cl-abc", "status": "pending"}
|
||||
)
|
||||
sb = client.sandboxes.create()
|
||||
resp = sb.http_client.get("/api/kernels")
|
||||
assert resp.status_code == 200
|
||||
assert route.called
|
||||
|
||||
|
||||
class TestCreateReturnsBoundSandbox:
|
||||
@respx.mock
|
||||
def test_create_returns_sandbox_subclass(self, client):
|
||||
respx.post("https://api.wrenn.dev/v1/sandboxes").respond(
|
||||
201, json={"id": "cl-1", "status": "pending", "template": "minimal"}
|
||||
)
|
||||
sb = client.sandboxes.create(template="minimal")
|
||||
assert isinstance(sb, Sandbox)
|
||||
assert sb.id == "cl-1"
|
||||
assert hasattr(sb, "exec")
|
||||
assert hasattr(sb, "run_code")
|
||||
assert hasattr(sb, "get_url")
|
||||
|
||||
@respx.mock
|
||||
def test_create_context_manager(self, client):
|
||||
route = respx.delete("https://api.wrenn.dev/v1/sandboxes/cl-1").respond(204)
|
||||
respx.post("https://api.wrenn.dev/v1/sandboxes").respond(
|
||||
201, json={"id": "cl-1", "status": "pending"}
|
||||
)
|
||||
sb = client.sandboxes.create()
|
||||
with sb:
|
||||
assert sb.id == "cl-1"
|
||||
assert route.called
|
||||
|
||||
|
||||
class TestCodeResult:
|
||||
def test_defaults(self):
|
||||
r = CodeResult()
|
||||
assert r.text is None
|
||||
assert r.data is None
|
||||
assert r.stdout == ""
|
||||
assert r.stderr == ""
|
||||
assert r.error is None
|
||||
|
||||
def test_with_values(self):
|
||||
r = CodeResult(
|
||||
text="84",
|
||||
data={"text/plain": "84"},
|
||||
stdout="",
|
||||
stderr="",
|
||||
error=None,
|
||||
)
|
||||
assert r.text == "84"
|
||||
assert r.data["text/plain"] == "84"
|
||||
|
||||
def test_error_result(self):
|
||||
r = CodeResult(error="ZeroDivisionError: division by zero\n...")
|
||||
assert r.error is not None
|
||||
assert "ZeroDivisionError" in r.error
|
||||
|
||||
|
||||
class TestRunCodeAuthGuard:
|
||||
def test_jwt_only_run_code_raises(self):
|
||||
with WrennClient(token="jwt-abc") as c:
|
||||
sb = Sandbox(id="cl-abc")
|
||||
sb._bind(c._http, str(c._http.base_url), api_key=None, token="jwt-abc")
|
||||
with pytest.raises(WrennAuthenticationError):
|
||||
sb.run_code("print(1)")
|
||||
|
||||
|
||||
class TestJupyterMessageFormat:
|
||||
def test_execute_request_structure(self):
|
||||
sb = Sandbox(id="test")
|
||||
msg = sb._jupyter_execute_request("x = 42")
|
||||
assert msg["msg_type"] == "execute_request"
|
||||
assert msg["content"]["code"] == "x = 42"
|
||||
assert msg["content"]["silent"] is False
|
||||
assert "msg_id" in msg
|
||||
assert "header" in msg
|
||||
assert msg["header"]["msg_type"] == "execute_request"
|
||||
|
||||
def test_execute_request_unique_ids(self):
|
||||
sb = Sandbox(id="test")
|
||||
m1 = sb._jupyter_execute_request("a")
|
||||
m2 = sb._jupyter_execute_request("b")
|
||||
assert m1["msg_id"] != m2["msg_id"]
|
||||
Reference in New Issue
Block a user