diff --git a/api/services/okta_service.py b/api/services/okta_service.py index 91137c3d..4637407a 100644 --- a/api/services/okta_service.py +++ b/api/services/okta_service.py @@ -32,6 +32,14 @@ # Page size for cursor-paginated list endpoints. Okta caps most list endpoints # at 200 per page; the facade drives the ``after`` cursor to walk pages. LIST_PAGE_LIMIT = 200 +# Extra attempts for a call that raises ``OktaTransientError``. Transient Okta +# failures (a 429/5xx or timeout, and — via the reclassification in ``_call`` — a +# pooled keep-alive connection Okta dropped server-side, which the SDK surfaces as +# a None-response error) are worth one immediate retry: it gets a fresh connection +# and usually succeeds, recovering within the run instead of skipping the resource +# until the next reconcile. Kept at 1 so a genuinely-degraded endpoint, or a +# lingering 429, still gives up quickly rather than amplifying load. +OKTA_TRANSIENT_RETRIES = 1 logger = logging.getLogger(__name__) @@ -113,7 +121,17 @@ def __getattr__(self, name: str) -> Any: @functools.wraps(attr) async def wrapper(*args: Any, **kwargs: Any) -> Any: - return await OktaService._call(attr(*args, **kwargs)) + # Rebuild the coroutine each attempt — a spent coroutine can't be + # re-awaited — and retry a bounded number of times on a transient + # failure, so a stale pooled connection is recovered within the run + # rather than deferred to the next reconcile. + for attempt in range(OKTA_TRANSIENT_RETRIES + 1): + try: + return await OktaService._call(attr(*args, **kwargs)) + except OktaTransientError: + if attempt == OKTA_TRANSIENT_RETRIES: + raise + raise AssertionError("unreachable") # pragma: no cover return wrapper diff --git a/tests/test_okta_service.py b/tests/test_okta_service.py index a877f6af..851daf75 100644 --- a/tests/test_okta_service.py +++ b/tests/test_okta_service.py @@ -3,12 +3,20 @@ from typing import Any, Optional from unittest.mock import AsyncMock, MagicMock, patch +import pytest + from okta.models.add_group_request import AddGroupRequest from okta.models.group import Group as OktaGroupType from okta.models.group_rule import GroupRule as OktaGroupRuleType from okta.models.user_schema import UserSchema as OktaUserSchemaType -from api.services.okta_service import OktaService, UserSchema, is_managed_group +from api.services.okta_service import ( + OKTA_TRANSIENT_RETRIES, + OktaService, + OktaTransientError, + UserSchema, + is_managed_group, +) from tests.factories import UserFactory @@ -215,3 +223,40 @@ def tracking_build_client() -> Any: # Each concurrent call built its own client; nothing shared. assert len(clients) == call_count assert len({id(client) for client in clients}) == call_count + + +async def test_transient_error_is_retried_then_succeeds() -> None: + """A transient Okta failure is retried within the same call and can recover. + + The SDK surfaces a transient outcome (here a 503) that ``_call`` turns into an + ``OktaTransientError``; ``_WrapperClient`` retries it, and the next attempt — + on a fresh connection in production — succeeds. + """ + service = OktaService() + service.initialize("fake.domain", "fake.token") + + transient = (None, MagicMock(status_code=503, headers={}), MagicMock(status=503)) + success = (UserFactory(), MagicMock(status_code=200, headers={}), None) + get_user = AsyncMock(side_effect=[transient, success]) + + with patch("okta.client.Client.get_user", get_user): + user = await service.get_user("okta_id") + + assert user is not None + # One transient failure, retried once, then success. + assert get_user.call_count == 2 + + +async def test_transient_error_gives_up_after_bounded_retries() -> None: + """A persistently transient endpoint stops after the bounded retries.""" + service = OktaService() + service.initialize("fake.domain", "fake.token") + + transient = (None, MagicMock(status_code=503, headers={}), MagicMock(status=503)) + get_user = AsyncMock(return_value=transient) + + with patch("okta.client.Client.get_user", get_user): + with pytest.raises(OktaTransientError): + await service.get_user("okta_id") + + assert get_user.call_count == OKTA_TRANSIENT_RETRIES + 1