diff --git a/tests/test_browser_routing.py b/tests/test_browser_routing.py index ce9a726f..9b37908e 100644 --- a/tests/test_browser_routing.py +++ b/tests/test_browser_routing.py @@ -1,7 +1,9 @@ from __future__ import annotations import os -from typing import Any, cast +import asyncio +from typing import Any, AsyncIterator, cast +from typing_extensions import override import httpx import respx @@ -119,6 +121,56 @@ def test_telemetry_stream_routes_directly_to_vm(monkeypatch: pytest.MonkeyPatch) assert request.headers.get("Authorization") is None +@pytest.mark.asyncio +async def test_async_telemetry_stream_cancellation_survives_direct_routing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", "telemetry/stream") + read_started = asyncio.Event() + read_stopped = asyncio.Event() + + class BlockingSSEStream(httpx.AsyncByteStream): + @override + async def __aiter__(self) -> AsyncIterator[bytes]: + read_started.set() + try: + await asyncio.Event().wait() + finally: + read_stopped.set() + yield b"" + + @override + async def aclose(self) -> None: + read_stopped.set() + + async def handle_request(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/browser/kernel/telemetry/stream" + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + stream=BlockingSSEStream(), + ) + + http_client = httpx.AsyncClient(transport=httpx.MockTransport(handle_request)) + async with AsyncKernel( + base_url=base_url, + api_key=api_key, + http_client=http_client, + _strict_response_validation=True, + ) as client: + route = browser_route_from_browser(_fake_browser()) + assert route is not None + client.browser_route_cache.set(route) + stream = await client.browsers.telemetry.stream("sess-1") + consumer = asyncio.create_task(stream.__anext__()) + await asyncio.wait_for(read_started.wait(), timeout=1) + + consumer.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(consumer, timeout=1) + await asyncio.wait_for(read_stopped.wait(), timeout=1) + + @respx.mock def test_browser_request_params_cannot_override_target_url_or_jwt() -> None: route = respx.get("http://browser-session.test/browser/kernel/curl/raw").mock(