diff --git a/src/kernel/lib/browser_pools/__init__.py b/src/kernel/lib/browser_pools/__init__.py new file mode 100644 index 0000000..8bbaaf1 --- /dev/null +++ b/src/kernel/lib/browser_pools/__init__.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from .acquire import ( + Acquired, + TimedOut, + PoolNotFound, + AcquireResult, + acquire, + acquire_async, +) + +__all__ = [ + "Acquired", + "TimedOut", + "PoolNotFound", + "AcquireResult", + "acquire", + "acquire_async", +] diff --git a/src/kernel/lib/browser_pools/acquire.py b/src/kernel/lib/browser_pools/acquire.py new file mode 100644 index 0000000..a596a47 --- /dev/null +++ b/src/kernel/lib/browser_pools/acquire.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Union +from dataclasses import dataclass +from typing_extensions import TypeAlias + +import httpx + +from ..._types import Body, Omit, Query, Headers, NotGiven, omit, not_given +from ..._exceptions import NotFoundError +from ...types.browser_pool_acquire_response import BrowserPoolAcquireResponse + +if TYPE_CHECKING: + from ..._client import Kernel, AsyncKernel + + +@dataclass +class Acquired: + """A browser was leased from the pool.""" + + browser: BrowserPoolAcquireResponse + + +@dataclass +class TimedOut: + """The long poll expired before a browser became available. Retry to keep waiting.""" + + +@dataclass +class PoolNotFound: + """No pool exists with the given id or name.""" + + +AcquireResult: TypeAlias = Union[Acquired, TimedOut, PoolNotFound] + + +def acquire( + client: "Kernel", + id_or_name: str, + *, + acquire_timeout_seconds: Union[int, Omit] = omit, + extra_headers: Union[Headers, None] = None, + extra_query: Union[Query, None] = None, + extra_body: Union[Body, None] = None, + timeout: Union[float, httpx.Timeout, None, NotGiven] = not_given, +) -> AcquireResult: + """Long-polling acquire that surfaces the HTTP outcome as a typed result. + + Returns one of: + + * :class:`Acquired` — a browser was leased from the pool. + * :class:`TimedOut` — the long poll expired without a browser becoming available. + Retry to keep waiting. + * :class:`PoolNotFound` — no pool exists with the given id or name. + + Other API errors (auth, server errors, etc.) still raise. + """ + try: + raw = client.browser_pools.with_raw_response.acquire( + id_or_name, + acquire_timeout_seconds=acquire_timeout_seconds, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) + except NotFoundError: + return PoolNotFound() + if raw.http_response.status_code == 204: + return TimedOut() + return Acquired(browser=raw.parse()) + + +async def acquire_async( + client: "AsyncKernel", + id_or_name: str, + *, + acquire_timeout_seconds: Union[int, Omit] = omit, + extra_headers: Union[Headers, None] = None, + extra_query: Union[Query, None] = None, + extra_body: Union[Body, None] = None, + timeout: Union[float, httpx.Timeout, None, NotGiven] = not_given, +) -> AcquireResult: + """Async variant of :func:`acquire`.""" + try: + raw = await client.browser_pools.with_raw_response.acquire( + id_or_name, + acquire_timeout_seconds=acquire_timeout_seconds, + extra_headers=extra_headers, + extra_query=extra_query, + extra_body=extra_body, + timeout=timeout, + ) + except NotFoundError: + return PoolNotFound() + if raw.http_response.status_code == 204: + return TimedOut() + return Acquired(browser=await raw.parse()) diff --git a/tests/test_browser_pools_typed_acquire.py b/tests/test_browser_pools_typed_acquire.py new file mode 100644 index 0000000..094ad3b --- /dev/null +++ b/tests/test_browser_pools_typed_acquire.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +import os + +import httpx +import respx +import pytest + +from kernel import Kernel, AsyncKernel +from kernel.lib.browser_pools import Acquired, TimedOut, PoolNotFound, acquire, acquire_async + +base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010") +api_key = "sk-123" + + +def _browser_payload() -> dict[str, object]: + return { + "session_id": "sess-1", + "base_url": "http://browser-session.test/browser/kernel", + "cdp_ws_url": "wss://browser-session.test/browser/cdp?jwt=t", + "webdriver_ws_url": "wss://x", + "created_at": "2020-01-01T00:00:00Z", + "headless": True, + "stealth": False, + "timeout_seconds": 60, + } + + +@respx.mock +def test_acquire_returns_acquired_on_200() -> None: + respx.post(f"{base_url}/browser_pools/my-pool/acquire").mock( + return_value=httpx.Response(200, json=_browser_payload()) + ) + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + result = acquire(client, "my-pool") + + assert isinstance(result, Acquired) + assert result.browser.session_id == "sess-1" + + +@respx.mock +def test_acquire_returns_timed_out_on_204() -> None: + respx.post(f"{base_url}/browser_pools/my-pool/acquire").mock(return_value=httpx.Response(204)) + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + result = acquire(client, "my-pool") + + assert isinstance(result, TimedOut) + + +@respx.mock +def test_acquire_returns_pool_not_found_on_404() -> None: + respx.post(f"{base_url}/browser_pools/missing/acquire").mock( + return_value=httpx.Response(404, json={"code": "not_found", "message": "pool not found"}) + ) + with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + result = acquire(client, "missing") + + assert isinstance(result, PoolNotFound) + + +@pytest.mark.asyncio +@respx.mock +async def test_acquire_async_returns_timed_out_on_204() -> None: + respx.post(f"{base_url}/browser_pools/my-pool/acquire").mock(return_value=httpx.Response(204)) + async with AsyncKernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client: + result = await acquire_async(client, "my-pool") + + assert isinstance(result, TimedOut)