update

Signed-off-by: Ryan S <267728323+ironcommit@users.noreply.github.com>

Author: ironcommit

services/core/auth/tests/integration/test_scoped_access_keys.py

@@ -2,23 +2,52 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import uuid
+from collections.abc import Callable
 from pathlib import Path
+from time import monotonic, sleep
 from unittest.mock import patch
 
+import pytest
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa
 from fastapi.testclient import TestClient
+from httpx import Response
 from nmp.common.auth.access_keys import public_jwk_from_private_key_pem, validate_access_key_token
 from nmp.common.auth.jwt import TokenClaims
 from nmp.common.config import AuthConfig
 from nmp.common.config.base import AccessKeyConfig, TokenSigningConfig
 from nmp.core.auth.config import AuthServiceConfig
 from nmp.testing.client import create_test_client
 
+# Keep this test in its own xdist group while embedded PDP uses process-wide
+# wasmtime state. This limits cross-test scheduling noise under --dist loadgroup.
+pytestmark = pytest.mark.xdist_group("auth_scoped_access_keys")
+
 ACCESS_KEYS_PATH = "/apis/auth/v2/access-keys"
 IAM_ROLE_BINDINGS_PATH = "/apis/auth/v2/iam/role-bindings"
 WORKSPACES_PATH = "/apis/entities/v2/workspaces"
 SERVICE_HEADERS = {"X-NMP-Principal-Id": "service:integration-test"}
+AUTHZ_PROPAGATION_TIMEOUT_SECONDS = 5.0
+AUTHZ_PROPAGATION_POLL_INTERVAL_SECONDS = 0.05
+
+
+def _wait_for_authorization_response(
+    request: Callable[[], Response],
+    *,
+    expected_status_code: int,
+) -> Response:
+    deadline = monotonic() + AUTHZ_PROPAGATION_TIMEOUT_SECONDS
+    response = request()
+    while response.status_code != expected_status_code and monotonic() < deadline:
+        sleep(AUTHZ_PROPAGATION_POLL_INTERVAL_SECONDS)
+        response = request()
+
+    if response.status_code != expected_status_code:
+        raise AssertionError(
+            f"Timed out waiting for authorization response {expected_status_code}; "
+            f"last response {response.status_code}: {response.text}"
+        )
+    return response
 
 
 def _tamper_jwt(token: str) -> str:
@@ -57,8 +86,8 @@ def _auth_configs(private_key_file: str) -> tuple[AuthConfig, AuthServiceConfig]
     )
     service_config = AuthServiceConfig(
         **shared_config.model_dump(),
-        policy_data_refresh_interval=0.05,
-        bundle_cache_seconds=0,
+        policy_data_refresh_interval=0.2,
+        bundle_cache_seconds=0.1,
         admin_email="admin@example.com",
     )
     return shared_config, service_config
@@ -117,35 +146,37 @@ def test_scoped_access_key_created_by_auth_service_authenticates_platform_reques
             async def validate_with_local_jwks(config: AuthConfig, token: str) -> TokenClaims | None:
                 return await validate_access_key_token(config, token, jwks_override=jwks)
 
-            with patch("nmp.common.auth.access_keys.validate_access_key_token", validate_with_local_jwks):
-                response = client.get(
-                    f"{WORKSPACES_PATH}/{workspace}",
-                    headers={"Authorization": f"Bearer {access_key}"},
-                )
+            def get_workspace_with_access_key(token: str) -> Response:
+                with patch("nmp.common.auth.access_keys.validate_access_key_token", validate_with_local_jwks):
+                    return client.get(
+                        f"{WORKSPACES_PATH}/{workspace}",
+                        headers={"Authorization": f"Bearer {token}"},
+                    )
+
+            def authenticate_with_access_key(token: str) -> Response:
+                with patch("nmp.common.auth.access_keys.validate_access_key_token", validate_with_local_jwks):
+                    return client.get(
+                        "/apis/auth/authenticate",
+                        headers={"Authorization": f"Bearer {token}"},
+                    )
+
+            response = _wait_for_authorization_response(
+                lambda: get_workspace_with_access_key(access_key),
+                expected_status_code=200,
+            )
 
             assert response.status_code == 200, response.text
             assert response.json()["name"] == workspace
 
-            with patch("nmp.common.auth.access_keys.validate_access_key_token", validate_with_local_jwks):
-                authenticate_response = client.get(
-                    "/apis/auth/authenticate",
-                    headers={"Authorization": f"Bearer {access_key}"},
-                )
-                invalid_authenticate_response = client.get(
-                    "/apis/auth/authenticate",
-                    headers={"Authorization": f"Bearer {_tamper_jwt(access_key)}"},
-                )
+            authenticate_response = authenticate_with_access_key(access_key)
+            invalid_authenticate_response = authenticate_with_access_key(_tamper_jwt(access_key))
 
             assert authenticate_response.status_code == 200, authenticate_response.text
             assert authenticate_response.json()["principal"] == user
             assert authenticate_response.json()["token_kind"] == "access_key"
             assert invalid_authenticate_response.status_code == 401, invalid_authenticate_response.text
 
-            with patch("nmp.common.auth.access_keys.validate_access_key_token", validate_with_local_jwks):
-                invalid_workspace_response = client.get(
-                    f"{WORKSPACES_PATH}/{workspace}",
-                    headers={"Authorization": f"Bearer {_tamper_jwt(access_key)}"},
-                )
+            invalid_workspace_response = get_workspace_with_access_key(_tamper_jwt(access_key))
 
             assert invalid_workspace_response.status_code == 401, invalid_workspace_response.text
 
@@ -155,15 +186,14 @@ async def validate_with_local_jwks(config: AuthConfig, token: str) -> TokenClaim
             assert len(revoked_keys) == 1
             assert revoked_keys[0]["status"] == "REVOKED"
 
-            with patch("nmp.common.auth.access_keys.validate_access_key_token", validate_with_local_jwks):
-                revoked_authenticate_response = client.get(
-                    "/apis/auth/authenticate",
-                    headers={"Authorization": f"Bearer {access_key}"},
-                )
-                revoked_workspace_response = client.get(
-                    f"{WORKSPACES_PATH}/{workspace}",
-                    headers={"Authorization": f"Bearer {access_key}"},
-                )
+            revoked_authenticate_response = _wait_for_authorization_response(
+                lambda: authenticate_with_access_key(access_key),
+                expected_status_code=401,
+            )
+            revoked_workspace_response = _wait_for_authorization_response(
+                lambda: get_workspace_with_access_key(access_key),
+                expected_status_code=401,
+            )
 
             assert revoked_authenticate_response.status_code == 401, revoked_authenticate_response.text
             assert revoked_workspace_response.status_code == 401, revoked_workspace_response.text