from __future__ import annotations

import json

from typing import Any

import httpx2

import pytest

import openai._base_client as base_client

import openai.auth._workload as workload

from openai import (

OpenAI,

OAuthError,

AsyncOpenAI,

APITimeoutError,

APIConnectionError,

DefaultHttpx2Client,

DefaultAsyncHttpx2Client,

)

from openai.auth import WorkloadIdentity

def workload_identity(get_token: Any = lambda: "subject-token") -> WorkloadIdentity:

return {

"identity_provider_id": "idp_123",

"service_account_id": "sa_123",

"provider": {"get_token": get_token, "token_type": "jwt"},

}

def exchange_payload(access_token: str) -> dict[str, object]:

return {"access_token": access_token, "expires_in": 3600}

def test_sync_httpx2_workload_exchange_is_native_and_cached(monkeypatch: pytest.MonkeyPatch) -> None:

exchange_requests: list[Any] = []

api_requests: list[Any] = []

provider_calls = 0

def get_token() -> str:

nonlocal provider_calls

provider_calls += 1

return "subject-token"

def exchange_handler(request: Any) -> Any:

exchange_requests.append(request)

return httpx2.Response(200, request=request, json=exchange_payload("access-token"))

def exchange_client(**kwargs: Any) -> Any:

assert kwargs == {"follow_redirects": False}

return httpx2.Client(transport=httpx2.MockTransport(exchange_handler), trust_env=False)

def api_handler(request: Any) -> Any:

api_requests.append(request)

return httpx2.Response(200, request=request, json={"object": "list", "data": []})

monkeypatch.setattr(workload, "DefaultHttpx2Client", exchange_client)

with OpenAI(

workload_identity=workload_identity(get_token),

base_url="https://api.example.test/v1",

http_client=httpx2.Client(transport=httpx2.MockTransport(api_handler), trust_env=False),

max_retries=0,

) as client:

assert client.models.list().object == "list"

assert client.models.list().object == "list"

assert provider_calls == 1

assert len(exchange_requests) == 1

assert len(api_requests) == 2

assert isinstance(exchange_requests[0], httpx2.Request)

assert str(exchange_requests[0].url) == "https://auth.openai.com/oauth/token"

assert json.loads(exchange_requests[0].content) == {

"grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",

"subject_token": "subject-token",

"subject_token_type": "urn:ietf:params:oauth:token-type:jwt",

"identity_provider_id": "idp_123",

"service_account_id": "sa_123",

}

assert all(isinstance(request, httpx2.Request) for request in api_requests)

assert [request.headers["authorization"] for request in api_requests] == ["Bearer access-token"] * 2

async def test_async_httpx2_workload_401_reexchanges_with_sync_native_client(monkeypatch: pytest.MonkeyPatch) -> None:

exchange_requests: list[Any] = []

api_requests: list[Any] = []

api_authorizations: list[str] = []

tokens = iter(["access-token-1", "access-token-2"])

def exchange_handler(request: Any) -> Any:

exchange_requests.append(request)

return httpx2.Response(200, request=request, json=exchange_payload(next(tokens)))

def exchange_client(**kwargs: Any) -> Any:

assert kwargs == {"follow_redirects": False}

return httpx2.Client(transport=httpx2.MockTransport(exchange_handler), trust_env=False)

async def api_handler(request: Any) -> Any:

api_requests.append(request)

api_authorizations.append(request.headers["authorization"])

status_code = 401 if len(api_requests) == 1 else 200

return httpx2.Response(status_code, request=request, json={"object": "list", "data": []})

monkeypatch.setattr(workload, "DefaultHttpx2Client", exchange_client)

async with AsyncOpenAI(

workload_identity=workload_identity(),

base_url="https://api.example.test/v1",

http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(api_handler), trust_env=False),

max_retries=0,

) as client:

assert (await client.models.list()).object == "list"

assert len(exchange_requests) == 2

assert all(isinstance(request, httpx2.Request) for request in exchange_requests)

assert len(api_requests) == 2

assert all(isinstance(request, httpx2.Request) for request in api_requests)

assert api_authorizations == ["Bearer access-token-1", "Bearer access-token-2"]

def test_sync_httpx2_default_workload_exchange_is_native(monkeypatch: pytest.MonkeyPatch) -> None:

exchange_requests: list[Any] = []

api_requests: list[Any] = []

def exchange_handler(request: Any) -> Any:

exchange_requests.append(request)

return httpx2.Response(200, request=request, json=exchange_payload("access-token"))

def exchange_client(**kwargs: Any) -> Any:

assert kwargs == {"follow_redirects": False}

return httpx2.Client(transport=httpx2.MockTransport(exchange_handler), trust_env=False)

def api_handler(request: Any) -> Any:

api_requests.append(request)

return httpx2.Response(200, request=request, json={"object": "list", "data": []})

