"""Unit tests for the concrete transport Kernel.""" from __future__ import annotations import httpx import pytest import notebooklm._kernel as kernel_module from notebooklm._kernel import Kernel from notebooklm.auth import AuthTokens from notebooklm.types import ConnectionLimits def _auth_tokens() -> AuthTokens: return AuthTokens( csrf_token="csrf", session_id="sid", cookies={"SID": "cookie-value"}, storage_path=None, ) @pytest.mark.asyncio async def test_open_builds_http_client_and_captures_live_cookie_snapshot() -> None: kernel = Kernel() captured: list[httpx.Cookies] = [] await kernel.open( auth=_auth_tokens(), timeout=30.0, connect_timeout=10.0, limits=ConnectionLimits(), capture_cookie_snapshot=captured.append, ) try: assert kernel.http_client is not None assert len(captured) == 1 assert captured[0] is kernel.cookies finally: await kernel.aclose() @pytest.mark.asyncio async def test_open_is_idempotent() -> None: kernel = Kernel() captured: list[httpx.Cookies] = [] await kernel.open( auth=_auth_tokens(), timeout=30.0, connect_timeout=10.0, limits=ConnectionLimits(), capture_cookie_snapshot=captured.append, ) first_client = kernel.http_client await kernel.open( auth=_auth_tokens(), timeout=30.0, connect_timeout=10.0, limits=ConnectionLimits(), capture_cookie_snapshot=captured.append, ) try: assert kernel.http_client is first_client assert len(captured) == 1 finally: await kernel.aclose() @pytest.mark.asyncio async def test_open_preserves_explicit_empty_cookie_jar(monkeypatch: pytest.MonkeyPatch) -> None: def fail_build_cookie_jar(**_: object) -> httpx.Cookies: raise AssertionError("explicit cookie_jar should not be rebuilt") monkeypatch.setattr(kernel_module, "build_cookie_jar", fail_build_cookie_jar) kernel = Kernel() await kernel.open( auth=AuthTokens( csrf_token="csrf", session_id="sid", cookies={}, cookie_jar=httpx.Cookies(), storage_path=None, ), timeout=30.0, connect_timeout=10.0, limits=ConnectionLimits(), capture_cookie_snapshot=lambda _: None, ) try: assert kernel.http_client is not None finally: await kernel.aclose() @pytest.mark.asyncio async def test_open_closes_client_when_cookie_snapshot_raises() -> None: closed: list[httpx.AsyncClient] = [] class _TrackingClient(httpx.AsyncClient): async def aclose(self) -> None: closed.append(self) await super().aclose() def async_client_factory(**kwargs: object) -> httpx.AsyncClient: return _TrackingClient(**kwargs) # type: ignore[arg-type] kernel = Kernel(async_client_factory=async_client_factory) def boom(_: httpx.Cookies) -> None: raise RuntimeError("snapshot failed") with pytest.raises(RuntimeError, match="snapshot failed"): await kernel.open( auth=_auth_tokens(), timeout=30.0, connect_timeout=10.0, limits=ConnectionLimits(), capture_cookie_snapshot=boom, ) # The partially opened client must have been closed and the kernel reset so # it does not leak a live connection pool (issue #1163). assert kernel.http_client is None assert len(closed) == 1 assert closed[0].is_closed @pytest.mark.asyncio async def test_post_uses_live_http_client_streaming_post() -> None: seen_requests: list[httpx.Request] = [] async def handler(request: httpx.Request) -> httpx.Response: seen_requests.append(request) return httpx.Response(200, content=b"ok") transport = httpx.MockTransport(handler) def async_client_factory(**kwargs: object) -> httpx.AsyncClient: return httpx.AsyncClient(transport=transport, **kwargs) # type: ignore[arg-type] kernel = Kernel(async_client_factory=async_client_factory) await kernel.open( auth=_auth_tokens(), timeout=30.0, connect_timeout=10.0, limits=ConnectionLimits(), capture_cookie_snapshot=lambda _: None, ) try: response = await kernel.post( "https://example.com/batchexecute", headers={"X-Test": "yes"}, body=b"payload", ) finally: await kernel.aclose() assert response.text == "ok" assert len(seen_requests) == 1 request = seen_requests[0] assert request.method == "POST" assert str(request.url) == "https://example.com/batchexecute" assert request.headers["X-Test"] == "yes" assert request.content == b"payload" @pytest.mark.asyncio async def test_post_read_timeout_override_preserves_other_timeout_slots() -> None: seen_requests: list[httpx.Request] = [] async def handler(request: httpx.Request) -> httpx.Response: seen_requests.append(request) return httpx.Response(200, content=b"ok") transport = httpx.MockTransport(handler) def async_client_factory(**kwargs: object) -> httpx.AsyncClient: return httpx.AsyncClient(transport=transport, **kwargs) # type: ignore[arg-type] kernel = Kernel(async_client_factory=async_client_factory) await kernel.open( auth=_auth_tokens(), timeout=30.0, connect_timeout=10.0, limits=ConnectionLimits(), capture_cookie_snapshot=lambda _: None, ) try: await kernel.post( "https://example.com/chat", headers={}, body=b"payload", read_timeout=300.0, ) finally: await kernel.aclose() timeout = seen_requests[0].extensions["timeout"] assert timeout["connect"] == 10.0 assert timeout["read"] == 300.0 assert timeout["write"] == 30.0 assert timeout["pool"] == 30.0 @pytest.mark.asyncio async def test_post_forwards_response_cap_to_streaming_helper( monkeypatch: pytest.MonkeyPatch, ) -> None: captured: dict[str, object] = {} async def fake_stream_post( client: httpx.AsyncClient, url: str, *, body: bytes | str, headers: dict[str, str] | None, timeout: httpx.Timeout | float | None = None, max_bytes: int | None = None, ) -> httpx.Response: captured.update( { "client": client, "url": url, "body": body, "headers": headers, "timeout": timeout, "max_bytes": max_bytes, } ) return httpx.Response(200, content=b"ok") monkeypatch.setattr(kernel_module, "stream_post_with_size_cap", fake_stream_post) kernel = Kernel() await kernel.open( auth=_auth_tokens(), timeout=30.0, connect_timeout=10.0, limits=ConnectionLimits(), capture_cookie_snapshot=lambda _: None, ) try: response = await kernel.post( "https://example.com/chat", headers={"X-Test": "yes"}, body=b"payload", read_timeout=300.0, max_response_bytes=123456, ) finally: await kernel.aclose() assert response.text == "ok" assert captured["url"] == "https://example.com/chat" assert captured["body"] == b"payload" assert captured["headers"] == {"X-Test": "yes"} assert captured["max_bytes"] == 123456 timeout = captured["timeout"] assert isinstance(timeout, httpx.Timeout) assert timeout.as_dict()["read"] == 300.0 @pytest.mark.asyncio async def test_aclose_marks_kernel_closed_and_is_idempotent() -> None: kernel = Kernel() await kernel.open( auth=_auth_tokens(), timeout=30.0, connect_timeout=10.0, limits=ConnectionLimits(), capture_cookie_snapshot=lambda _: None, ) await kernel.aclose() await kernel.aclose() assert kernel.http_client is None with pytest.raises(RuntimeError, match="Client not initialized"): kernel.get_http_client()