test(maas): update auth enforcement tests for API key flow (#1197)

Author: SB159

tests/model_serving/model_server/maas_billing/maas_subscription/test_maas_auth_enforcement.py

@@ -4,7 +4,10 @@
 import requests
 from simple_logger.logger import get_logger
 
-from tests.model_serving.model_server.maas_billing.maas_subscription.utils import chat_payload_for_url
+from tests.model_serving.model_server.maas_billing.maas_subscription.utils import (
+    chat_payload_for_url,
+    poll_expected_status,
+)
 from tests.model_serving.model_server.maas_billing.utils import build_maas_headers
 from utilities.plugins.constant import RestHeader
 
@@ -24,44 +27,46 @@
     "maas_subscription_tinyllama_premium",
 )
 class TestMaaSAuthPolicyEnforcementTinyLlama:
-    @pytest.mark.sanity
+    @pytest.mark.smoke
     @pytest.mark.parametrize("ocp_token_for_actor", [{"type": "free"}], indirect=True)
     def test_authorized_user_gets_200(
         self,
         request_session_http: requests.Session,
         model_url_tinyllama_free: str,
-        ocp_token_for_actor: str,
+        maas_headers_for_actor_api_key: dict[str, str],
     ) -> None:
-        headers = build_maas_headers(token=ocp_token_for_actor)
         payload = chat_payload_for_url(model_url=model_url_tinyllama_free)
 
-        resp = request_session_http.post(
-            url=model_url_tinyllama_free,
-            headers=headers,
-            json=payload,
-            timeout=60,
+        resp = poll_expected_status(
+            request_session_http=request_session_http,
+            model_url=model_url_tinyllama_free,
+            headers=maas_headers_for_actor_api_key,
+            payload=payload,
+            expected_statuses={200},
         )
+
         LOGGER.info(f"test_authorized_user_gets_200 -> POST {model_url_tinyllama_free} returned {resp.status_code}")
         assert resp.status_code == 200, f"Expected 200, got {resp.status_code}: {resp.text[:200]}"
 
-    @pytest.mark.sanity
+    @pytest.mark.smoke
     def test_no_auth_header_gets_401(
         self,
         request_session_http: requests.Session,
         model_url_tinyllama_free: str,
     ) -> None:
         payload = chat_payload_for_url(model_url=model_url_tinyllama_free)
 
-        resp = request_session_http.post(
-            url=model_url_tinyllama_free,
+        resp = poll_expected_status(
+            request_session_http=request_session_http,
+            model_url=model_url_tinyllama_free,
             headers=RestHeader.HEADERS,
-            json=payload,
-            timeout=60,
+            payload=payload,
+            expected_statuses={401},
         )
         LOGGER.info(f"test_no_auth_header_gets_401 -> POST {model_url_tinyllama_free} returned {resp.status_code}")
         assert resp.status_code == 401, f"Expected 401, got {resp.status_code}: {resp.text[:200]}"
 
-    @pytest.mark.sanity
+    @pytest.mark.smoke
     def test_invalid_token_gets_401(
         self,
         request_session_http: requests.Session,
@@ -76,10 +81,11 @@ def test_invalid_token_gets_401(
             json=payload,
             timeout=60,
         )
+
         LOGGER.info(f"test_invalid_token_gets_401 -> POST {model_url_tinyllama_free} returned {resp.status_code}")
-        assert resp.status_code == 401, f"Expected 401, got {resp.status_code}: {resp.text[:200]}"
+        assert resp.status_code in (401, 403), f"Expected 401 or 403, got {resp.status_code}: {(resp.text or '')[:200]}"
 
-    @pytest.mark.sanity
+    @pytest.mark.smoke
     def test_wrong_group_sa_denied_on_premium_model(
         self,
         request_session_http: requests.Session,
@@ -88,11 +94,12 @@ def test_wrong_group_sa_denied_on_premium_model(
     ) -> None:
         payload = chat_payload_for_url(model_url=model_url_tinyllama_premium)
 
-        resp = request_session_http.post(
-            url=model_url_tinyllama_premium,
+        resp = poll_expected_status(
+            request_session_http=request_session_http,
+            model_url=model_url_tinyllama_premium,
             headers=maas_headers_for_wrong_group_sa,
-            json=payload,
-            timeout=60,
+            payload=payload,
+            expected_statuses={403},
         )
         LOGGER.info(
             "test_wrong_group_sa_denied_on_premium_model -> "

tests/model_serving/model_server/maas_billing/maas_subscription/utils.py

@@ -85,7 +85,6 @@ def chat_payload_for_url(model_url: str, *, prompt: str = "Hello", max_tokens: i
 
 
 def poll_expected_status(
-    *,
     request_session_http: requests.Session,
     model_url: str,
     headers: dict[str, str],