Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions _test_unstructured_client/unit/auth/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@

95 changes: 95 additions & 0 deletions _test_unstructured_client/unit/auth/_mock_transport.py
Original file line number Diff line number Diff line change
@@ -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"))
167 changes: 167 additions & 0 deletions _test_unstructured_client/unit/auth/test_async_client_credentials.py
Original file line number Diff line number Diff line change
@@ -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"
Loading
Loading