First iteration of changes to SDK handling new client secrets and token-exchange

Author: mateuszkuprowski

_test_unstructured_client/unit/auth/__init__.py

@@ -0,0 +1 @@
+

_test_unstructured_client/unit/auth/_mock_transport.py

@@ -0,0 +1,95 @@
+"""Shared helpers for exercising token-exchange auth callables.
+
+The mock transports here let tests script a sequence of responses / exceptions
+for the ``POST /auth/token-exchange`` endpoint without standing up a real
+account-service.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any, Callable, Iterable, List, Optional, Union
+
+import httpx
+
+
+ResponseStep = Union[httpx.Response, Exception, Callable[[httpx.Request], httpx.Response]]
+
+
+class ScriptedTransport(httpx.MockTransport):
+    """A MockTransport that walks through a scripted sequence of responses.
+
+    Each element can be an :class:`httpx.Response`, an ``Exception`` instance
+    (raised instead of returned), or a callable that accepts the request and
+    returns a response. Tests can inspect :attr:`requests` to assert how many
+    exchanges took place and what bodies were sent.
+    """
+
+    def __init__(self, steps: Iterable[ResponseStep]) -> None:
+        self._steps: List[ResponseStep] = list(steps)
+        self.requests: List[httpx.Request] = []
+        super().__init__(self._handler)
+
+    def _handler(self, request: httpx.Request) -> httpx.Response:
+        self.requests.append(request)
+        if not self._steps:
+            raise AssertionError(
+                "ScriptedTransport exhausted; unexpected extra request to "
+                f"{request.url}",
+            )
+        step = self._steps.pop(0)
+        if isinstance(step, Exception):
+            raise step
+        if callable(step):
+            return step(request)
+        return step
+
+
+class AsyncScriptedTransport(httpx.MockTransport):
+    """Async counterpart to :class:`ScriptedTransport`."""
+
+    def __init__(self, steps: Iterable[ResponseStep]) -> None:
+        self._steps: List[ResponseStep] = list(steps)
+        self.requests: List[httpx.Request] = []
+
+        async def _handler(request: httpx.Request) -> httpx.Response:
+            self.requests.append(request)
+            if not self._steps:
+                raise AssertionError(
+                    "AsyncScriptedTransport exhausted; unexpected extra "
+                    f"request to {request.url}",
+                )
+            step = self._steps.pop(0)
+            if isinstance(step, Exception):
+                raise step
+            if callable(step):
+                return step(request)
+            return step
+
+        super().__init__(_handler)
+
+
+def exchange_response(
+    access_token: Optional[str] = "jwt-1",
+    *,
+    expires_in: int = 900,
+    token_exchange_enabled: bool = True,
+    token_type: str = "bearer",
+    status_code: int = 200,
+    extra: Optional[dict] = None,
+) -> httpx.Response:
+    """Build a canned ``/auth/token-exchange`` response body."""
+    body: dict[str, Any] = {
+        "access_token": access_token,
+        "token_type": token_type,
+        "expires_in": expires_in,
+        "token_exchange_enabled": token_exchange_enabled,
+    }
+    if extra:
+        body.update(extra)
+    return httpx.Response(status_code, json=body)
+
+
+def body_of(request: httpx.Request) -> dict:
+    """Decode the JSON body from an outgoing exchange request."""
+    return json.loads(request.content.decode("utf-8"))

_test_unstructured_client/unit/auth/test_async_client_credentials.py

@@ -0,0 +1,167 @@
+"""Unit tests for :class:`unstructured_client.auth.AsyncClientCredentials`."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import List
+
+import httpx
+import pytest
+
+from unstructured_client.auth import (
+    AsyncClientCredentials,
+    InvalidCredentialError,
+    TokenExchangeError,
+)
+
+from ._mock_transport import AsyncScriptedTransport, body_of, exchange_response
+
+SERVER_URL = "https://accounts.example.test"
+SECRET = "uns_sk_async_example"
+
+
+@pytest.fixture(autouse=True)
+def _no_sleep(monkeypatch):
+    async def _noop(*_args, **_kwargs):
+        return None
+
+    monkeypatch.setattr(
+        "unstructured_client.auth.client_credentials.asyncio.sleep",
+        _noop,
+    )
+
+
+@pytest.fixture
+def fake_clock(monkeypatch):
+    state = {"now": 2_000_000.0}
+
+    def _now() -> float:
+        return state["now"]
+
+    monkeypatch.setattr("unstructured_client.auth._base.time.monotonic", _now)
+    monkeypatch.setattr(
+        "unstructured_client.auth.client_credentials.time.monotonic", _now
+    )
+    return state
+
+
+class DescribeAsyncClientCredentials:
+    @pytest.mark.asyncio
+    async def it_exchanges_then_caches(self, fake_clock):
+        transport = AsyncScriptedTransport(
+            [exchange_response(access_token="jwt-1", expires_in=900)]
+        )
+        http_client = httpx.AsyncClient(transport=transport)
+        acc = AsyncClientCredentials(
+            client_secret=SECRET,
+            server_url=SERVER_URL,
+            http_client=http_client,
+        )
+
+        first = await acc.acquire()
+        second = await acc.acquire()
+
+        assert first == second == "jwt-1"
+        assert len(transport.requests) == 1
+        assert body_of(transport.requests[0]) == {
+            "grant_type": "client_credentials",
+            "client_secret": SECRET,
+        }
+
+    @pytest.mark.asyncio
+    async def it_raises_invalid_credential_on_401(self, fake_clock):
+        transport = AsyncScriptedTransport(
+            [httpx.Response(401, json={"detail": "bad"})]
+        )
+        http_client = httpx.AsyncClient(transport=transport)
+        acc = AsyncClientCredentials(
+            client_secret=SECRET,
+            server_url=SERVER_URL,
+            http_client=http_client,
+            max_retries=5,
+        )
+
+        with pytest.raises(InvalidCredentialError):
+            await acc.acquire()
+
+    @pytest.mark.asyncio
+    async def it_retries_5xx_then_succeeds(self, fake_clock):
+        transport = AsyncScriptedTransport(
+            [
+                httpx.Response(500),
+                httpx.Response(502),
+                exchange_response(access_token="jwt-1", expires_in=900),
+            ]
+        )
+        http_client = httpx.AsyncClient(transport=transport)
+        acc = AsyncClientCredentials(
+            client_secret=SECRET,
+            server_url=SERVER_URL,
+            http_client=http_client,
+            max_retries=3,
+        )
+
+        assert await acc.acquire() == "jwt-1"
+        assert len(transport.requests) == 3
+
+    @pytest.mark.asyncio
+    async def it_serializes_concurrent_acquires(self, fake_clock):
+        """Ten concurrent ``acquire()`` calls must share one exchange."""
+        transport = AsyncScriptedTransport(
+            [exchange_response(access_token="jwt-1", expires_in=900)]
+        )
+        http_client = httpx.AsyncClient(transport=transport)
+        acc = AsyncClientCredentials(
+            client_secret=SECRET,
+            server_url=SERVER_URL,
+            http_client=http_client,
+        )
+
+        results: List[str] = await asyncio.gather(*(acc.acquire() for _ in range(10)))
+
+        assert results == ["jwt-1"] * 10
+        assert len(transport.requests) == 1
+
+    @pytest.mark.asyncio
+    async def it_raises_outage_error_without_cached_token(self, fake_clock):
+        transport = AsyncScriptedTransport([httpx.Response(500)] * 4)
+        http_client = httpx.AsyncClient(transport=transport)
+        acc = AsyncClientCredentials(
+            client_secret=SECRET,
+            server_url=SERVER_URL,
+            http_client=http_client,
+            max_retries=3,
+        )
+
+        with pytest.raises(TokenExchangeError):
+            await acc.acquire()
+
+    def it_sync_call_works_outside_running_loop(self, fake_clock):
+        """``__call__`` is the SDK entry point; must work without a loop."""
+        transport = AsyncScriptedTransport(
+            [exchange_response(access_token="jwt-1", expires_in=900)]
+        )
+        http_client = httpx.AsyncClient(transport=transport)
+        acc = AsyncClientCredentials(
+            client_secret=SECRET,
+            server_url=SERVER_URL,
+            http_client=http_client,
+        )
+
+        assert acc() == "jwt-1"
+
+    @pytest.mark.asyncio
+    async def it_sync_call_works_inside_running_loop(self, fake_clock):
+        """Driving __call__ from a running loop offloads to a worker thread."""
+        transport = AsyncScriptedTransport(
+            [exchange_response(access_token="jwt-1", expires_in=900)]
+        )
+        http_client = httpx.AsyncClient(transport=transport)
+        acc = AsyncClientCredentials(
+            client_secret=SECRET,
+            server_url=SERVER_URL,
+            http_client=http_client,
+        )
+
+        token = await asyncio.to_thread(acc)
+        assert token == "jwt-1"