From 3fba10159ca9aae6d0e8be713041ba6f7dc66679 Mon Sep 17 00:00:00 2001 From: badtst Date: Mon, 31 Aug 2026 15:34:30 +0000 Subject: [PATCH 1/5] feat(auth): add QPUAuthConfig for Keycloak client credentials feat(auth): add Keycloak client_credentials httpx auth flow --- tests/lib/qpu_client/test_auth.py | 186 ++++++++++++++++++++++++++++++ tests/test_config.py | 47 ++++++++ warden/lib/config/config.py | 30 +++++ warden/lib/qpu_client/auth.py | 129 +++++++++++++++++++++ 4 files changed, 392 insertions(+) create mode 100644 tests/lib/qpu_client/test_auth.py create mode 100644 warden/lib/qpu_client/auth.py diff --git a/tests/lib/qpu_client/test_auth.py b/tests/lib/qpu_client/test_auth.py new file mode 100644 index 0000000..8ab5920 --- /dev/null +++ b/tests/lib/qpu_client/test_auth.py @@ -0,0 +1,186 @@ +"""Testing lib/qpu_client/auth""" + +import httpx +import pytest +from pytest_httpx import HTTPXMock + +from warden.lib.config.config import QPUAuthConfig +from warden.lib.qpu_client.auth import ( + KeycloakClientCredentialsAuth, + TokenRequestError, +) + +TOKEN_URL = "http://keycloak:8080/realms/pasqos/protocol/openid-connect/token" +QPU_URL = "http://qpu:4300/api/v1/system" + + +@pytest.fixture +def auth_conf() -> QPUAuthConfig: + return QPUAuthConfig( + url="http://keycloak:8080", realm="pasqos", id="warden", secret="s3cret" + ) + + +def test_token_is_fetched_once_and_reused(httpx_mock: HTTPXMock, auth_conf): + httpx_mock.add_response( + url=TOKEN_URL, + json={"access_token": "tok-1", "expires_in": 300}, + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + first = client.get(QPU_URL) + second = client.get(QPU_URL) + + assert first.request.headers["Authorization"] == "Bearer tok-1" + assert second.request.headers["Authorization"] == "Bearer tok-1" + token_requests = [r for r in httpx_mock.get_requests() if str(r.url) == TOKEN_URL] + assert len(token_requests) == 1 + + +def test_token_request_uses_client_credentials_grant(httpx_mock: HTTPXMock, auth_conf): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-1", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + client.get(QPU_URL) + + token_request = next( + r for r in httpx_mock.get_requests() if str(r.url) == TOKEN_URL + ) + body = token_request.read().decode() + assert "grant_type=client_credentials" in body + assert "client_id=warden" in body + assert "client_secret=s3cret" in body + + +def test_expired_token_is_refreshed(httpx_mock: HTTPXMock, auth_conf, monkeypatch): + # expires_in 300 with leeway 30 means the token is stale after 270s. + clock = {"now": 1_000.0} + monkeypatch.setattr("warden.lib.qpu_client.auth.monotonic", lambda: clock["now"]) + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-1", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-2", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + first = client.get(QPU_URL) + clock["now"] += 271 + second = client.get(QPU_URL) + + assert first.request.headers["Authorization"] == "Bearer tok-1" + assert second.request.headers["Authorization"] == "Bearer tok-2" + + +def test_401_triggers_one_refresh_and_one_retry(httpx_mock: HTTPXMock, auth_conf): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "stale", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, status_code=401) + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "fresh", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + response = client.get(QPU_URL) + + assert response.status_code == 200 + assert response.request.headers["Authorization"] == "Bearer fresh" + qpu_requests = [r for r in httpx_mock.get_requests() if str(r.url) == QPU_URL] + assert len(qpu_requests) == 2 + + +def test_persistent_401_is_not_retried_forever(httpx_mock: HTTPXMock, auth_conf): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-1", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, status_code=401) + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-2", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, status_code=401) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + response = client.get(QPU_URL) + + # The second 401 is surfaced, not retried again. + assert response.status_code == 401 + qpu_requests = [r for r in httpx_mock.get_requests() if str(r.url) == QPU_URL] + assert len(qpu_requests) == 2 + + +@pytest.mark.parametrize("status_code", [400, 401]) +def test_bad_credentials_raise_token_request_error( + httpx_mock: HTTPXMock, auth_conf, status_code +): + httpx_mock.add_response( + url=TOKEN_URL, + status_code=status_code, + json={"error": "invalid_client"}, + ) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + with pytest.raises(TokenRequestError, match="invalid_client"): + client.get(QPU_URL) + + +def test_keycloak_5xx_raises_retryable_http_status_error( + httpx_mock: HTTPXMock, auth_conf +): + # 503 must stay an httpx.HTTPStatusError so the existing retry decorator + # recognises it as transient. + httpx_mock.add_response(url=TOKEN_URL, status_code=503) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + with pytest.raises(httpx.HTTPStatusError): + client.get(QPU_URL) + + +@pytest.mark.asyncio +async def test_async_flow_attaches_token(httpx_mock: HTTPXMock, auth_conf): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-async", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + async with httpx.AsyncClient(auth=auth) as client: + response = await client.get(QPU_URL) + + assert response.request.headers["Authorization"] == "Bearer tok-async" + + +@pytest.mark.asyncio +async def test_async_flow_reuses_token_cached_by_sync_flow( + httpx_mock: HTTPXMock, auth_conf +): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "shared", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with httpx.Client(auth=auth) as client: + client.get(QPU_URL) + async with httpx.AsyncClient(auth=auth) as client: + response = await client.get(QPU_URL) + + assert response.request.headers["Authorization"] == "Bearer shared" + token_requests = [r for r in httpx_mock.get_requests() if str(r.url) == TOKEN_URL] + assert len(token_requests) == 1 diff --git a/tests/test_config.py b/tests/test_config.py index 6c111da..8b41eef 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -8,6 +8,7 @@ from warden.lib.config.config import ( APIConfig, Config, + QPUAuthConfig, SchedulerConfig, SchedulerStrategy, SqliteConfig, @@ -150,3 +151,49 @@ def test_admin_users_must_not_be_empty(): """ with pytest.raises(ValidationError): APIConfig(admin_users=[]) + + +def test_qpu_auth_absent_by_default(): + assert Config().qpu.auth is None + + +def test_qpu_auth_token_url_is_built_from_base_and_realm(): + auth = QPUAuthConfig( + url="http://keycloak:8080", realm="pasqos", id="warden", secret="s" + ) + + assert ( + auth.token_url + == "http://keycloak:8080/realms/pasqos/protocol/openid-connect/token" + ) + + +def test_qpu_auth_token_url_tolerates_trailing_slash(): + auth = QPUAuthConfig( + url="http://keycloak:8080/", realm="pasqos", id="warden", secret="s" + ) + + assert ( + auth.token_url + == "http://keycloak:8080/realms/pasqos/protocol/openid-connect/token" + ) + + +def test_qpu_auth_rejects_partial_configuration(): + # A half-configured auth section must fail loudly rather than silently + # falling back to unauthenticated requests. + with pytest.raises(ValidationError): + QPUAuthConfig(url="http://keycloak:8080", id="warden") + + +def test_qpu_auth_secret_read_from_env(monkeypatch, tmp_path): + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("WARDEN_QPU_AUTH_URL", "http://keycloak:8080") + monkeypatch.setenv("WARDEN_QPU_AUTH_ID", "warden") + monkeypatch.setenv("WARDEN_QPU_AUTH_SECRET", "from-env") + + config = Config() + + assert config.qpu.auth is not None + assert config.qpu.auth.id == "warden" + assert config.qpu.auth.secret == "from-env" diff --git a/warden/lib/config/config.py b/warden/lib/config/config.py index fa8ce2f..70ea2c0 100644 --- a/warden/lib/config/config.py +++ b/warden/lib/config/config.py @@ -181,6 +181,34 @@ class SchedulerConfig(WardenSettings): ) +class QPUAuthConfig(WardenSettings): + """Keycloak client_credentials configuration for outbound QPU API calls. + + Presence of this section is what enables authentication. There is + deliberately no separate ``enabled`` flag: a second switch can drift out of + sync with the credentials it guards. ``url``, ``id`` and ``secret`` have no + defaults, so a partially configured section is a startup validation error + rather than a silent fallback to unauthenticated requests. + """ + + url: str = Field(description="Keycloak base URL, e.g. http://keycloak:8080") + + realm: str = Field(default="pasqos") + + id: str = Field(description="OIDC client_id") + + secret: str = Field(description="OIDC client_secret. Provide via WARDEN_QPU_AUTH_SECRET, never in YAML.") + + leeway_s: float = Field(default=30, description="Refresh this many seconds before the token actually expires.") + + @property + def token_url(self) -> str: + """Keycloak's OIDC token endpoint for this realm.""" + return ( + f"{self.url.rstrip('/')}/realms/{self.realm}/protocol/openid-connect/token" + ) + + class QPUConfig(WardenSettings): """QPU backend connection configuration.""" @@ -188,6 +216,8 @@ class QPUConfig(WardenSettings): default="http://localhost:8000", description="Local Pasqal QPU API URI." ) + auth: QPUAuthConfig | None = None + retry_max: int = Field( default=10, description=( diff --git a/warden/lib/qpu_client/auth.py b/warden/lib/qpu_client/auth.py new file mode 100644 index 0000000..23b239b --- /dev/null +++ b/warden/lib/qpu_client/auth.py @@ -0,0 +1,129 @@ +"""Keycloak client_credentials authentication for outbound QPU API calls.""" + +import logging +import ssl +from time import monotonic +from typing import AsyncGenerator, Generator + +import httpx + +from warden.lib.config.config import QPUAuthConfig +from warden.lib.qpu_client.retry import QPUClientRequestError + +logger = logging.getLogger(__name__) + +# Token-endpoint statuses that will never succeed on retry: the credentials or +# the grant itself are wrong. Anything else (transport errors, 5xx) is left to +# propagate so the existing retry decorator can treat it as transient. +FATAL_TOKEN_STATUSES = (400, 401, 403) + + +class TokenRequestError(QPUClientRequestError): + """Keycloak refused to issue a token and retrying cannot help.""" + + +class KeycloakClientCredentialsAuth(httpx.Auth): + """Attach a Keycloak service-account bearer token to each request. + + Implemented as an ``httpx.Auth`` so it runs inside the transport, below + Warden's ``retry`` decorator. That matters because 401 is not in + ``RETRY_HTTP_EXIT_CODES``: a token expiring mid-job would otherwise surface + as an immediate, non-retryable ``NotRetriedHTTPStatus``. Here it is just a + refresh. + + Args: + conf: Keycloak credentials and endpoint. + verify: httpx TLS verification setting for the token request. + """ + + def __init__( + self, + conf: QPUAuthConfig, + verify: bool | str | ssl.SSLContext = True, + ) -> None: + self.conf = conf + self.verify = verify + self._token: str | None = None + # monotonic() deadline after which the cached token is considered stale. + self._expires_at: float = 0.0 + + def sync_auth_flow( + self, request: httpx.Request + ) -> Generator[httpx.Request, httpx.Response, None]: + request.headers["Authorization"] = f"Bearer {self._sync_token()}" + response = yield request + if response.status_code == httpx.codes.UNAUTHORIZED: + logger.info("QPU API returned 401, refreshing token and retrying once") + request.headers["Authorization"] = f"Bearer {self._sync_token(force=True)}" + yield request + + async def async_auth_flow( + self, request: httpx.Request + ) -> AsyncGenerator[httpx.Request, httpx.Response]: + request.headers["Authorization"] = f"Bearer {await self._async_token()}" + response = yield request + if response.status_code == httpx.codes.UNAUTHORIZED: + logger.info("QPU API returned 401, refreshing token and retrying once") + request.headers["Authorization"] = ( + f"Bearer {await self._async_token(force=True)}" + ) + yield request + + # ponytail: unlocked cache. Two concurrent requests in one process can both + # miss and both fetch a token; one wins and the loser wasted a request. Add + # a lock only if token-endpoint traffic ever becomes a problem. + def _is_fresh(self) -> bool: + return self._token is not None and monotonic() < self._expires_at + + def _token_request(self) -> tuple[str, dict[str, str]]: + """Return the (url, form data) for a client_credentials token request.""" + return self.conf.token_url, { + "grant_type": "client_credentials", + "client_id": self.conf.id, + "client_secret": self.conf.secret, + } + + def _store(self, response: httpx.Response) -> str: + """Validate a token response, cache the token and return it.""" + if response.status_code in FATAL_TOKEN_STATUSES: + # Never log the response body of a token request: it may echo + # credentials. The error field alone is the useful part. + try: + error = response.json().get("error", "unknown_error") + except ValueError: + error = "unknown_error" + raise TokenRequestError( + f"Keycloak refused to issue a token for client " + f"'{self.conf.id}' at {self.conf.token_url}: " + f"{response.status_code} {error}" + ) + # Transport errors and 5xx stay as httpx exceptions so the existing + # retry decorator sees them as transient. + response.raise_for_status() + + payload = response.json() + token = payload["access_token"] + expires_in = float(payload.get("expires_in", 0)) + self._token = token + self._expires_at = monotonic() + max(expires_in - self.conf.leeway_s, 0.0) + logger.debug( + f"Obtained QPU API token for client '{self.conf.id}', " + f"expires in {expires_in}s" + ) + return token + + def _sync_token(self, force: bool = False) -> str: + if not force and self._is_fresh(): + assert self._token is not None + return self._token + url, data = self._token_request() + with httpx.Client(verify=self.verify) as client: + return self._store(client.post(url, data=data)) + + async def _async_token(self, force: bool = False) -> str: + if not force and self._is_fresh(): + assert self._token is not None + return self._token + url, data = self._token_request() + async with httpx.AsyncClient(verify=self.verify) as client: + return self._store(await client.post(url, data=data)) From 2932533e9179e14ba10c61f24e356ec077f756e1 Mon Sep 17 00:00:00 2001 From: badtst Date: Mon, 31 Aug 2026 15:35:23 +0000 Subject: [PATCH 2/5] # This is a combination of 6 commits. # This is the 1st commit message: feat(auth): add QPUAuthConfig for Keycloak client credentials feat(auth): add Keycloak client_credentials httpx auth flow # This is the commit message #2: feat(auth): add QPUAuthConfig for Keycloak client credentials feat(auth): add Keycloak client_credentials httpx auth flow feat(auth): attach Keycloak auth flow to QPU API clients fix(retry): do not rewrap already-classified QPU client errors docs(auth): document qpu.auth config block # This is the commit message #3: feat(auth): add QPUAuthConfig for Keycloak client credentials feat(auth): add Keycloak client_credentials httpx auth flow feat(auth): attach Keycloak auth flow to QPU API clients fix(retry): do not rewrap already-classified QPU client errors docs(auth): document qpu.auth config block fix(auth): handle short token lifespans and address review findings Cache for half the lifespan (with a WARNING) instead of disabling the cache entirely when expires_in <= leeway_s. Also catch AttributeError alongside ValueError when parsing a non-object token error body, narrow the unlocked-cache ponytail comment to the async path where the race is actually reachable, tighten config.py's auth_flow typing, note that QPUClientRequestError subclasses can propagate through retry(), and add coverage for the no-auth-header and POST-401-replay paths. fix(typechecking): Fix type-checking error # This is the commit message #4: feat(auth): add QPUAuthConfig for Keycloak client credentials feat(auth): add Keycloak client_credentials httpx auth flow feat(auth): attach Keycloak auth flow to QPU API clients fix(retry): do not rewrap already-classified QPU client errors docs(auth): document qpu.auth config block fix(auth): handle short token lifespans and address review findings Cache for half the lifespan (with a WARNING) instead of disabling the cache entirely when expires_in <= leeway_s. Also catch AttributeError alongside ValueError when parsing a non-object token error body, narrow the unlocked-cache ponytail comment to the async path where the race is actually reachable, tighten config.py's auth_flow typing, note that QPUClientRequestError subclasses can propagate through retry(), and add coverage for the no-auth-header and POST-401-replay paths. fix(typechecking): Fix type-checking error fix(test): Fix test flakiness feat(auth): add QPUAuthConfig for Keycloak client credentials feat(auth): add Keycloak client_credentials httpx auth flow feat(auth): attach Keycloak auth flow to QPU API clients fix(retry): do not rewrap already-classified QPU client errors docs(auth): document qpu.auth config block fix(auth): handle short token lifespans and address review findings Cache for half the lifespan (with a WARNING) instead of disabling the cache entirely when expires_in <= leeway_s. Also catch AttributeError alongside ValueError when parsing a non-object token error body, narrow the unlocked-cache ponytail comment to the async path where the race is actually reachable, tighten config.py's auth_flow typing, note that QPUClientRequestError subclasses can propagate through retry(), and add coverage for the no-auth-header and POST-401-replay paths. fix(typechecking): Fix type-checking error fix(test): Fix test flakiness # This is the commit message #5: Fix test and refacto # This is the commit message #6: PR comments --- tests/lib/qpu_client/test_auth.py | 89 ++++++-- tests/lib/qpu_client/test_retry.py | 32 +++ tests/scheduler/test_scheduler.py | 213 +++++++++--------- tests/scheduler/test_scheduler_integration.py | 17 +- tests/scheduler/utils.py | 32 +++ tests/test_config.py | 31 ++- warden/lib/config/config.py | 21 +- warden/lib/qpu_client/auth.py | 42 ++-- warden/lib/qpu_client/retry.py | 6 + warden/scheduler/worker.py | 48 ++-- 10 files changed, 367 insertions(+), 164 deletions(-) create mode 100644 tests/lib/qpu_client/test_retry.py diff --git a/tests/lib/qpu_client/test_auth.py b/tests/lib/qpu_client/test_auth.py index 8ab5920..c829623 100644 --- a/tests/lib/qpu_client/test_auth.py +++ b/tests/lib/qpu_client/test_auth.py @@ -1,10 +1,13 @@ """Testing lib/qpu_client/auth""" -import httpx +import json +import logging + import pytest -from pytest_httpx import HTTPXMock +from httpx2 import AsyncClient, Client, HTTPStatusError +from pytest_httpx2 import HTTPXMock -from warden.lib.config.config import QPUAuthConfig +from warden.lib.config.config import QPUAuthConfig, QPUConfig from warden.lib.qpu_client.auth import ( KeycloakClientCredentialsAuth, TokenRequestError, @@ -30,7 +33,7 @@ def test_token_is_fetched_once_and_reused(httpx_mock: HTTPXMock, auth_conf): httpx_mock.add_response(url=QPU_URL, json={"data": {}}) auth = KeycloakClientCredentialsAuth(auth_conf) - with httpx.Client(auth=auth) as client: + with Client(auth=auth) as client: first = client.get(QPU_URL) second = client.get(QPU_URL) @@ -47,7 +50,7 @@ def test_token_request_uses_client_credentials_grant(httpx_mock: HTTPXMock, auth httpx_mock.add_response(url=QPU_URL, json={"data": {}}) auth = KeycloakClientCredentialsAuth(auth_conf) - with httpx.Client(auth=auth) as client: + with Client(auth=auth) as client: client.get(QPU_URL) token_request = next( @@ -73,7 +76,7 @@ def test_expired_token_is_refreshed(httpx_mock: HTTPXMock, auth_conf, monkeypatc httpx_mock.add_response(url=QPU_URL, json={"data": {}}) auth = KeycloakClientCredentialsAuth(auth_conf) - with httpx.Client(auth=auth) as client: + with Client(auth=auth) as client: first = client.get(QPU_URL) clock["now"] += 271 second = client.get(QPU_URL) @@ -93,7 +96,7 @@ def test_401_triggers_one_refresh_and_one_retry(httpx_mock: HTTPXMock, auth_conf httpx_mock.add_response(url=QPU_URL, json={"data": {}}) auth = KeycloakClientCredentialsAuth(auth_conf) - with httpx.Client(auth=auth) as client: + with Client(auth=auth) as client: response = client.get(QPU_URL) assert response.status_code == 200 @@ -113,7 +116,7 @@ def test_persistent_401_is_not_retried_forever(httpx_mock: HTTPXMock, auth_conf) httpx_mock.add_response(url=QPU_URL, status_code=401) auth = KeycloakClientCredentialsAuth(auth_conf) - with httpx.Client(auth=auth) as client: + with Client(auth=auth) as client: response = client.get(QPU_URL) # The second 401 is surfaced, not retried again. @@ -133,7 +136,7 @@ def test_bad_credentials_raise_token_request_error( ) auth = KeycloakClientCredentialsAuth(auth_conf) - with httpx.Client(auth=auth) as client: + with Client(auth=auth) as client: with pytest.raises(TokenRequestError, match="invalid_client"): client.get(QPU_URL) @@ -146,8 +149,8 @@ def test_keycloak_5xx_raises_retryable_http_status_error( httpx_mock.add_response(url=TOKEN_URL, status_code=503) auth = KeycloakClientCredentialsAuth(auth_conf) - with httpx.Client(auth=auth) as client: - with pytest.raises(httpx.HTTPStatusError): + with Client(auth=auth) as client: + with pytest.raises(HTTPStatusError): client.get(QPU_URL) @@ -159,7 +162,7 @@ async def test_async_flow_attaches_token(httpx_mock: HTTPXMock, auth_conf): httpx_mock.add_response(url=QPU_URL, json={"data": {}}) auth = KeycloakClientCredentialsAuth(auth_conf) - async with httpx.AsyncClient(auth=auth) as client: + async with AsyncClient(auth=auth) as client: response = await client.get(QPU_URL) assert response.request.headers["Authorization"] == "Bearer tok-async" @@ -176,11 +179,69 @@ async def test_async_flow_reuses_token_cached_by_sync_flow( httpx_mock.add_response(url=QPU_URL, json={"data": {}}) auth = KeycloakClientCredentialsAuth(auth_conf) - with httpx.Client(auth=auth) as client: + with Client(auth=auth) as client: client.get(QPU_URL) - async with httpx.AsyncClient(auth=auth) as client: + async with AsyncClient(auth=auth) as client: response = await client.get(QPU_URL) assert response.request.headers["Authorization"] == "Bearer shared" token_requests = [r for r in httpx_mock.get_requests() if str(r.url) == TOKEN_URL] assert len(token_requests) == 1 + + +def test_short_lived_token_still_caches_with_warning( + httpx_mock: HTTPXMock, auth_conf, caplog +): + # expires_in 30 with the default leeway_s 30 would clamp to 0 without the + # half-lifespan fallback, disabling the cache entirely and forcing a + # Keycloak round-trip on every request. + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "tok-1", "expires_in": 30} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + auth = KeycloakClientCredentialsAuth(auth_conf) + with caplog.at_level(logging.WARNING, logger="warden.lib.qpu_client.auth"): + with Client(auth=auth) as client: + client.get(QPU_URL) + client.get(QPU_URL) + + token_requests = [r for r in httpx_mock.get_requests() if str(r.url) == TOKEN_URL] + assert len(token_requests) == 1 + assert any(record.levelno == logging.WARNING for record in caplog.records) + + +def test_client_sends_no_authorization_header_without_auth_config( + httpx_mock: HTTPXMock, +): + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + response = QPUConfig(uri="http://qpu:4300").client.get(QPU_URL) + + assert "Authorization" not in response.request.headers + + +def test_401_on_post_retries_with_fresh_token_and_identical_body( + httpx_mock: HTTPXMock, auth_conf +): + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "stale", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, status_code=401) + httpx_mock.add_response( + url=TOKEN_URL, json={"access_token": "fresh", "expires_in": 300} + ) + httpx_mock.add_response(url=QPU_URL, json={"data": {}}) + + body = {"circuit": "bell", "shots": 100} + auth = KeycloakClientCredentialsAuth(auth_conf) + with Client(auth=auth) as client: + response = client.post(QPU_URL, json=body) + + assert response.status_code == 200 + assert response.request.headers["Authorization"] == "Bearer fresh" + qpu_requests = [r for r in httpx_mock.get_requests() if str(r.url) == QPU_URL] + assert len(qpu_requests) == 2 + for request in qpu_requests: + assert json.loads(request.read()) == body diff --git a/tests/lib/qpu_client/test_retry.py b/tests/lib/qpu_client/test_retry.py new file mode 100644 index 0000000..2ebb342 --- /dev/null +++ b/tests/lib/qpu_client/test_retry.py @@ -0,0 +1,32 @@ +"""Testing lib/qpu_client/retry""" + +import pytest + +from warden.lib.qpu_client.auth import TokenRequestError +from warden.lib.qpu_client.retry import UnhandledError, retry + + +@pytest.mark.asyncio +async def test_already_classified_errors_are_not_rewrapped(): + calls = {"n": 0} + + @retry(max=5, sleep_s=0) + async def fails_with_bad_credentials(): + calls["n"] += 1 + raise TokenRequestError("invalid_client") + + with pytest.raises(TokenRequestError): + await fails_with_bad_credentials() + + # Fail fast: a wrong secret will not fix itself. + assert calls["n"] == 1 + + +@pytest.mark.asyncio +async def test_unknown_errors_are_still_wrapped(): + @retry(max=5, sleep_s=0) + async def fails_with_value_error(): + raise ValueError("something unexpected") + + with pytest.raises(UnhandledError): + await fails_with_value_error() diff --git a/tests/scheduler/test_scheduler.py b/tests/scheduler/test_scheduler.py index 2705ccf..8e4d358 100644 --- a/tests/scheduler/test_scheduler.py +++ b/tests/scheduler/test_scheduler.py @@ -18,7 +18,8 @@ from warden.lib.config import Config, SchedulerStrategy from warden.lib.models import Job from warden.scheduler.main import run_scheduler -from warden.scheduler.worker import LocalQPUWorker +from warden.scheduler.types import JobUpdateQueue +from warden.scheduler.worker import TERMINAL_STATUSES, LocalQPUWorker NOW = datetime.now() @@ -35,8 +36,6 @@ SYSTEM_API = API_URI + "/system" PROGRAM_API = API_URI + "/programs" -SUCCESS_CHECK_INTERVAL_S = 0.1 - DUMMY_RESULTS = json.dumps([{"counter": {"0001": 1, "0010": 2, "0100": 3, "1000": 4}}]) @@ -60,7 +59,7 @@ async def test_run_nominal( - To return "RUNNING" and then "DONE" status for each job - Run scheduler until: - All jobs have a "DONE" status is DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Check n (jobs with status "DONE") = N_JOBS - Check "DONE" jobs have the right results and non-empty logs - Check those jobs have `scheduled_at` set @@ -73,7 +72,6 @@ async def test_run_nominal( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 10 conf: Config = build_conf(strategy, QPU_URI) @@ -147,8 +145,6 @@ async def test_run_nominal( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status == "DONE") - ################## ### TEST RUN ### ################## @@ -157,13 +153,7 @@ async def test_run_nominal( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_count, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled(session, main_task, count=N_JOBS) stmt_all = select(Job).where(Job.status == "DONE") jobs_done = (await session.execute(stmt_all)).scalars().all() @@ -200,7 +190,7 @@ async def test_run_resume_job( - To return "DONE" status for ALREADY_DONE_BACKEND_ID - Run scheduler until: - All jobs have a "DONE" status in DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Check n (jobs with status "DONE") = 3 - Check "DONE" jobs have the right results and non-empty logs - Check those jobs have `scheduled_at` set @@ -213,7 +203,6 @@ async def test_run_resume_job( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 5 NORMAL_BACKEND_ID = "1" NON_EXISTING_BACKEND_ID = "9999" NEW_BACKEND_ID = "2" @@ -330,8 +319,6 @@ async def test_run_resume_job( ], ) - stmt_count = select(func.count(Job.id)).where(Job.status == "DONE") - ################## ### TEST RUN ### ################## @@ -340,13 +327,7 @@ async def test_run_resume_job( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_count, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled(session, main_task, count=N_JOBS) stmt_all = select(Job).where(Job.status == "DONE") jobs_done = (await session.execute(stmt_all)).scalars().all() @@ -377,7 +358,7 @@ async def test_run_qpu_down( - No need to mock jobs calls - Run scheduler until: - All jobs have an "ERROR" status - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Check n (jobs with status "ERROR") = N_JOBS - Check those jobs have non-empty logs """ @@ -389,7 +370,6 @@ async def test_run_qpu_down( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 3 EXPECTED_STATUS = "ERROR" @@ -413,8 +393,6 @@ async def test_run_qpu_down( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt = select(func.count(Job.id)).where(Job.status == EXPECTED_STATUS) - ################## ### TEST RUN ### ################## @@ -423,13 +401,9 @@ async def test_run_qpu_down( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=(EXPECTED_STATUS,) + ) stmt_all = select(Job).where(Job.status == EXPECTED_STATUS) all_jobs = (await session.execute(stmt_all)).scalars().all() @@ -497,7 +471,6 @@ async def test_run_job_timeout( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 8 N_JOBS_TIMEOUT = 4 @@ -650,13 +623,9 @@ async def test_run_job_timeout( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_processed, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=("DONE", "CANCELED") + ) n_processed = (await session.execute(stmt_processed)).scalar() assert n_processed == N_JOBS @@ -693,7 +662,7 @@ async def test_run_resume_job_timeout( - Accept the job's cancelation request - Run scheduler until: - All jobs have a "DONE" status in DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Check n (jobs with status "CANCELED") == 1 - Check "CANCELED" jobs have the right results and non-empty logs - Check those jobs have `scheduled_at` set @@ -706,7 +675,6 @@ async def test_run_resume_job_timeout( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 5 N_JOBS = 1 BACKEND_ID = "1" # Setting the job's created_at at a time that is already timedout @@ -776,8 +744,6 @@ async def test_run_resume_job_timeout( backend_ids=[BACKEND_ID], ) - stmt_count = select(func.count(Job.id)).where(Job.status == EXPECTED_JOB_STATUS) - ################## ### TEST RUN ### ################## @@ -786,13 +752,9 @@ async def test_run_resume_job_timeout( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_count, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=(EXPECTED_JOB_STATUS,) + ) stmt_all = select(Job).where(Job.status == EXPECTED_JOB_STATUS) jobs_done = (await session.execute(stmt_all)).scalars().all() @@ -828,7 +790,7 @@ async def test_run_retry_transient_errors( - To return "RUNNING" and then "DONE" status for each job - Run scheduler until: - All jobs have a "DONE" status in DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Checks: - n (jobs with status "DONE") = N_JOBS - jobs have non-empty logs @@ -841,7 +803,6 @@ async def test_run_retry_transient_errors( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 1 conf: Config = build_conf(strategy, QPU_URI) @@ -934,7 +895,6 @@ def _add_transient_errors(httpx_mock: HTTPXMock, url: str, method: str): # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status == "DONE") stmt = select(Job).where(Job.status == "DONE") ################## @@ -945,13 +905,7 @@ def _add_transient_errors(httpx_mock: HTTPXMock, url: str, method: str): main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_count, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled(session, main_task, count=N_JOBS) jobs_done = (await session.execute(stmt)).scalars().all() assert len(jobs_done) == N_JOBS @@ -979,7 +933,7 @@ async def test_run_qpu_api_unreachable( - To return QPU status as "Down" - Run scheduler until: - All jobs have a "ERROR" status is DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Checks - n(jobs with status "ERROR") = N_JOBS - All jobs have non-empty logs and an "ERROR" message @@ -992,7 +946,6 @@ async def test_run_qpu_api_unreachable( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 1 EXPECTED_STATUS = "ERROR" @@ -1008,7 +961,6 @@ async def test_run_qpu_api_unreachable( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status == EXPECTED_STATUS) stmt = select(Job).where(Job.status == EXPECTED_STATUS) ################## @@ -1019,13 +971,9 @@ async def test_run_qpu_api_unreachable( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_count, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=(EXPECTED_STATUS,) + ) error_jobs = (await session.execute(stmt)).scalars().all() assert len(error_jobs) == N_JOBS @@ -1055,7 +1003,7 @@ async def test_run_job_creation_client_error( - Return exceptions when attempting to create a job - Run scheduler until: - All jobs have a "ERROR" status is DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Checks - n(jobs with status "ERROR") = N_JOBS - All jobs have non-empty logs and an "ERROR" message @@ -1068,7 +1016,6 @@ async def test_run_job_creation_client_error( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 3 EXPECTED_STATUS = "ERROR" @@ -1097,7 +1044,6 @@ async def test_run_job_creation_client_error( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status == EXPECTED_STATUS) stmt = select(Job).where(Job.status == EXPECTED_STATUS) ################## @@ -1108,13 +1054,9 @@ async def test_run_job_creation_client_error( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_count, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=(EXPECTED_STATUS,) + ) jobs = (await session.execute(stmt)).scalars().all() assert len(jobs) == N_JOBS @@ -1149,7 +1091,7 @@ async def test_run_job_client_error_timeout( (it's the same backend request in QPU) - Run scheduler until: - All jobs have an "ERROR" status is DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Check n (jobs with status "ERROR") = N_JOBS - Check those jobs have `ended_at` set, even though the job was never reported as ended by the QPU: `to_error` must backfill it. @@ -1162,7 +1104,6 @@ async def test_run_job_client_error_timeout( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 1 EXPECTED_STATUS = "ERROR" @@ -1215,7 +1156,6 @@ async def test_run_job_client_error_timeout( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status == EXPECTED_STATUS) stmt = select(Job).where(Job.status == EXPECTED_STATUS) ################## @@ -1226,19 +1166,16 @@ async def test_run_job_client_error_timeout( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_count, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=(EXPECTED_STATUS,) + ) jobs = (await session.execute(stmt)).scalars().all() assert len(jobs) == N_JOBS for job in jobs: assert len(job.logs) > 0 assert "ERROR" in job.logs + assert "Job execution ended with status 'ERROR'" in job.logs assert job.ended_at is not None @@ -1264,7 +1201,7 @@ async def test_run_job_canceled_by_cancellation_worker( - For JOB_ID_CANCELED return "CANCELED" status - Run scheduler until: - All jobs have a "DONE" or "CANCELED" status in DB - - Test timeout after TEST_TIMEOUT_S + - Test timeout after utils.JOB_WAIT_TIMEOUT_S - Check n (jobs with status "DONE") = N_JOBS-1 - Check "DONE" jobs have the right results and non-empty logs - Check n (jobs with status "CANCELED") = 1 @@ -1278,7 +1215,6 @@ async def test_run_job_canceled_by_cancellation_worker( # Enable warden logging for jobs 'logs' field to be populated caplog.set_level(logging.INFO, logger="warden") - TEST_TIMEOUT_S = 3 N_JOBS = 10 JOB_ID_CANCELED = 5 @@ -1384,8 +1320,6 @@ async def test_run_job_canceled_by_cancellation_worker( # Populate DB with jobs to run await utils.create_n_jobs(db_session_maker, N_JOBS) - stmt_count = select(func.count(Job.id)).where(Job.status.in_(("DONE", "CANCELED"))) - ################## ### TEST RUN ### ################## @@ -1394,13 +1328,9 @@ async def test_run_job_canceled_by_cancellation_worker( main_task = asyncio.create_task(run_scheduler(db_engine, conf)) async with db_session_maker() as session: - async with utils.scheduler_task_timeout(TEST_TIMEOUT_S, main_task): - await utils.wait_until_scalar_equals( - session, - stmt_count, - N_JOBS, - interval=SUCCESS_CHECK_INTERVAL_S, - ) + await utils.wait_until_jobs_settled( + session, main_task, count=N_JOBS, statuses=("DONE", "CANCELED") + ) stmt_done = select(Job).where(Job.status == "DONE") jobs_done = (await session.execute(stmt_done)).scalars().all() @@ -1555,3 +1485,74 @@ async def crash(*args, **kwargs): job = (await session.execute(select(Job))).scalar_one() assert job.backend_id == "0" assert job.status == "RUNNING" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("final_status", ["DONE", "ERROR"]) +async def test_terminal_status_and_closing_log_are_committed_together( + final_status: str, + httpx_mock: HTTPXMock, + caplog, +): + """A job's terminal status must be queued together with its closing log line. + + One JobUpdate is one transaction, so pushing the status and the closing + "Job execution ended with status '...'" line separately leaves a window + where the DB holds a finished job whose logs are truncated. Anything that + stops polling once the status is terminal - the API, and every test that + waits on a status then asserts on logs - reads incomplete logs. + + This guards the invariant: keep the closing line in the same update as the + terminal status, whether the job reaches it through `update_job` (DONE) + or through `to_error` (ERROR, here triggered by a failing job creation). + """ + + # Enable warden logging for jobs 'logs' field to be populated + caplog.set_level(logging.INFO, logger="warden") + + def job_json(status: str) -> dict: + return { + "data": { + "uid": 0, + "batch_id": SLURM_USER_ID, + "status": status, + "result": DUMMY_RESULTS if status == "DONE" else None, + "program_id": QPU_PROGRAM_UID, + "created_datetime": NOW.isoformat(), + "start_datetime": (NOW + timedelta(seconds=1)).isoformat(), + "end_datetime": (NOW + timedelta(seconds=2)).isoformat(), + } + } + + httpx_mock.add_response( + method="GET", + url=SYSTEM_OPERATIONAL_API, + json={"data": {"operational_status": "UP"}}, + ) + if final_status == "ERROR": + # Job creation fails outright -> create_job() catches + # QPUClientRequestError and calls to_error() directly. + # is_reusable=True: the QPU client retries on 500s before giving up. + httpx_mock.add_response( + method="POST", status_code=500, url=JOB_API, is_reusable=True + ) + else: + httpx_mock.add_response( + method="POST", status_code=200, url=JOB_API, json=job_json("RUNNING") + ) + httpx_mock.add_response( + method="GET", status_code=200, url=JOB_API + "/0", json=job_json("RUNNING") + ) + httpx_mock.add_response( + method="GET", status_code=200, url=JOB_API + "/0", json=job_json("DONE") + ) + + queue: JobUpdateQueue = JobUpdateQueue() + worker = LocalQPUWorker(conf=build_conf(SchedulerStrategy.FIFO, QPU_URI)) + await worker.execute_job(queue=queue, nb_run=100, sequence="{}") + + updates = [queue.get_nowait() for _ in range(queue.qsize())] + first_terminal = next(u for u in updates if u.status in TERMINAL_STATUSES) + assert ( + f"Job execution ended with status '{final_status}'" in first_terminal.new_logs + ) diff --git a/tests/scheduler/test_scheduler_integration.py b/tests/scheduler/test_scheduler_integration.py index 778b77e..52780c8 100644 --- a/tests/scheduler/test_scheduler_integration.py +++ b/tests/scheduler/test_scheduler_integration.py @@ -80,12 +80,10 @@ async def test_run_scheduler_integration( await utils.create_n_jobs(db_session_maker, N_JOBS) - # The terminal status is committed before the closing "Job execution ended - # with status 'DONE'" log line is flushed, so waiting on the status alone - # races with the log assertions below. Wait for the logs too. - stmt = select(func.count(Job.id)).where( - Job.status == "DONE", Job.logs.contains("DONE") - ) + # Safe to wait on the status alone: the scheduler commits a job's terminal + # status and its complete logs in the same transaction, so the log + # assertions below cannot race it + stmt = select(func.count(Job.id)).where(Job.status == "DONE") ################## ### TEST RUN ### @@ -191,11 +189,8 @@ async def test_run_scheduler_integration_cancellation_worker( JOB_TO_CANCEL_ID = job_to_cancel.id - # Same race as above: wait for the closing log line, not just the status - stmt = select(func.count(Job.id)).where( - Job.status.in_(("CANCELED", "DONE")), - Job.logs.contains("Job execution ended"), - ) + # Status alone is enough here too, see the comment in the test above + stmt = select(func.count(Job.id)).where(Job.status.in_(("CANCELED", "DONE"))) ################## ### TEST RUN ### diff --git a/tests/scheduler/utils.py b/tests/scheduler/utils.py index 255503d..a925c91 100644 --- a/tests/scheduler/utils.py +++ b/tests/scheduler/utils.py @@ -2,15 +2,22 @@ import asyncio from asyncio import Task, timeout +from collections.abc import Sequence from contextlib import asynccontextmanager from typing import Any import pytest +from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from warden.lib.config import Config, QPUConfig, SchedulerConfig, SchedulerStrategy from warden.lib.models import Job, Session +# Deliberately generous: this budget is only ever spent by a test that is +# already failing, so it costs nothing on the happy path. Tight per-test budgets +# turned a slow CI runner into a flake instead of catching anything. +JOB_WAIT_TIMEOUT_S = 30 + async def wait_until_scalar_equals( session: AsyncSession, @@ -123,6 +130,31 @@ async def scheduler_task_timeout(delay: float, scheduler_task: Task): pass +async def wait_until_jobs_settled( + session: AsyncSession, + scheduler_task: Task, + *, + count: int, + statuses: Sequence[str] = ("DONE",), + timeout_s: float = JOB_WAIT_TIMEOUT_S, + interval: float = 0.1, +) -> None: + """Wait until ``count`` jobs reached one of ``statuses``, then stop the scheduler. + + Use this rather than hand-rolling a wait predicate. The scheduler commits a + job's terminal status and its complete logs in a single transaction (see + ``JobExecutionTracker.update_job``), so waiting on the status alone is enough + to make the assertions that follow - including any on ``logs`` - safe. + + Fails the test on timeout, and always cancels ``scheduler_task`` so it cannot + outlive the test body and interfere with fixture teardown. + """ + + stmt = select(func.count(Job.id)).where(Job.status.in_(tuple(statuses))) + async with scheduler_task_timeout(timeout_s, scheduler_task): + await wait_until_scalar_equals(session, stmt, count, interval=interval) + + def build_conf(strategy: SchedulerStrategy, qpu_uri: str) -> Config: return Config( scheduler=SchedulerConfig( diff --git a/tests/test_config.py b/tests/test_config.py index 8b41eef..8275382 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -9,6 +9,7 @@ APIConfig, Config, QPUAuthConfig, + QPUConfig, SchedulerConfig, SchedulerStrategy, SqliteConfig, @@ -183,7 +184,9 @@ def test_qpu_auth_rejects_partial_configuration(): # A half-configured auth section must fail loudly rather than silently # falling back to unauthenticated requests. with pytest.raises(ValidationError): - QPUAuthConfig(url="http://keycloak:8080", id="warden") + # model_validate, not the constructor: omitting a required field is the + # point of the test, and a static type checker rejects the direct call. + QPUAuthConfig.model_validate({"url": "http://keycloak:8080", "id": "warden"}) def test_qpu_auth_secret_read_from_env(monkeypatch, tmp_path): @@ -197,3 +200,29 @@ def test_qpu_auth_secret_read_from_env(monkeypatch, tmp_path): assert config.qpu.auth is not None assert config.qpu.auth.id == "warden" assert config.qpu.auth.secret == "from-env" + + +def test_auth_flow_is_none_without_auth_config(): + assert Config().qpu.auth_flow is None + + +def test_auth_flow_is_memoized(): + qpu = QPUConfig( + uri="http://qpu:4300", + auth=QPUAuthConfig(url="http://keycloak:8080", id="warden", secret="s"), + ) + + assert qpu.auth_flow is qpu.auth_flow + + +def test_client_is_given_the_auth_flow(): + qpu = QPUConfig( + uri="http://qpu:4300", + auth=QPUAuthConfig(url="http://keycloak:8080", id="warden", secret="s"), + ) + + assert qpu.client.auth is qpu.auth_flow + + +def test_client_has_no_auth_without_auth_config(): + assert QPUConfig(uri="http://qpu:4300").client.auth is None diff --git a/warden/lib/config/config.py b/warden/lib/config/config.py index 70ea2c0..6d45ebd 100644 --- a/warden/lib/config/config.py +++ b/warden/lib/config/config.py @@ -241,6 +241,7 @@ class QPUConfig(WardenSettings): ) _client: httpx2.AsyncClient | None = PrivateAttr(default=None) + _auth_flow: httpx2.Auth | None = PrivateAttr(default=None) @property def verify(self) -> bool | ssl.SSLContext: @@ -249,10 +250,28 @@ def verify(self) -> bool | ssl.SSLContext: return ssl.create_default_context(cafile=self.tls_verify) return self.tls_verify + @property + def auth_flow(self) -> httpx2.Auth | None: + """Memoized Keycloak auth flow, or None when auth is not configured. + + Memoized so that every client built from this config shares a single + cached token. Imported lazily because ``qpu_client.auth`` imports this + module. + """ + if self.auth is None: + return None + if self._auth_flow is None: + from warden.lib.qpu_client.auth import KeycloakClientCredentialsAuth + + self._auth_flow = KeycloakClientCredentialsAuth( + self.auth, verify=self.verify + ) + return self._auth_flow + @property def client(self) -> httpx2.AsyncClient: if self._client is None: - self._client = httpx2.AsyncClient(verify=self.verify) + self._client = httpx2.AsyncClient(verify=self.verify, auth=self.auth_flow) self._client.base_url = self.uri + API_PREFIX return self._client diff --git a/warden/lib/qpu_client/auth.py b/warden/lib/qpu_client/auth.py index 23b239b..6f52e2d 100644 --- a/warden/lib/qpu_client/auth.py +++ b/warden/lib/qpu_client/auth.py @@ -5,7 +5,7 @@ from time import monotonic from typing import AsyncGenerator, Generator -import httpx +import httpx2 from warden.lib.config.config import QPUAuthConfig from warden.lib.qpu_client.retry import QPUClientRequestError @@ -22,7 +22,7 @@ class TokenRequestError(QPUClientRequestError): """Keycloak refused to issue a token and retrying cannot help.""" -class KeycloakClientCredentialsAuth(httpx.Auth): +class KeycloakClientCredentialsAuth(httpx2.Auth): """Attach a Keycloak service-account bearer token to each request. Implemented as an ``httpx.Auth`` so it runs inside the transport, below @@ -48,30 +48,27 @@ def __init__( self._expires_at: float = 0.0 def sync_auth_flow( - self, request: httpx.Request - ) -> Generator[httpx.Request, httpx.Response, None]: + self, request: httpx2.Request + ) -> Generator[httpx2.Request, httpx2.Response, None]: request.headers["Authorization"] = f"Bearer {self._sync_token()}" response = yield request - if response.status_code == httpx.codes.UNAUTHORIZED: + if response.status_code == httpx2.codes.UNAUTHORIZED: logger.info("QPU API returned 401, refreshing token and retrying once") request.headers["Authorization"] = f"Bearer {self._sync_token(force=True)}" yield request async def async_auth_flow( - self, request: httpx.Request - ) -> AsyncGenerator[httpx.Request, httpx.Response]: + self, request: httpx2.Request + ) -> AsyncGenerator[httpx2.Request, httpx2.Response]: request.headers["Authorization"] = f"Bearer {await self._async_token()}" response = yield request - if response.status_code == httpx.codes.UNAUTHORIZED: + if response.status_code == httpx2.codes.UNAUTHORIZED: logger.info("QPU API returned 401, refreshing token and retrying once") request.headers["Authorization"] = ( f"Bearer {await self._async_token(force=True)}" ) yield request - # ponytail: unlocked cache. Two concurrent requests in one process can both - # miss and both fetch a token; one wins and the loser wasted a request. Add - # a lock only if token-endpoint traffic ever becomes a problem. def _is_fresh(self) -> bool: return self._token is not None and monotonic() < self._expires_at @@ -83,14 +80,14 @@ def _token_request(self) -> tuple[str, dict[str, str]]: "client_secret": self.conf.secret, } - def _store(self, response: httpx.Response) -> str: + def _store(self, response: httpx2.Response) -> str: """Validate a token response, cache the token and return it.""" if response.status_code in FATAL_TOKEN_STATUSES: # Never log the response body of a token request: it may echo # credentials. The error field alone is the useful part. try: error = response.json().get("error", "unknown_error") - except ValueError: + except (ValueError, AttributeError): error = "unknown_error" raise TokenRequestError( f"Keycloak refused to issue a token for client " @@ -105,7 +102,15 @@ def _store(self, response: httpx.Response) -> str: token = payload["access_token"] expires_in = float(payload.get("expires_in", 0)) self._token = token - self._expires_at = monotonic() + max(expires_in - self.conf.leeway_s, 0.0) + configured_ttl = expires_in - self.conf.leeway_s + ttl = max(configured_ttl, expires_in / 2) + if ttl > configured_ttl: + logger.warning( + f"Configured leeway {self.conf.leeway_s}s leaves less than " + f"half of the {expires_in}s token lifespan; caching for " + f"{ttl}s (half the lifespan) instead." + ) + self._expires_at = monotonic() + ttl logger.debug( f"Obtained QPU API token for client '{self.conf.id}', " f"expires in {expires_in}s" @@ -117,13 +122,18 @@ def _sync_token(self, force: bool = False) -> str: assert self._token is not None return self._token url, data = self._token_request() - with httpx.Client(verify=self.verify) as client: + with httpx2.Client(verify=self.verify) as client: return self._store(client.post(url, data=data)) + # Note: unlocked check-then-fetch. Two concurrent async requests in one + # process can both miss and both fetch a token; one wins and the loser + # wasted a request. The sync path (scheduler) blocks its single event loop + # per fetch, so this race is only reachable via awaited callers here. Add + # a lock only if token-endpoint traffic ever becomes a problem. async def _async_token(self, force: bool = False) -> str: if not force and self._is_fresh(): assert self._token is not None return self._token url, data = self._token_request() - async with httpx.AsyncClient(verify=self.verify) as client: + async with httpx2.AsyncClient(verify=self.verify) as client: return self._store(await client.post(url, data=data)) diff --git a/warden/lib/qpu_client/retry.py b/warden/lib/qpu_client/retry.py index d14648b..f3bd945 100644 --- a/warden/lib/qpu_client/retry.py +++ b/warden/lib/qpu_client/retry.py @@ -50,6 +50,8 @@ def retry(max: int, sleep_s: float, no_retry: bool = False) -> Callable: UnhandledError: If decorator encounters an unnexpected exception. NotRetriedHTTPStatus: If the HTTP request returns with a non-retryable error code. MaxRetryError: If the maximum number of retries without success has been reached. + QPUClientRequestError: If `no_retry=True` or any subclass already classified as + non-retryable by the wrapped function (e.g. TokenRequestError) propagates unchanged. """ def decorator(func: Callable): @@ -60,6 +62,10 @@ def _handle_exception(e: Exception): elif isinstance(e, HTTPStatusError): if e.response.status_code not in RETRY_HTTP_EXIT_CODES: raise NotRetriedHTTPStatus(e) from e + elif isinstance(e, QPUClientRequestError): + # Already classified as non-retryable by the raiser (e.g. bad + # Keycloak credentials). Do not rewrap it as UnhandledError. + raise else: raise UnhandledError(e) from e diff --git a/warden/scheduler/worker.py b/warden/scheduler/worker.py index d6ce6fa..89e41aa 100644 --- a/warden/scheduler/worker.py +++ b/warden/scheduler/worker.py @@ -19,6 +19,8 @@ logger = logging.getLogger(__name__) +TERMINAL_STATUSES: tuple[JobStatus, ...] = ("ERROR", "DONE", "CANCELED") + class JobExecutionTracker: """Handles current job status and sends updates to db""" @@ -44,10 +46,6 @@ def job(self) -> QPUJobInfo: def is_error(self) -> bool: return self.status == "ERROR" - @property - def is_in_terminal_state(self) -> bool: - return self.status in ("DONE", "CANCELED", "ERROR") - @property def created_datetime(self) -> UTCDatetime: return self.job.created_datetime @@ -55,23 +53,41 @@ def created_datetime(self) -> UTCDatetime: async def update_job( self, qpu_job_info: QPUJobInfo, enforce_end_datetime: bool = False ): + was_terminal = self._status in TERMINAL_STATUSES self._qpu_job_info = qpu_job_info self._status = qpu_job_info.status or "ERROR" - if enforce_end_datetime and self._qpu_job_info.end_datetime is None: - self._qpu_job_info.end_datetime = datetime.now(timezone.utc) - await self.push_update() + await self.push_update(was_terminal, enforce_end_datetime) async def to_error(self): + was_terminal = self._status in TERMINAL_STATUSES self._status = "ERROR" - if self._qpu_job_info and self._qpu_job_info.end_datetime is None: - self._qpu_job_info.end_datetime = datetime.now(timezone.utc) - await self.push_update() + await self.push_update(was_terminal, enforce_end_datetime=True) def log(self, msg: str) -> None: self._log_buffer.append(msg + "\n") - async def push_update(self): - """Push update of job execution to db commit task through queue""" + async def push_update( + self, was_terminal: bool | None = None, enforce_end_datetime: bool = False + ): + """Push update of job execution to db commit task through queue + + `was_terminal` and `enforce_end_datetime` are set by `update_job`/ + `to_error` on a status transition, so the closing log line and the + backfilled `end_datetime` land in the same `JobUpdate` - hence the + same transaction - as the terminal status itself. Flushed separately, + the DB would briefly hold a finished job whose logs/dates are + incomplete, and anything that stops polling once the status is + terminal would read that incomplete state. + """ + if ( + enforce_end_datetime + and self._qpu_job_info + and self._qpu_job_info.end_datetime is None + ): + self._qpu_job_info.end_datetime = datetime.now(timezone.utc) + if was_terminal is False and self._status in TERMINAL_STATUSES: + logger.info("Job execution ended with status '%s'", self._status) + new_logs = "".join(self._log_buffer) self._log_buffer = [] @@ -172,7 +188,6 @@ async def execute_job( return await self.await_job_execution(job_tracker) - logger.info("Job execution ended with status '%s'", job_tracker.status) # Flush potential last updates before return await job_tracker.push_update() @@ -247,7 +262,7 @@ async def await_job_execution(self, job_tracker: JobExecutionTracker) -> None: polling_start = job_tracker.created_datetime await self._get_job_poll(job_tracker) - while not job_tracker.is_in_terminal_state: + while job_tracker.status not in TERMINAL_STATUSES: if self.is_timed_out(self.conf_sched.job_polling_timeout_s, polling_start): logger.warning( f"Job timed out (max {self.conf_sched.job_polling_timeout_s} s). " @@ -256,6 +271,10 @@ async def await_job_execution(self, job_tracker: JobExecutionTracker) -> None: ) try: qpu_job_info = await self.qpu_client.cancel_job(job_tracker.job.uid) + # Logged before the update so it is buffered into the same + # JobUpdate, and stays ahead of the closing line that a + # terminal status appends + logger.info("Job cancellation done") await job_tracker.update_job( qpu_job_info, enforce_end_datetime=True ) @@ -263,7 +282,6 @@ async def await_job_execution(self, job_tracker: JobExecutionTracker) -> None: logger.error(f"Failed cancelling job: {e}") await job_tracker.to_error() continue - logger.info("Job cancellation done") continue await asyncio.sleep(self.conf_sched.job_polling_interval_s) await self._get_job_poll(job_tracker) From c3d490d4d798d525b52314350e3dbf96d5cb4208 Mon Sep 17 00:00:00 2001 From: badtst Date: Mon, 31 Aug 2026 15:36:30 +0000 Subject: [PATCH 3/5] feat(auth): add QPUAuthConfig for Keycloak client credentials feat(auth): add Keycloak client_credentials httpx auth flow feat(auth): add QPUAuthConfig for Keycloak client credentials feat(auth): add Keycloak client_credentials httpx auth flow feat(auth): attach Keycloak auth flow to QPU API clients fix(retry): do not rewrap already-classified QPU client errors docs(auth): document qpu.auth config block feat(auth): add QPUAuthConfig for Keycloak client credentials feat(auth): add Keycloak client_credentials httpx auth flow feat(auth): attach Keycloak auth flow to QPU API clients fix(retry): do not rewrap already-classified QPU client errors docs(auth): document qpu.auth config block fix(auth): handle short token lifespans and address review findings Cache for half the lifespan (with a WARNING) instead of disabling the cache entirely when expires_in <= leeway_s. Also catch AttributeError alongside ValueError when parsing a non-object token error body, narrow the unlocked-cache ponytail comment to the async path where the race is actually reachable, tighten config.py's auth_flow typing, note that QPUClientRequestError subclasses can propagate through retry(), and add coverage for the no-auth-header and POST-401-replay paths. fix(typechecking): Fix type-checking error feat(auth): add QPUAuthConfig for Keycloak client credentials feat(auth): add Keycloak client_credentials httpx auth flow feat(auth): attach Keycloak auth flow to QPU API clients fix(retry): do not rewrap already-classified QPU client errors docs(auth): document qpu.auth config block fix(auth): handle short token lifespans and address review findings Cache for half the lifespan (with a WARNING) instead of disabling the cache entirely when expires_in <= leeway_s. Also catch AttributeError alongside ValueError when parsing a non-object token error body, narrow the unlocked-cache ponytail comment to the async path where the race is actually reachable, tighten config.py's auth_flow typing, note that QPUClientRequestError subclasses can propagate through retry(), and add coverage for the no-auth-header and POST-401-replay paths. fix(typechecking): Fix type-checking error fix(test): Fix test flakiness feat(auth): add QPUAuthConfig for Keycloak client credentials feat(auth): add Keycloak client_credentials httpx auth flow feat(auth): attach Keycloak auth flow to QPU API clients fix(retry): do not rewrap already-classified QPU client errors docs(auth): document qpu.auth config block fix(auth): handle short token lifespans and address review findings Cache for half the lifespan (with a WARNING) instead of disabling the cache entirely when expires_in <= leeway_s. Also catch AttributeError alongside ValueError when parsing a non-object token error body, narrow the unlocked-cache ponytail comment to the async path where the race is actually reachable, tighten config.py's auth_flow typing, note that QPUClientRequestError subclasses can propagate through retry(), and add coverage for the no-auth-header and POST-401-replay paths. fix(typechecking): Fix type-checking error fix(test): Fix test flakiness Fix test and refacto PR comments test(auth): Add QPU Auth tests feat(auth): add QPUAuthConfig for Keycloak client credentials feat(auth): add Keycloak client_credentials httpx auth flow feat(auth): add QPUAuthConfig for Keycloak client credentials feat(auth): add Keycloak client_credentials httpx auth flow feat(auth): attach Keycloak auth flow to QPU API clients fix(retry): do not rewrap already-classified QPU client errors docs(auth): document qpu.auth config block feat(auth): add QPUAuthConfig for Keycloak client credentials feat(auth): add Keycloak client_credentials httpx auth flow feat(auth): attach Keycloak auth flow to QPU API clients fix(retry): do not rewrap already-classified QPU client errors docs(auth): document qpu.auth config block fix(auth): handle short token lifespans and address review findings Cache for half the lifespan (with a WARNING) instead of disabling the cache entirely when expires_in <= leeway_s. Also catch AttributeError alongside ValueError when parsing a non-object token error body, narrow the unlocked-cache ponytail comment to the async path where the race is actually reachable, tighten config.py's auth_flow typing, note that QPUClientRequestError subclasses can propagate through retry(), and add coverage for the no-auth-header and POST-401-replay paths. fix(typechecking): Fix type-checking error feat(auth): add QPUAuthConfig for Keycloak client credentials feat(auth): add Keycloak client_credentials httpx auth flow feat(auth): attach Keycloak auth flow to QPU API clients fix(retry): do not rewrap already-classified QPU client errors docs(auth): document qpu.auth config block fix(auth): handle short token lifespans and address review findings Cache for half the lifespan (with a WARNING) instead of disabling the cache entirely when expires_in <= leeway_s. Also catch AttributeError alongside ValueError when parsing a non-object token error body, narrow the unlocked-cache ponytail comment to the async path where the race is actually reachable, tighten config.py's auth_flow typing, note that QPUClientRequestError subclasses can propagate through retry(), and add coverage for the no-auth-header and POST-401-replay paths. fix(typechecking): Fix type-checking error fix(test): Fix test flakiness feat(auth): add QPUAuthConfig for Keycloak client credentials feat(auth): add Keycloak client_credentials httpx auth flow feat(auth): attach Keycloak auth flow to QPU API clients fix(retry): do not rewrap already-classified QPU client errors docs(auth): document qpu.auth config block fix(auth): handle short token lifespans and address review findings Cache for half the lifespan (with a WARNING) instead of disabling the cache entirely when expires_in <= leeway_s. Also catch AttributeError alongside ValueError when parsing a non-object token error body, narrow the unlocked-cache ponytail comment to the async path where the race is actually reachable, tighten config.py's auth_flow typing, note that QPUClientRequestError subclasses can propagate through retry(), and add coverage for the no-auth-header and POST-401-replay paths. fix(typechecking): Fix type-checking error fix(test): Fix test flakiness Fix test and refacto PR comments test(auth): Add QPU Auth tests --- tests/lib/qpu_client/test_auth.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/lib/qpu_client/test_auth.py b/tests/lib/qpu_client/test_auth.py index c829623..8a9a9cb 100644 --- a/tests/lib/qpu_client/test_auth.py +++ b/tests/lib/qpu_client/test_auth.py @@ -212,12 +212,13 @@ def test_short_lived_token_still_caches_with_warning( assert any(record.levelno == logging.WARNING for record in caplog.records) -def test_client_sends_no_authorization_header_without_auth_config( +@pytest.mark.asyncio +async def test_client_sends_no_authorization_header_without_auth_config( httpx_mock: HTTPXMock, ): httpx_mock.add_response(url=QPU_URL, json={"data": {}}) - response = QPUConfig(uri="http://qpu:4300").client.get(QPU_URL) + response = await QPUConfig(uri="http://qpu:4300").client.get(QPU_URL) assert "Authorization" not in response.request.headers From fa2b18ca50d5bb6e451c9546bc9bb02f93c74bb7 Mon Sep 17 00:00:00 2001 From: badtst Date: Mon, 31 Aug 2026 16:14:42 +0000 Subject: [PATCH 4/5] Fix httpx2.auth design adhering more closely to the documentation --- tests/lib/qpu_client/test_auth.py | 94 ++++++++++++++----------------- warden/lib/config/config.py | 13 +++-- warden/lib/qpu_client/auth.py | 89 +++++++++++------------------ 3 files changed, 83 insertions(+), 113 deletions(-) diff --git a/tests/lib/qpu_client/test_auth.py b/tests/lib/qpu_client/test_auth.py index 8a9a9cb..28c7677 100644 --- a/tests/lib/qpu_client/test_auth.py +++ b/tests/lib/qpu_client/test_auth.py @@ -4,7 +4,7 @@ import logging import pytest -from httpx2 import AsyncClient, Client, HTTPStatusError +from httpx2 import AsyncClient, HTTPStatusError from pytest_httpx2 import HTTPXMock from warden.lib.config.config import QPUAuthConfig, QPUConfig @@ -24,7 +24,8 @@ def auth_conf() -> QPUAuthConfig: ) -def test_token_is_fetched_once_and_reused(httpx_mock: HTTPXMock, auth_conf): +@pytest.mark.asyncio +async def test_token_is_fetched_once_and_reused(httpx_mock: HTTPXMock, auth_conf): httpx_mock.add_response( url=TOKEN_URL, json={"access_token": "tok-1", "expires_in": 300}, @@ -33,9 +34,9 @@ def test_token_is_fetched_once_and_reused(httpx_mock: HTTPXMock, auth_conf): httpx_mock.add_response(url=QPU_URL, json={"data": {}}) auth = KeycloakClientCredentialsAuth(auth_conf) - with Client(auth=auth) as client: - first = client.get(QPU_URL) - second = client.get(QPU_URL) + async with AsyncClient(auth=auth) as client: + first = await client.get(QPU_URL) + second = await client.get(QPU_URL) assert first.request.headers["Authorization"] == "Bearer tok-1" assert second.request.headers["Authorization"] == "Bearer tok-1" @@ -43,15 +44,18 @@ def test_token_is_fetched_once_and_reused(httpx_mock: HTTPXMock, auth_conf): assert len(token_requests) == 1 -def test_token_request_uses_client_credentials_grant(httpx_mock: HTTPXMock, auth_conf): +@pytest.mark.asyncio +async def test_token_request_uses_client_credentials_grant( + httpx_mock: HTTPXMock, auth_conf +): httpx_mock.add_response( url=TOKEN_URL, json={"access_token": "tok-1", "expires_in": 300} ) httpx_mock.add_response(url=QPU_URL, json={"data": {}}) auth = KeycloakClientCredentialsAuth(auth_conf) - with Client(auth=auth) as client: - client.get(QPU_URL) + async with AsyncClient(auth=auth) as client: + await client.get(QPU_URL) token_request = next( r for r in httpx_mock.get_requests() if str(r.url) == TOKEN_URL @@ -62,7 +66,10 @@ def test_token_request_uses_client_credentials_grant(httpx_mock: HTTPXMock, auth assert "client_secret=s3cret" in body -def test_expired_token_is_refreshed(httpx_mock: HTTPXMock, auth_conf, monkeypatch): +@pytest.mark.asyncio +async def test_expired_token_is_refreshed( + httpx_mock: HTTPXMock, auth_conf, monkeypatch +): # expires_in 300 with leeway 30 means the token is stale after 270s. clock = {"now": 1_000.0} monkeypatch.setattr("warden.lib.qpu_client.auth.monotonic", lambda: clock["now"]) @@ -76,16 +83,17 @@ def test_expired_token_is_refreshed(httpx_mock: HTTPXMock, auth_conf, monkeypatc httpx_mock.add_response(url=QPU_URL, json={"data": {}}) auth = KeycloakClientCredentialsAuth(auth_conf) - with Client(auth=auth) as client: - first = client.get(QPU_URL) + async with AsyncClient(auth=auth) as client: + first = await client.get(QPU_URL) clock["now"] += 271 - second = client.get(QPU_URL) + second = await client.get(QPU_URL) assert first.request.headers["Authorization"] == "Bearer tok-1" assert second.request.headers["Authorization"] == "Bearer tok-2" -def test_401_triggers_one_refresh_and_one_retry(httpx_mock: HTTPXMock, auth_conf): +@pytest.mark.asyncio +async def test_401_triggers_one_refresh_and_one_retry(httpx_mock: HTTPXMock, auth_conf): httpx_mock.add_response( url=TOKEN_URL, json={"access_token": "stale", "expires_in": 300} ) @@ -96,8 +104,8 @@ def test_401_triggers_one_refresh_and_one_retry(httpx_mock: HTTPXMock, auth_conf httpx_mock.add_response(url=QPU_URL, json={"data": {}}) auth = KeycloakClientCredentialsAuth(auth_conf) - with Client(auth=auth) as client: - response = client.get(QPU_URL) + async with AsyncClient(auth=auth) as client: + response = await client.get(QPU_URL) assert response.status_code == 200 assert response.request.headers["Authorization"] == "Bearer fresh" @@ -105,7 +113,8 @@ def test_401_triggers_one_refresh_and_one_retry(httpx_mock: HTTPXMock, auth_conf assert len(qpu_requests) == 2 -def test_persistent_401_is_not_retried_forever(httpx_mock: HTTPXMock, auth_conf): +@pytest.mark.asyncio +async def test_persistent_401_is_not_retried_forever(httpx_mock: HTTPXMock, auth_conf): httpx_mock.add_response( url=TOKEN_URL, json={"access_token": "tok-1", "expires_in": 300} ) @@ -116,8 +125,8 @@ def test_persistent_401_is_not_retried_forever(httpx_mock: HTTPXMock, auth_conf) httpx_mock.add_response(url=QPU_URL, status_code=401) auth = KeycloakClientCredentialsAuth(auth_conf) - with Client(auth=auth) as client: - response = client.get(QPU_URL) + async with AsyncClient(auth=auth) as client: + response = await client.get(QPU_URL) # The second 401 is surfaced, not retried again. assert response.status_code == 401 @@ -125,8 +134,9 @@ def test_persistent_401_is_not_retried_forever(httpx_mock: HTTPXMock, auth_conf) assert len(qpu_requests) == 2 +@pytest.mark.asyncio @pytest.mark.parametrize("status_code", [400, 401]) -def test_bad_credentials_raise_token_request_error( +async def test_bad_credentials_raise_token_request_error( httpx_mock: HTTPXMock, auth_conf, status_code ): httpx_mock.add_response( @@ -136,12 +146,13 @@ def test_bad_credentials_raise_token_request_error( ) auth = KeycloakClientCredentialsAuth(auth_conf) - with Client(auth=auth) as client: + async with AsyncClient(auth=auth) as client: with pytest.raises(TokenRequestError, match="invalid_client"): - client.get(QPU_URL) + await client.get(QPU_URL) -def test_keycloak_5xx_raises_retryable_http_status_error( +@pytest.mark.asyncio +async def test_keycloak_5xx_raises_retryable_http_status_error( httpx_mock: HTTPXMock, auth_conf ): # 503 must stay an httpx.HTTPStatusError so the existing retry decorator @@ -149,9 +160,9 @@ def test_keycloak_5xx_raises_retryable_http_status_error( httpx_mock.add_response(url=TOKEN_URL, status_code=503) auth = KeycloakClientCredentialsAuth(auth_conf) - with Client(auth=auth) as client: + async with AsyncClient(auth=auth) as client: with pytest.raises(HTTPStatusError): - client.get(QPU_URL) + await client.get(QPU_URL) @pytest.mark.asyncio @@ -169,27 +180,7 @@ async def test_async_flow_attaches_token(httpx_mock: HTTPXMock, auth_conf): @pytest.mark.asyncio -async def test_async_flow_reuses_token_cached_by_sync_flow( - httpx_mock: HTTPXMock, auth_conf -): - httpx_mock.add_response( - url=TOKEN_URL, json={"access_token": "shared", "expires_in": 300} - ) - httpx_mock.add_response(url=QPU_URL, json={"data": {}}) - httpx_mock.add_response(url=QPU_URL, json={"data": {}}) - - auth = KeycloakClientCredentialsAuth(auth_conf) - with Client(auth=auth) as client: - client.get(QPU_URL) - async with AsyncClient(auth=auth) as client: - response = await client.get(QPU_URL) - - assert response.request.headers["Authorization"] == "Bearer shared" - token_requests = [r for r in httpx_mock.get_requests() if str(r.url) == TOKEN_URL] - assert len(token_requests) == 1 - - -def test_short_lived_token_still_caches_with_warning( +async def test_short_lived_token_still_caches_with_warning( httpx_mock: HTTPXMock, auth_conf, caplog ): # expires_in 30 with the default leeway_s 30 would clamp to 0 without the @@ -203,9 +194,9 @@ def test_short_lived_token_still_caches_with_warning( auth = KeycloakClientCredentialsAuth(auth_conf) with caplog.at_level(logging.WARNING, logger="warden.lib.qpu_client.auth"): - with Client(auth=auth) as client: - client.get(QPU_URL) - client.get(QPU_URL) + async with AsyncClient(auth=auth) as client: + await client.get(QPU_URL) + await client.get(QPU_URL) token_requests = [r for r in httpx_mock.get_requests() if str(r.url) == TOKEN_URL] assert len(token_requests) == 1 @@ -223,7 +214,8 @@ async def test_client_sends_no_authorization_header_without_auth_config( assert "Authorization" not in response.request.headers -def test_401_on_post_retries_with_fresh_token_and_identical_body( +@pytest.mark.asyncio +async def test_401_on_post_retries_with_fresh_token_and_identical_body( httpx_mock: HTTPXMock, auth_conf ): httpx_mock.add_response( @@ -237,8 +229,8 @@ def test_401_on_post_retries_with_fresh_token_and_identical_body( body = {"circuit": "bell", "shots": 100} auth = KeycloakClientCredentialsAuth(auth_conf) - with Client(auth=auth) as client: - response = client.post(QPU_URL, json=body) + async with AsyncClient(auth=auth) as client: + response = await client.post(QPU_URL, json=body) assert response.status_code == 200 assert response.request.headers["Authorization"] == "Bearer fresh" diff --git a/warden/lib/config/config.py b/warden/lib/config/config.py index 6d45ebd..62da152 100644 --- a/warden/lib/config/config.py +++ b/warden/lib/config/config.py @@ -197,9 +197,14 @@ class QPUAuthConfig(WardenSettings): id: str = Field(description="OIDC client_id") - secret: str = Field(description="OIDC client_secret. Provide via WARDEN_QPU_AUTH_SECRET, never in YAML.") + secret: str = Field( + description="OIDC client_secret. Provide via WARDEN_QPU_AUTH_SECRET, never in YAML." + ) - leeway_s: float = Field(default=30, description="Refresh this many seconds before the token actually expires.") + leeway_s: float = Field( + default=30, + description="Refresh this many seconds before the token actually expires.", + ) @property def token_url(self) -> str: @@ -263,9 +268,7 @@ def auth_flow(self) -> httpx2.Auth | None: if self._auth_flow is None: from warden.lib.qpu_client.auth import KeycloakClientCredentialsAuth - self._auth_flow = KeycloakClientCredentialsAuth( - self.auth, verify=self.verify - ) + self._auth_flow = KeycloakClientCredentialsAuth(self.auth) return self._auth_flow @property diff --git a/warden/lib/qpu_client/auth.py b/warden/lib/qpu_client/auth.py index 6f52e2d..fad4323 100644 --- a/warden/lib/qpu_client/auth.py +++ b/warden/lib/qpu_client/auth.py @@ -1,9 +1,8 @@ """Keycloak client_credentials authentication for outbound QPU API calls.""" import logging -import ssl from time import monotonic -from typing import AsyncGenerator, Generator +from typing import Generator import httpx2 @@ -31,57 +30,56 @@ class KeycloakClientCredentialsAuth(httpx2.Auth): as an immediate, non-retryable ``NotRetriedHTTPStatus``. Here it is just a refresh. + ``requires_response_body`` tells httpx to read the token response before + handing it back, since ``_store`` needs the JSON body. + Args: conf: Keycloak credentials and endpoint. - verify: httpx TLS verification setting for the token request. """ - def __init__( - self, - conf: QPUAuthConfig, - verify: bool | str | ssl.SSLContext = True, - ) -> None: + requires_response_body = True + + def __init__(self, conf: QPUAuthConfig) -> None: self.conf = conf - self.verify = verify self._token: str | None = None # monotonic() deadline after which the cached token is considered stale. self._expires_at: float = 0.0 - def sync_auth_flow( + # Note: unlocked check-then-fetch. Two concurrent requests can both miss + # and both fetch a token; one wins and the loser wasted a request. Add a + # lock only if token-endpoint traffic ever becomes a problem. + def auth_flow( self, request: httpx2.Request ) -> Generator[httpx2.Request, httpx2.Response, None]: - request.headers["Authorization"] = f"Bearer {self._sync_token()}" - response = yield request - if response.status_code == httpx2.codes.UNAUTHORIZED: - logger.info("QPU API returned 401, refreshing token and retrying once") - request.headers["Authorization"] = f"Bearer {self._sync_token(force=True)}" - yield request - - async def async_auth_flow( - self, request: httpx2.Request - ) -> AsyncGenerator[httpx2.Request, httpx2.Response]: - request.headers["Authorization"] = f"Bearer {await self._async_token()}" + if not self._is_fresh(): + token_response = yield self._token_request() + self._store(token_response) + assert self._token is not None + request.headers["Authorization"] = f"Bearer {self._token}" response = yield request if response.status_code == httpx2.codes.UNAUTHORIZED: logger.info("QPU API returned 401, refreshing token and retrying once") - request.headers["Authorization"] = ( - f"Bearer {await self._async_token(force=True)}" - ) + token_response = yield self._token_request() + self._store(token_response) + request.headers["Authorization"] = f"Bearer {self._token}" yield request def _is_fresh(self) -> bool: return self._token is not None and monotonic() < self._expires_at - def _token_request(self) -> tuple[str, dict[str, str]]: - """Return the (url, form data) for a client_credentials token request.""" - return self.conf.token_url, { - "grant_type": "client_credentials", - "client_id": self.conf.id, - "client_secret": self.conf.secret, - } + def _token_request(self) -> httpx2.Request: + return httpx2.Request( + "POST", + self.conf.token_url, + data={ + "grant_type": "client_credentials", + "client_id": self.conf.id, + "client_secret": self.conf.secret, + }, + ) - def _store(self, response: httpx2.Response) -> str: - """Validate a token response, cache the token and return it.""" + def _store(self, response: httpx2.Response) -> None: + """Validate a token response and cache the token.""" if response.status_code in FATAL_TOKEN_STATUSES: # Never log the response body of a token request: it may echo # credentials. The error field alone is the useful part. @@ -99,9 +97,8 @@ def _store(self, response: httpx2.Response) -> str: response.raise_for_status() payload = response.json() - token = payload["access_token"] + self._token = payload["access_token"] expires_in = float(payload.get("expires_in", 0)) - self._token = token configured_ttl = expires_in - self.conf.leeway_s ttl = max(configured_ttl, expires_in / 2) if ttl > configured_ttl: @@ -115,25 +112,3 @@ def _store(self, response: httpx2.Response) -> str: f"Obtained QPU API token for client '{self.conf.id}', " f"expires in {expires_in}s" ) - return token - - def _sync_token(self, force: bool = False) -> str: - if not force and self._is_fresh(): - assert self._token is not None - return self._token - url, data = self._token_request() - with httpx2.Client(verify=self.verify) as client: - return self._store(client.post(url, data=data)) - - # Note: unlocked check-then-fetch. Two concurrent async requests in one - # process can both miss and both fetch a token; one wins and the loser - # wasted a request. The sync path (scheduler) blocks its single event loop - # per fetch, so this race is only reachable via awaited callers here. Add - # a lock only if token-endpoint traffic ever becomes a problem. - async def _async_token(self, force: bool = False) -> str: - if not force and self._is_fresh(): - assert self._token is not None - return self._token - url, data = self._token_request() - async with httpx2.AsyncClient(verify=self.verify) as client: - return self._store(await client.post(url, data=data)) From 968452516d06e0fa52bda7e0ef2315fdf66b49f6 Mon Sep 17 00:00:00 2001 From: badtst Date: Tue, 1 Sep 2026 08:46:03 +0000 Subject: [PATCH 5/5] Fix auth section config generation --- tests/test_generate_config.py | 46 +++++++++++++++++++++++++--- warden/lib/config/config.py | 6 ++-- warden/lib/config/generate_config.py | 27 +++++++++++----- 3 files changed, 65 insertions(+), 14 deletions(-) diff --git a/tests/test_generate_config.py b/tests/test_generate_config.py index 161db8c..3f6c387 100644 --- a/tests/test_generate_config.py +++ b/tests/test_generate_config.py @@ -8,6 +8,7 @@ Config, MariadbConfig, PostgresConfig, + QPUAuthConfig, QPUConfig, SchedulerConfig, SqliteConfig, @@ -25,12 +26,17 @@ APIConfig, SchedulerConfig, QPUConfig, + QPUAuthConfig, SqliteConfig, PostgresConfig, MariadbConfig, ) for name in model.model_fields } +# Fields whose type is itself a nested model (e.g. QPUConfig.auth, an Optional +# QPUAuthConfig) are rendered as an uncommented "name:" header, with their own +# fields commented out beneath, rather than a single "# name:" scalar line. +NESTED_MODEL_FIELD_NAMES = {"auth"} def test_generate_config_is_valid_yaml(): @@ -47,7 +53,8 @@ def test_generate_config_documents_every_field(): generated = generate_config() for name in ALL_FIELD_NAMES: - assert f"# {name}:" in generated + prefix = "" if name in NESTED_MODEL_FIELD_NAMES else "# " + assert f"{prefix}{name}:" in generated def test_generate_config_preserves_existing_overrides(): @@ -101,9 +108,40 @@ class Outer(BaseModel): overridden = _render_section_fields(Outer.model_fields, {"inner": {"value": 42}}, 1) assert " value: 42" in overridden - assert yaml.safe_load(f"outer:\n{overridden}") == { - "outer": {"inner": {"value": 42}} - } + + +def test_render_fields_indents_optional_nested_models(): + """Test that an Optional (``Model | None``) nested field is recursed into + just like a plain nested model, falling back to the nested model's own + docstring since it has no field-level description of its own""" + + class Inner(BaseModel): + """Inner docstring.""" + + value: int = Field(default=1, description="An inner value.") + + class Outer(BaseModel): + inner: Inner | None = None + + generated = _render_section_fields(Outer.model_fields, {}, 1) + + assert ( + " # Inner docstring.\n inner:\n # An inner value.\n # value: 1" + in generated + ) + + +def test_generate_config_documents_qpu_auth_section(): + """Test that qpu.auth (an Optional QPUAuthConfig) is recursed into: its + own docstring and every one of its fields must appear, not just an opaque + "# auth: null" line""" + generated = generate_config() + + assert "auth:" in generated + assert "# auth: null" not in generated + assert "Keycloak client_credentials configuration" in generated + for name in QPUAuthConfig.model_fields: + assert f"# {name}:" in generated def test_generate_writes_directly_when_no_previous_file(tmp_path, monkeypatch): diff --git a/warden/lib/config/config.py b/warden/lib/config/config.py index 62da152..c89a26e 100644 --- a/warden/lib/config/config.py +++ b/warden/lib/config/config.py @@ -185,13 +185,13 @@ class QPUAuthConfig(WardenSettings): """Keycloak client_credentials configuration for outbound QPU API calls. Presence of this section is what enables authentication. There is - deliberately no separate ``enabled`` flag: a second switch can drift out of - sync with the credentials it guards. ``url``, ``id`` and ``secret`` have no + deliberately no separate `enabled` flag: a second switch can drift out of + sync with the credentials it guards. `url`, `id` and `secret` have no defaults, so a partially configured section is a startup validation error rather than a silent fallback to unauthenticated requests. """ - url: str = Field(description="Keycloak base URL, e.g. http://keycloak:8080") + url: str = Field(description="Keycloak base URL, for example http://keycloak:8080") realm: str = Field(default="pasqos") diff --git a/warden/lib/config/generate_config.py b/warden/lib/config/generate_config.py index 5631a4a..959e750 100644 --- a/warden/lib/config/generate_config.py +++ b/warden/lib/config/generate_config.py @@ -5,6 +5,8 @@ import re import sys import textwrap +import types +import typing from enum import Enum from pathlib import Path @@ -85,10 +87,14 @@ def _wrap_indented_text(text: str, indent: str) -> list[str]: def _get_nested_model(field: FieldInfo) -> type[BaseModel] | None: - """Returns if the field type is indeed a BaseModel subclass""" - field_annotation = field.annotation - if isinstance(field_annotation, type) and issubclass(field_annotation, BaseModel): - return field_annotation + """Returns the field's BaseModel type, unwrapping an Optional (``Model | None``).""" + annotation = field.annotation + if typing.get_origin(annotation) in (typing.Union, types.UnionType): + non_none = [a for a in typing.get_args(annotation) if a is not type(None)] + if len(non_none) == 1: + annotation = non_none[0] + if isinstance(annotation, type) and issubclass(annotation, BaseModel): + return annotation return None @@ -126,10 +132,19 @@ def _render_field( indent = INDENT_UNIT * depth + # Check if contains a nested model, to recursively render it below and, + # absent a field-level description, fall back to its docstring's summary + # line (the rest of a multi-paragraph docstring is skipped). + nested_model = _get_nested_model(field_info) + nested_doc = nested_model.__doc__ if nested_model else None + description = field_info.description or ( + nested_doc.strip().splitlines()[0] if nested_doc else None + ) + # Add comment lines lines = [ wrapped - for paragraph in (field_info.description or "").split("\n") + for paragraph in (description or "").split("\n") if paragraph for sentence in SENTENCE_RE.split(paragraph) if sentence @@ -139,8 +154,6 @@ def _render_field( # Get matching previously set data that we need to migrate to new config data_to_migrate = _existing_value(previous_section_data, name, field_info) - # Check if contains a nested model and recursively renders it - nested_model = _get_nested_model(field_info) if nested_model is not None: nested_existing = data_to_migrate if isinstance(data_to_migrate, dict) else {} lines.append(f"{indent}{name}:")