address review comment

Signed-off-by: Swati Mukund Bagal <sbagal@redhat.com>

Author: SB159

tests/model_serving/maas_billing/maas_subscription/conftest.py

@@ -203,6 +203,64 @@ def maas_headers_for_actor_api_key(maas_api_key_for_actor: str) -> dict[str, str
     return build_maas_headers(token=maas_api_key_for_actor)
 
 
+@pytest.fixture(scope="class")
+def api_key_bound_to_free_subscription(
+    request_session_http: requests.Session,
+    base_url: str,
+    ocp_token_for_actor: str,
+    maas_subscription_tinyllama_free: MaaSSubscription,
+    maas_subscription_controller_enabled_latest: None,
+    maas_gateway_api: None,
+    maas_api_gateway_reachable: None,
+) -> Generator[str, Any, Any]:
+    """
+    API key bound to the free subscription at mint time. Revoked on teardown.
+    """
+    _, body = create_api_key(
+        base_url=base_url,
+        ocp_user_token=ocp_token_for_actor,
+        request_session_http=request_session_http,
+        api_key_name=f"e2e-auth-enforce-{generate_random_name()}",
+        subscription=maas_subscription_tinyllama_free.name,
+    )
+    yield body["key"]
+    revoke_api_key(
+        request_session_http=request_session_http,
+        base_url=base_url,
+        key_id=body["id"],
+        ocp_user_token=ocp_token_for_actor,
+    )
+
+
+@pytest.fixture(scope="class")
+def api_key_bound_to_premium_subscription(
+    request_session_http: requests.Session,
+    base_url: str,
+    ocp_token_for_actor: str,
+    maas_subscription_tinyllama_premium: MaaSSubscription,
+    maas_subscription_controller_enabled_latest: None,
+    maas_gateway_api: None,
+    maas_api_gateway_reachable: None,
+) -> Generator[str, Any, Any]:
+    """
+    API key bound to the premium subscription at mint time. Revoked on teardown.
+    """
+    _, body = create_api_key(
+        base_url=base_url,
+        ocp_user_token=ocp_token_for_actor,
+        request_session_http=request_session_http,
+        api_key_name=f"e2e-sub-enforce-{generate_random_name()}",
+        subscription=maas_subscription_tinyllama_premium.name,
+    )
+    yield body["key"]
+    revoke_api_key(
+        request_session_http=request_session_http,
+        base_url=base_url,
+        key_id=body["id"],
+        ocp_user_token=ocp_token_for_actor,
+    )
+
+
 @pytest.fixture(scope="class")
 def maas_wrong_group_service_account_token(
     maas_api_server_url: str,

tests/model_serving/maas_billing/maas_subscription/test_maas_auth_enforcement.py

@@ -6,12 +6,9 @@
 
 from tests.model_serving.maas_billing.maas_subscription.utils import (
     chat_payload_for_url,
-    create_api_key,
     poll_expected_status,
-    revoke_api_key,
 )
 from tests.model_serving.maas_billing.utils import build_maas_headers
-from utilities.general import generate_random_name
 from utilities.plugins.constant import RestHeader
 
 LOGGER = structlog.get_logger(name=__name__)
@@ -35,41 +32,21 @@ class TestMaaSAuthPolicyEnforcementTinyLlama:
     def test_authorized_user_gets_200(
         self,
         request_session_http: requests.Session,
-        base_url: str,
-        ocp_token_for_actor: str,
         model_url_tinyllama_free: str,
-        maas_subscription_tinyllama_free,
+        api_key_bound_to_free_subscription: str,
     ) -> None:
         """
         Verify a free user with a subscription-bound API key can access the free model.
         """
-        _, body = create_api_key(
-            base_url=base_url,
-            ocp_user_token=ocp_token_for_actor,
+        resp = poll_expected_status(
             request_session_http=request_session_http,
-            api_key_name=f"e2e-auth-enforce-{generate_random_name()}",
-            subscription=maas_subscription_tinyllama_free.name,
+            model_url=model_url_tinyllama_free,
+            headers=build_maas_headers(token=api_key_bound_to_free_subscription),
+            payload=chat_payload_for_url(model_url=model_url_tinyllama_free),
+            expected_statuses={200},
         )
-        api_key = body["key"]
-        key_id = body["id"]
-
-        try:
-            resp = poll_expected_status(
-                request_session_http=request_session_http,
-                model_url=model_url_tinyllama_free,
-                headers=build_maas_headers(token=api_key),
-                payload=chat_payload_for_url(model_url=model_url_tinyllama_free),
-                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]}"
-        finally:
-            revoke_api_key(
-                request_session_http=request_session_http,
-                base_url=base_url,
-                key_id=key_id,
-                ocp_user_token=ocp_token_for_actor,
-            )
+        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.smoke
     def test_no_auth_header_gets_401(

tests/model_serving/maas_billing/maas_subscription/test_maas_sub_enforcement.py

@@ -6,11 +6,8 @@
 
 from tests.model_serving.maas_billing.maas_subscription.utils import (
     chat_payload_for_url,
-    create_api_key,
-    revoke_api_key,
 )
 from tests.model_serving.maas_billing.utils import build_maas_headers
-from utilities.general import generate_random_name
 
 LOGGER = structlog.get_logger(name=__name__)
 
@@ -36,40 +33,20 @@ class TestSubscriptionEnforcementTinyLlama:
     def test_subscribed_user_gets_200(
         self,
         request_session_http: requests.Session,
-        base_url: str,
-        ocp_token_for_actor: str,
         model_url_tinyllama_premium: str,
-        maas_subscription_tinyllama_premium,
+        api_key_bound_to_premium_subscription: str,
     ) -> None:
         """
         Verify a premium user with a subscription-bound API key can access the premium model.
         """
-        _, body = create_api_key(
-            base_url=base_url,
-            ocp_user_token=ocp_token_for_actor,
-            request_session_http=request_session_http,
-            api_key_name=f"e2e-sub-enforce-{generate_random_name()}",
-            subscription=maas_subscription_tinyllama_premium.name,
+        resp = request_session_http.post(
+            url=model_url_tinyllama_premium,
+            headers=build_maas_headers(token=api_key_bound_to_premium_subscription),
+            json=chat_payload_for_url(model_url=model_url_tinyllama_premium),
+            timeout=60,
         )
-        api_key = body["key"]
-        key_id = body["id"]
-
-        try:
-            resp = request_session_http.post(
-                url=model_url_tinyllama_premium,
-                headers=build_maas_headers(token=api_key),
-                json=chat_payload_for_url(model_url=model_url_tinyllama_premium),
-                timeout=60,
-            )
-            LOGGER.info(f"test_subscribed_user_gets_200 -> {resp.status_code}")
-            assert resp.status_code == 200, f"Expected 200, got {resp.status_code}: {resp.text[:200]}"
-        finally:
-            revoke_api_key(
-                request_session_http=request_session_http,
-                base_url=base_url,
-                key_id=key_id,
-                ocp_user_token=ocp_token_for_actor,
-            )
+        LOGGER.info(f"test_subscribed_user_gets_200 -> {resp.status_code}")
+        assert resp.status_code == 200, f"Expected 200, got {resp.status_code}: {resp.text[:200]}"
 
     @pytest.mark.parametrize("ocp_token_for_actor", [{"type": "premium"}], indirect=True)
     def test_explicit_subscription_header_works(