def default_client(**kwargs: Any) -> Any:

return DefaultHttpx2Client(transport=httpx2.MockTransport(api_handler), trust_env=False, **kwargs)

monkeypatch.setattr(base_client, "SyncHttpxClientWrapper", default_client)

monkeypatch.setattr(workload, "DefaultHttpx2Client", exchange_client)

with OpenAI(

workload_identity=workload_identity(),

base_url="https://api.example.test/v1",

max_retries=0,

) as client:

assert isinstance(client._client, httpx2.Client)

assert client.models.list().object == "list"

assert len(exchange_requests) == 1

assert len(api_requests) == 1

assert isinstance(exchange_requests[0], httpx2.Request)

assert isinstance(api_requests[0], httpx2.Request)

assert api_requests[0].headers["authorization"] == "Bearer access-token"

async def test_async_httpx2_default_workload_exchange_is_native(monkeypatch: pytest.MonkeyPatch) -> None:

exchange_requests: list[Any] = []

api_requests: list[Any] = []

def exchange_handler(request: Any) -> Any:

exchange_requests.append(request)

return httpx2.Response(200, request=request, json=exchange_payload("access-token"))

def exchange_client(**kwargs: Any) -> Any:

assert kwargs == {"follow_redirects": False}

return httpx2.Client(transport=httpx2.MockTransport(exchange_handler), trust_env=False)

async def api_handler(request: Any) -> Any:

api_requests.append(request)

return httpx2.Response(200, request=request, json={"object": "list", "data": []})

def default_client(**kwargs: Any) -> Any:

return DefaultAsyncHttpx2Client(transport=httpx2.MockTransport(api_handler), trust_env=False, **kwargs)

monkeypatch.setattr(base_client, "AsyncHttpxClientWrapper", default_client)

monkeypatch.setattr(workload, "DefaultHttpx2Client", exchange_client)

async with AsyncOpenAI(

workload_identity=workload_identity(),

base_url="https://api.example.test/v1",

max_retries=0,

) as client:

assert isinstance(client._client, httpx2.AsyncClient)

assert (await client.models.list()).object == "list"

assert len(exchange_requests) == 1

assert len(api_requests) == 1

assert isinstance(exchange_requests[0], httpx2.Request)

assert isinstance(api_requests[0], httpx2.Request)

assert api_requests[0].headers["authorization"] == "Bearer access-token"

def test_httpx2_workload_oauth_error_preserves_native_response(monkeypatch: pytest.MonkeyPatch) -> None:

api_calls = 0

def exchange_handler(request: Any) -> Any:

return httpx2.Response(

401,

request=request,

json={"error": "invalid_grant", "error_description": "invalid workload identity"},

)

def exchange_client(**kwargs: Any) -> Any:

assert kwargs == {"follow_redirects": False}

return httpx2.Client(transport=httpx2.MockTransport(exchange_handler), trust_env=False)

def api_handler(request: Any) -> Any:

nonlocal api_calls

api_calls += 1

return httpx2.Response(200, request=request, json={"object": "list", "data": []})

monkeypatch.setattr(workload, "DefaultHttpx2Client", exchange_client)

with OpenAI(

workload_identity=workload_identity(),

base_url="https://api.example.test/v1",

http_client=httpx2.Client(transport=httpx2.MockTransport(api_handler), trust_env=False),

max_retries=0,

) as client:

with pytest.raises(OAuthError) as exc_info:

client.models.list()

assert exc_info.value.message == "invalid workload identity"

assert isinstance(exc_info.value.response, httpx2.Response)

assert isinstance(exc_info.value.request, httpx2.Request)

assert api_calls == 0

@pytest.mark.parametrize(

("failure", "expected"),

[("timeout", APITimeoutError), ("connection", APIConnectionError)],

)

def test_httpx2_workload_exchange_transport_failure(

failure: str, expected: type[Exception], monkeypatch: pytest.MonkeyPatch

) -> None:

def exchange_handler(request: Any) -> Any:

if failure == "timeout":

raise httpx2.ReadTimeout("exchange timeout", request=request)

raise httpx2.ConnectError("exchange unavailable", request=request)

def exchange_client(**kwargs: Any) -> Any:

assert kwargs == {"follow_redirects": False}

return httpx2.Client(transport=httpx2.MockTransport(exchange_handler), trust_env=False)

def api_handler(request: Any) -> Any:

return httpx2.Response(200, request=request)

monkeypatch.setattr(workload, "DefaultHttpx2Client", exchange_client)

with OpenAI(

workload_identity=workload_identity(),

base_url="https://api.example.test/v1",

http_client=httpx2.Client(transport=httpx2.MockTransport(api_handler)),

max_retries=0,

) as client:

with pytest.raises(expected) as exc_info:

client.models.list()

assert type(exc_info.value.__cause__).__module__ == "httpx2"