diff --git a/src/wrenn/client.py b/src/wrenn/client.py index ac9770b..3723aaa 100644 --- a/src/wrenn/client.py +++ b/src/wrenn/client.py @@ -401,7 +401,9 @@ class EventsResource: print(ev.event, ev.resource) """ with self._http.stream("GET", "/v1/events/stream", timeout=None) as resp: - _raise_for_status(resp) + if resp.status_code >= 400: + resp.read() + _raise_for_status(resp) yield from _iter_sse_events(resp.iter_lines()) @@ -611,7 +613,9 @@ class AsyncEventsResource: print(ev.event, ev.resource) """ async with self._http.stream("GET", "/v1/events/stream", timeout=None) as resp: - _raise_for_status(resp) + if resp.status_code >= 400: + await resp.aread() + _raise_for_status(resp) data_lines: list[str] = [] async for raw in resp.aiter_lines(): if raw == "": diff --git a/tests/test_client.py b/tests/test_client.py index 78cc2bf..325ae28 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -13,9 +13,16 @@ from wrenn.exceptions import ( WrennValidationError, ) from wrenn.models import ( + Actor, Capsule, + CapsuleMetrics, + CapsuleStats, + MetricPoint, + Resource, + SSEEvent, Status, Template, + UsageResponse, ) BASE = "https://app.wrenn.dev/api" @@ -144,6 +151,69 @@ class TestCapsules: m = client.capsules.metrics("sb-1") assert m.sandbox_id == "sb-1" assert m.points and m.points[0].cpu_pct == 12.5 + assert isinstance(m, CapsuleMetrics) + assert isinstance(m.points[0], MetricPoint) + + @respx.mock + def test_stats_default_omits_range(self, client): + route = respx.get(f"{BASE}/v1/capsules/stats").respond( + 200, + json={"range": "1h", "current": {}, "peaks": {}, "series": {}}, + ) + stats = client.capsules.stats() + assert "range=" not in str(route.calls[0].request.url) + assert isinstance(stats, CapsuleStats) + + @respx.mock + def test_usage_default_omits_params(self, client): + route = respx.get(f"{BASE}/v1/capsules/usage").respond(200, json={"points": []}) + usage = client.capsules.usage() + url = str(route.calls[0].request.url) + assert "from=" not in url + assert "to=" not in url + assert isinstance(usage, UsageResponse) + + @respx.mock + def test_metrics_default_omits_range(self, client): + route = respx.get(f"{BASE}/v1/capsules/sb-1/metrics").respond( + 200, json={"sandbox_id": "sb-1", "points": []} + ) + client.capsules.metrics("sb-1") + assert "range=" not in str(route.calls[0].request.url) + + @respx.mock + def test_metrics_not_found(self, client): + respx.get(f"{BASE}/v1/capsules/nope/metrics").respond( + 404, + json={"error": {"code": "not_found", "message": "capsule not found"}}, + ) + with pytest.raises(WrennNotFoundError): + client.capsules.metrics("nope") + + @respx.mock + def test_stats_auth_error(self, client): + respx.get(f"{BASE}/v1/capsules/stats").respond( + 401, + json={"error": {"code": "unauthorized", "message": "bad key"}}, + ) + with pytest.raises(WrennAuthenticationError): + client.capsules.stats() + + @respx.mock + def test_usage_string_dates(self, client): + route = respx.get(f"{BASE}/v1/capsules/usage").respond(200, json={"points": []}) + client.capsules.usage(from_="2026-01-01", to="2026-01-31") + url = str(route.calls[0].request.url) + assert "from=2026-01-01" in url + assert "to=2026-01-31" in url + + @respx.mock + def test_usage_partial_dates(self, client): + route = respx.get(f"{BASE}/v1/capsules/usage").respond(200, json={"points": []}) + client.capsules.usage(from_="2026-01-01") + url = str(route.calls[0].request.url) + assert "from=2026-01-01" in url + assert "to=" not in url class TestEvents: @@ -170,6 +240,82 @@ class TestEvents: ] assert events[0].resource and events[0].resource.id == "sb-1" + @respx.mock + def test_stream_multiline_data(self, client): + body = ( + "event: capsule.create\n" + 'data: {"event":"capsule.create",\n' + 'data: "resource":{"id":"sb-1","type":"capsule"}}\n' + "\n" + ) + respx.get(f"{BASE}/v1/events/stream").respond( + 200, + headers={"content-type": "text/event-stream"}, + content=body, + ) + events = list(client.events.stream()) + assert len(events) == 1 + assert events[0].resource.id == "sb-1" + + @respx.mock + def test_stream_skips_keepalive_only(self, client): + body = ": keepalive\n\n: keepalive\n\n" + respx.get(f"{BASE}/v1/events/stream").respond( + 200, + headers={"content-type": "text/event-stream"}, + content=body, + ) + assert list(client.events.stream()) == [] + + @respx.mock + def test_stream_ignores_incomplete_trailing_frame(self, client): + body = ( + 'data: {"event":"capsule.create","resource":{"id":"sb-1","type":"capsule"}}\n' + "\n" + 'data: {"event":"capsule.destroy"' # no closing brace, no blank line + ) + respx.get(f"{BASE}/v1/events/stream").respond( + 200, + headers={"content-type": "text/event-stream"}, + content=body, + ) + events = list(client.events.stream()) + assert len(events) == 1 + assert events[0].event.value == "capsule.create" + + @respx.mock + def test_stream_full_payload_round_trip(self, client): + body = ( + 'data: {"event":"capsule.create",' + '"outcome":"success",' + '"resource":{"id":"sb-1","type":"capsule"},' + '"actor":{"type":"api_key","id":"key-1","name":"ci"},' + '"metadata":{"reason":"manual"}}\n' + "\n" + ) + respx.get(f"{BASE}/v1/events/stream").respond( + 200, + headers={"content-type": "text/event-stream"}, + content=body, + ) + events = list(client.events.stream()) + assert len(events) == 1 + ev = events[0] + assert isinstance(ev, SSEEvent) + assert isinstance(ev.resource, Resource) + assert isinstance(ev.actor, Actor) + assert ev.actor.name == "ci" + assert ev.metadata == {"reason": "manual"} + + @respx.mock + def test_stream_raises_on_4xx(self, client): + respx.get(f"{BASE}/v1/events/stream").respond( + 401, + json={"error": {"code": "unauthorized", "message": "bad key"}}, + ) + with pytest.raises(WrennAuthenticationError): + list(client.events.stream()) + class TestSnapshots: @respx.mock @@ -329,6 +475,156 @@ class TestAsyncClient: with pytest.raises(WrennNotFoundError): await async_client.capsules.get("nope") + @pytest.mark.asyncio + @respx.mock + async def test_async_stats(self, async_client): + async with async_client: + route = respx.get(f"{BASE}/v1/capsules/stats").respond( + 200, + json={ + "range": "24h", + "current": {"running_count": 1}, + "peaks": {"running_count": 5}, + "series": {"running": [1]}, + }, + ) + stats = await async_client.capsules.stats(range="24h") + assert "range=24h" in str(route.calls[0].request.url) + assert isinstance(stats, CapsuleStats) + assert stats.current.running_count == 1 + + @pytest.mark.asyncio + @respx.mock + async def test_async_usage(self, async_client): + import datetime as _dt + + async with async_client: + route = respx.get(f"{BASE}/v1/capsules/usage").respond( + 200, + json={ + "from": "2026-06-01", + "to": "2026-06-22", + "points": [ + { + "date": "2026-06-01", + "cpu_minutes": 1.5, + "ram_mb_minutes": 200.0, + } + ], + }, + ) + usage = await async_client.capsules.usage( + from_=_dt.date(2026, 6, 1), to="2026-06-22" + ) + url = str(route.calls[0].request.url) + assert "from=2026-06-01" in url + assert "to=2026-06-22" in url + assert isinstance(usage, UsageResponse) + assert usage.points and usage.points[0].cpu_minutes == 1.5 + + @pytest.mark.asyncio + @respx.mock + async def test_async_metrics(self, async_client): + async with async_client: + respx.get(f"{BASE}/v1/capsules/sb-1/metrics").respond( + 200, + json={ + "sandbox_id": "sb-1", + "range": "2h", + "points": [ + {"timestamp_unix": 1, "cpu_pct": 33.0, "mem_bytes": 1024} + ], + }, + ) + m = await async_client.capsules.metrics("sb-1", range="2h") + assert isinstance(m, CapsuleMetrics) + assert m.points[0].cpu_pct == 33.0 + + @pytest.mark.asyncio + @respx.mock + async def test_async_metrics_not_found(self, async_client): + async with async_client: + respx.get(f"{BASE}/v1/capsules/nope/metrics").respond( + 404, + json={"error": {"code": "not_found", "message": "not found"}}, + ) + with pytest.raises(WrennNotFoundError): + await async_client.capsules.metrics("nope") + + @pytest.mark.asyncio + @respx.mock + async def test_async_events_stream(self, async_client): + async with async_client: + body = ( + ": keepalive\n" + 'data: {"event":"capsule.create","resource":{"id":"sb-1","type":"capsule"}}\n' + "\n" + 'data: {"event":"capsule.destroy","resource":{"id":"sb-1","type":"capsule"}}\n' + "\n" + ) + respx.get(f"{BASE}/v1/events/stream").respond( + 200, + headers={"content-type": "text/event-stream"}, + content=body, + ) + events = [ev async for ev in async_client.events.stream()] + assert [e.event.value for e in events] == [ + "capsule.create", + "capsule.destroy", + ] + + @pytest.mark.asyncio + @respx.mock + async def test_async_events_raises_on_4xx(self, async_client): + async with async_client: + respx.get(f"{BASE}/v1/events/stream").respond( + 401, + json={"error": {"code": "unauthorized", "message": "bad key"}}, + ) + with pytest.raises(WrennAuthenticationError): + async for _ in async_client.events.stream(): + pass + + @pytest.mark.asyncio + @respx.mock + async def test_async_events_multiline_data(self, async_client): + async with async_client: + body = ( + 'data: {"event":"capsule.create",\n' + 'data: "resource":{"id":"sb-1","type":"capsule"}}\n' + "\n" + ) + respx.get(f"{BASE}/v1/events/stream").respond( + 200, + headers={"content-type": "text/event-stream"}, + content=body, + ) + events = [ev async for ev in async_client.events.stream()] + assert len(events) == 1 + assert events[0].resource.id == "sb-1" + + +class TestPackageSurface: + def test_version(self): + import wrenn + + assert wrenn.__version__ == "0.3.0" + + def test_new_models_exported(self): + import wrenn.models as m + + for name in ( + "Actor", + "CapsuleMetrics", + "CapsuleStats", + "MetricPoint", + "Resource", + "SSEEvent", + "UsageResponse", + ): + assert hasattr(m, name), name + assert name in m.__all__ + class TestClientResolution: def test_default_base_url_strips_app_subdomain(self): diff --git a/tests/test_code_runner_e2e.py b/tests/test_code_runner_e2e.py index 12cae7c..7189124 100644 --- a/tests/test_code_runner_e2e.py +++ b/tests/test_code_runner_e2e.py @@ -428,13 +428,14 @@ class TestCodeRunnerMimeTypes: assert "array([0, 1, 2, 3, 4])" in ex.text def test_requests_status_code(self): - ex = self._run( + ex = self.capsule.run_code( "import requests\n" "r = requests.get('https://httpbin.org/status/204', timeout=10)\n" - "r.status_code\n" + "r.status_code\n", + timeout=60, ) if ex.error is not None: - pytest.skip(f"network unavailable: {ex.error.name}") + pytest.skip(f"network unavailable: {ex.error.name}: {ex.error.value}") assert ex.text == "204" diff --git a/tests/test_code_runner_unit.py b/tests/test_code_runner_unit.py index d181749..c5be0da 100644 --- a/tests/test_code_runner_unit.py +++ b/tests/test_code_runner_unit.py @@ -388,6 +388,214 @@ class TestJupyterRequest: assert a["header"]["msg_id"] != b["header"]["msg_id"] +# ───────────────────────── _protocol direct helpers ───────────────────────── + + +class TestPickKernelId: + def test_returns_first_matching_kernel(self): + from wrenn.code_runner._protocol import pick_kernel_id + + kernels = [ + {"id": "k-1", "name": "python3"}, + {"id": "k-2", "name": "wrenn"}, + {"id": "k-3", "name": "wrenn"}, + ] + assert pick_kernel_id(kernels, "wrenn") == "k-2" + + def test_returns_none_when_no_match(self): + from wrenn.code_runner._protocol import pick_kernel_id + + kernels = [{"id": "k-1", "name": "python3"}] + assert pick_kernel_id(kernels, "wrenn") is None + + def test_returns_none_on_empty_list(self): + from wrenn.code_runner._protocol import pick_kernel_id + + assert pick_kernel_id([], "wrenn") is None + + def test_ignores_entries_missing_name(self): + from wrenn.code_runner._protocol import pick_kernel_id + + kernels = [{"id": "k-1"}, {"id": "k-2", "name": "wrenn"}] + assert pick_kernel_id(kernels, "wrenn") == "k-2" + + +class TestValidateLanguage: + def test_python_accepted(self): + from wrenn.code_runner._protocol import validate_language + + validate_language("python") # no raise + + def test_other_language_raises(self): + from wrenn.code_runner._protocol import validate_language + + with pytest.raises(ValueError, match="not supported"): + validate_language("r") + + +class TestBuildWsUrl: + def test_uses_proxy_domain_when_given(self): + from wrenn.code_runner._protocol import build_ws_url + + url = build_ws_url( + base_url="https://app.wrenn.dev/api", + capsule_id="sb-1", + kernel_id="k-1", + proxy_domain="wrenn.dev", + ) + assert url.startswith("wss://") + assert "8888-sb-1.wrenn.dev" in url + assert url.endswith("/api/kernels/k-1/channels") + + def test_falls_back_to_base_host(self): + from wrenn.code_runner._protocol import build_ws_url + + url = build_ws_url( + base_url="http://localhost:8080/api", + capsule_id="sb-1", + kernel_id="k-1", + ) + # localhost stays http→ws (not wss) only if helper preserves scheme + assert "8888-sb-1.localhost:8080" in url + assert url.endswith("/api/kernels/k-1/channels") + + +class TestApplyKernelMessage: + def _make_execution(self): + from wrenn.code_runner.models import Execution, Logs + + return Execution(results=[], logs=Logs(stdout=[], stderr=[]), error=None) + + def test_ignores_other_parent(self): + from wrenn.code_runner._protocol import apply_kernel_message + + execution = self._make_execution() + msg = { + "msg_type": "stream", + "header": {"msg_type": "stream"}, + "parent_header": {"msg_id": "other"}, + "content": {"name": "stdout", "text": "hi"}, + } + emitted: list = [] + done = apply_kernel_message( + msg, "mine", execution, emitted.append, None, None, None + ) + assert done is False + assert execution.logs.stdout == [] + + def test_stream_stdout_routed(self): + from wrenn.code_runner._protocol import apply_kernel_message + + execution = self._make_execution() + msg = { + "msg_type": "stream", + "parent_header": {"msg_id": "mine"}, + "content": {"name": "stdout", "text": "hi"}, + } + outs: list[str] = [] + apply_kernel_message( + msg, "mine", execution, lambda e: None, None, outs.append, None + ) + assert execution.logs.stdout == ["hi"] + assert outs == ["hi"] + + def test_stream_stderr_routed(self): + from wrenn.code_runner._protocol import apply_kernel_message + + execution = self._make_execution() + msg = { + "msg_type": "stream", + "parent_header": {"msg_id": "mine"}, + "content": {"name": "stderr", "text": "boom"}, + } + errs: list[str] = [] + apply_kernel_message( + msg, "mine", execution, lambda e: None, None, None, errs.append + ) + assert execution.logs.stderr == ["boom"] + assert errs == ["boom"] + + def test_execute_result_marks_main(self): + from wrenn.code_runner._protocol import apply_kernel_message + + execution = self._make_execution() + msg = { + "msg_type": "execute_result", + "parent_header": {"msg_id": "mine"}, + "content": { + "data": {"text/plain": "42"}, + "execution_count": 7, + }, + } + results: list = [] + apply_kernel_message( + msg, "mine", execution, lambda e: None, results.append, None, None + ) + assert execution.execution_count == 7 + assert execution.results[0].is_main_result is True + assert results and results[0].is_main_result is True + + def test_display_data_is_not_main(self): + from wrenn.code_runner._protocol import apply_kernel_message + + execution = self._make_execution() + msg = { + "msg_type": "display_data", + "parent_header": {"msg_id": "mine"}, + "content": {"data": {"text/plain": "x"}}, + } + apply_kernel_message(msg, "mine", execution, lambda e: None, None, None, None) + assert execution.results[0].is_main_result is False + assert execution.execution_count is None + + def test_error_message_emits_error(self): + from wrenn.code_runner._protocol import apply_kernel_message + + execution = self._make_execution() + msg = { + "msg_type": "error", + "parent_header": {"msg_id": "mine"}, + "content": { + "ename": "ValueError", + "evalue": "bad", + "traceback": ["line1", "line2"], + }, + } + emitted: list = [] + apply_kernel_message(msg, "mine", execution, emitted.append, None, None, None) + assert len(emitted) == 1 + assert emitted[0].name == "ValueError" + assert emitted[0].traceback == "line1\nline2" + + def test_idle_status_returns_true(self): + from wrenn.code_runner._protocol import apply_kernel_message + + execution = self._make_execution() + msg = { + "msg_type": "status", + "parent_header": {"msg_id": "mine"}, + "content": {"execution_state": "idle"}, + } + done = apply_kernel_message( + msg, "mine", execution, lambda e: None, None, None, None + ) + assert done is True + + def test_busy_status_returns_false(self): + from wrenn.code_runner._protocol import apply_kernel_message + + execution = self._make_execution() + msg = { + "msg_type": "status", + "parent_header": {"msg_id": "mine"}, + "content": {"execution_state": "busy"}, + } + done = apply_kernel_message( + msg, "mine", execution, lambda e: None, None, None, None + ) + assert done is False + + # ───────────────────────── run_code (WS-mocked) ───────────────────────── diff --git a/tests/test_integration.py b/tests/test_integration.py index 358065e..d98a695 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -7,8 +7,18 @@ from pathlib import Path import pytest from wrenn import Capsule, CommandResult +from wrenn.client import WrennClient from wrenn.commands import CommandHandle, ProcessInfo -from wrenn.models import Capsule as CapsuleModel, FileEntry, Status +from wrenn.exceptions import WrennNotFoundError +from wrenn.models import ( + Capsule as CapsuleModel, + CapsuleMetrics, + CapsuleStats, + FileEntry, + SSEEvent, + Status, + UsageResponse, +) pytestmark = pytest.mark.integration @@ -406,3 +416,124 @@ class TestGit: def test_get_config_missing_returns_none(self): value = self.capsule.git.get_config("nonexistent.key") assert value is None + + +class TestStatsUsageMetrics: + """Account-scoped + per-capsule observability endpoints.""" + + def setup_method(self): + _ensure_env() + + def test_stats_default(self): + with WrennClient() as client: + stats = client.capsules.stats() + assert isinstance(stats, CapsuleStats) + # current.running_count is always present even when zero + assert stats.current is not None + assert stats.peaks is not None + + def test_stats_range(self): + with WrennClient() as client: + stats = client.capsules.stats(range="24h") + assert isinstance(stats, CapsuleStats) + if stats.range is not None: + assert stats.range.value == "24h" + + def test_usage_no_args(self): + with WrennClient() as client: + usage = client.capsules.usage() + assert isinstance(usage, UsageResponse) + assert usage.points is not None + + def test_usage_date_range(self): + from datetime import date, timedelta + + today = date.today() + with WrennClient() as client: + usage = client.capsules.usage(from_=today - timedelta(days=7), to=today) + assert isinstance(usage, UsageResponse) + + def test_metrics_running_capsule(self): + capsule = Capsule(wait=True) + try: + with WrennClient() as client: + m = client.capsules.metrics(capsule.capsule_id) + assert isinstance(m, CapsuleMetrics) + assert m.sandbox_id == capsule.capsule_id + assert m.points is not None + finally: + capsule.destroy() + + def test_metrics_range_2h(self): + capsule = Capsule(wait=True) + try: + with WrennClient() as client: + m = client.capsules.metrics(capsule.capsule_id, range="2h") + assert isinstance(m, CapsuleMetrics) + finally: + capsule.destroy() + + def test_metrics_destroyed_capsule_raises(self): + # Create, destroy, then ask for metrics — guaranteed valid-shape ID + # that no longer exists. Server returns 404 for the missing capsule + # (avoids the 400 validation path for malformed IDs). + capsule = Capsule(wait=True) + capsule_id = capsule.capsule_id + capsule.destroy(wait=True) + + with WrennClient() as client: + with pytest.raises(WrennNotFoundError): + client.capsules.metrics(capsule_id) + + +class TestEventsStream: + """SSE lifecycle stream — drives a capsule create and watches for it.""" + + def setup_method(self): + _ensure_env() + + def test_stream_receives_capsule_create_event(self): + import threading + + seen: list[SSEEvent] = [] + ready = threading.Event() + stop = threading.Event() + + def _reader() -> None: + with WrennClient() as client: + stream = client.events.stream() + ready.set() + for ev in stream: + seen.append(ev) + if stop.is_set() or len(seen) > 20: + return + + t = threading.Thread(target=_reader, daemon=True) + t.start() + assert ready.wait(timeout=5) + + # Give the server a beat to fully open the SSE stream before + # the action that should appear on it. + time.sleep(0.5) + + capsule = Capsule() + capsule_id = capsule.capsule_id + try: + capsule.wait_ready(timeout=60) + finally: + capsule.destroy() + + # Wait for either our capsule's event or a timeout. + deadline = time.monotonic() + 15 + while time.monotonic() < deadline: + if any( + ev.resource is not None and ev.resource.id == capsule_id for ev in seen + ): + break + time.sleep(0.5) + + stop.set() + assert any(isinstance(ev, SSEEvent) for ev in seen) + # Best-effort: assert we saw a capsule.* event for our capsule. + matching = [ev for ev in seen if ev.resource and ev.resource.id == capsule_id] + assert matching, f"no events for {capsule_id} in {len(seen)} received"