[WIP] add test for Raw kserve auth

Signed-off-by: Milind Waykole <[email protected]>

tests/model_serving/model_server/authentication/conftest.py

@@ -88,7 +88,27 @@ def grpc_s3_inference_service(
         enable_auth=True,
     ) as isvc:
         yield isvc
-
+@pytest.fixture(scope="class")
+def grpc_s3_raw_inference_service(
+    admin_client: DynamicClient,
+    model_namespace: Namespace,
+    grpc_s3_caikit_serving_runtime: ServingRuntime,
+    s3_models_storage_uri: str,
+    grpc_model_service_account: ServiceAccount,
+) -> InferenceService:
+    with create_isvc(
+        client=admin_client,
+        name=f"{Protocols.GRPC}-{ModelFormat.CAIKIT}",
+        namespace=model_namespace.name,
+        runtime=grpc_s3_caikit_serving_runtime.name,
+        storage_uri=s3_models_storage_uri,
+        model_format=grpc_s3_caikit_serving_runtime.instance.spec.supportedModelFormats[0].name,
+        deployment_mode=KServeDeploymentType.RAW_DEPLOYMENT,
+        model_service_account=grpc_model_service_account.name,
+        enable_auth=True,
+        external_route=True,
+    ) as isvc:
+        yield isvc
 
 @pytest.fixture(scope="class")
 def http_view_role(
@@ -103,6 +123,18 @@ def http_view_role(
     ) as role:
         yield role
 
+@pytest.fixture(scope="class")
+def http_raw_view_role(
+    admin_client: DynamicClient,
+    http_s3_caikit_raw_inference_service: InferenceService,
+) -> Role:
+    with create_isvc_view_role(
+        client=admin_client,
+        isvc=http_s3_caikit_raw_inference_service,
+        name=f"{http_s3_caikit_raw_inference_service.name}-view",
+        resource_names=[http_s3_caikit_raw_inference_service.name],
+    ) as role:
+        yield role
 
 @pytest.fixture(scope="class")
 def http_role_binding(
@@ -122,13 +154,36 @@ def http_role_binding(
     ) as rb:
         yield rb
 
+@pytest.fixture(scope="class")
+def http_raw_role_binding(
+    admin_client: DynamicClient,
+    http_raw_view_role: Role,
+    model_service_account: ServiceAccount,
+    http_s3_caikit_raw_inference_service: InferenceService,
+) -> RoleBinding:
+    with RoleBinding(
+        client=admin_client,
+        namespace=model_service_account.namespace,
+        name=f"{Protocols.HTTP}-{model_service_account.name}-view",
+        role_ref_name=http_raw_view_role.name,
+        role_ref_kind=http_raw_view_role.kind,
+        subjects_kind=model_service_account.kind,
+        subjects_name=model_service_account.name,
+    ) as rb:
+        yield rb
+
 
 @pytest.fixture(scope="class")
 def http_inference_token(model_service_account: ServiceAccount, http_role_binding: RoleBinding) -> str:
     return run_command(
         command=shlex.split(f"oc create token -n {model_service_account.namespace} {model_service_account.name}")
     )[1].strip()
 
+@pytest.fixture(scope="class")
+def http_raw_inference_token(model_service_account: ServiceAccount, http_raw_role_binding: RoleBinding) -> str:
+    return run_command(
+        command=shlex.split(f"oc create token -n {model_service_account.namespace} {model_service_account.name}")
+    )[1].strip()
 
 @pytest.fixture()
 def patched_remove_authentication_isvc(
@@ -151,6 +206,27 @@ def patched_remove_authentication_isvc(
         predictor_pod.wait_deleted()
 
         yield http_s3_caikit_serverless_inference_service
+@pytest.fixture()
+def patched_remove_authentication_isvc(
+    admin_client: DynamicClient,
+    http_s3_caikit_serverless_inference_service: InferenceService,
+) -> InferenceService:
+    with ResourceEditor(
+        patches={
+            http_s3_caikit_raw_inference_service: {
+                "metadata": {
+                    "labels": {"security.opendatahub.io/enable-auth": "false"},
+                }
+            }
+        }
+    ):
+        predictor_pod = get_pods_by_isvc_label(
+            client=admin_client,
+            isvc=http_s3_caikit_raw_inference_service,
+        )[0]
+        predictor_pod.wait_deleted()
+
+        yield http_s3_caikit_raw_inference_service
 
 
 @pytest.fixture(scope="class")
@@ -162,7 +238,15 @@ def grpc_view_role(admin_client: DynamicClient, grpc_s3_inference_service: Infer
         resource_names=[grpc_s3_inference_service.name],
     ) as role:
         yield role
-
+@pytest.fixture(scope="class")
+def grpc_raw_view_role(admin_client: DynamicClient, grpc_s3_raw_inference_service: InferenceService) -> Role:
+    with create_isvc_view_role(
+        client=admin_client,
+        isvc=grpc_s3_raw_inference_service,
+        name=f"{grpc_s3_raw_inference_service.name}-view",
+        resource_names=[grpc_s3_raw_inference_service.name],
+    ) as role:
+        yield role
 
 @pytest.fixture(scope="class")
 def grpc_role_binding(
@@ -183,6 +267,23 @@ def grpc_role_binding(
         yield rb
 
 
+@pytest.fixture(scope="class")
+def grpc_raw_role_binding(
+    admin_client: DynamicClient,
+    grpc_raw_view_role: Role,
+    grpc_model_service_account: ServiceAccount,
+    grpc_s3_inference_service: InferenceService,
+) -> RoleBinding:
+    with RoleBinding(
+        client=admin_client,
+        namespace=grpc_model_service_account.namespace,
+        name=f"{Protocols.GRPC}-{grpc_model_service_account.name}-view",
+        role_ref_name=grpc_raw_view_role.name,
+        role_ref_kind=grpc_raw_view_role.kind,
+        subjects_kind=grpc_model_service_account.kind,
+        subjects_name=grpc_model_service_account.name,
+    ) as rb:
+        yield rb
 @pytest.fixture(scope="class")
 def grpc_inference_token(grpc_model_service_account: ServiceAccount, grpc_role_binding: RoleBinding) -> str:
     return run_command(
@@ -191,6 +292,13 @@ def grpc_inference_token(grpc_model_service_account: ServiceAccount, grpc_role_b
         )
     )[1].strip()
 
+@pytest.fixture(scope="class")
+def grpc_raw_inference_token(grpc_model_service_account: ServiceAccount, grpc_role_binding: RoleBinding) -> str:
+    return run_command(
+        command=shlex.split(
+            f"oc create token -n {grpc_model_service_account.namespace} {grpc_model_service_account.name}"
+        )
+    )[1].strip()
 
 @pytest.fixture(scope="class")
 def http_s3_caikit_serverless_inference_service(
@@ -215,6 +323,31 @@ def http_s3_caikit_serverless_inference_service(
         yield isvc
 
 
+@pytest.fixture(scope="class")
+def http_s3_caikit_raw_inference_service(
+    request: FixtureRequest,
+    admin_client: DynamicClient,
+    model_namespace: Namespace,
+    http_s3_caikit_tgis_serving_runtime: ServingRuntime,
+    s3_models_storage_uri: str,
+    model_service_account: ServiceAccount,
+) -> InferenceService:
+    with create_isvc(
+        client=admin_client,
+        name=f"{Protocols.HTTP}-{ModelFormat.CAIKIT}",
+        namespace=model_namespace.name,
+        runtime=http_s3_caikit_tgis_serving_runtime.name,
+        storage_uri=s3_models_storage_uri,
+        model_format=http_s3_caikit_tgis_serving_runtime.instance.spec.supportedModelFormats[0].name,
+        deployment_mode=KServeDeploymentType.RAW_DEPLOYMENT,
+        model_service_account=model_service_account.name,
+        enable_auth=True,
+        external_route=True,
+    ) as isvc:
+        yield isvc
+
+
+
 # Unprivileged user tests
 @pytest.fixture(scope="class")
 def unprivileged_model_namespace(

tests/model_serving/model_server/authentication/test_kserve_token_authentication_raw.py

@@ -0,0 +1,59 @@
+import pytest
+
+from tests.model_serving.model_server.utils import verify_inference_response
+from utilities.constants import ModelFormat, ModelStoragePath, Protocols, ModelInferenceRuntime
+from utilities.inference_utils import Inference
+
+pytestmark = pytest.mark.usefixtures("valid_aws_config")
+
+
+@pytest.mark.serverless
+@pytest.mark.parametrize(
+    "model_namespace, s3_models_storage_uri",
+    [
+        pytest.param(
+            {"name": "kserve-token-authentication"},
+            {"model-dir": ModelStoragePath.FLAN_T5_SMALL},
+        )
+    ],
+    indirect=True,
+)
+class TestKserveTokenAuthenticationRaw:
+    @pytest.mark.smoke
+    @pytest.mark.dependency(name="test_model_authentication_using_rest_raw")
+    def test_model_authentication_using_rest_raw(self, http_s3_caikit_raw_inference_service, http_raw_inference_token):
+        """Verify RAW Kserve model query with token using REST"""
+        verify_inference_response(
+            inference_service=http_s3_caikit_raw_inference_service,
+            runtime=ModelInferenceRuntime.CAIKIT_TGIS_RUNTIME,
+            inference_type=Inference.ALL_TOKENS,
+            protocol=Protocols.HTTPS,
+            model_name=ModelFormat.CAIKIT,
+            use_default_query=True,
+            token=http_raw_inference_token,
+        )
+
+    @pytest.mark.smoke
+    def test_model_raw_authentication_using_grpc(self, grpc_s3_raw_inference_service, grpc_raw_inference_token):
+        """Verify model query with token using GRPC"""
+        verify_inference_response(
+            inference_service=grpc_s3_raw_inference_service,
+            runtime=ModelInferenceRuntime.CAIKIT_TGIS_RUNTIME,
+            inference_type=Inference.STREAMING,
+            protocol=Protocols.GRPC,
+            model_name=ModelFormat.CAIKIT,
+            use_default_query=True,
+            token=grpc_raw_inference_token,
+        )
+
+    @pytest.mark.dependency(name="test_disabled_raw_model_authentication")
+    def test_disabled_raw_model_authentication(self, patched_remove_authentication_isvc):
+        """Verify model query after authentication is disabled"""
+        verify_inference_response(
+            inference_service=patched_remove_authentication_isvc,
+            runtime=ModelInferenceRuntime.CAIKIT_TGIS_RUNTIME,
+            inference_type=Inference.ALL_TOKENS,
+            protocol=Protocols.HTTP,
+            model_name=ModelFormat.CAIKIT,
+            use_default_query=True,
+        )

tests/model_serving/model_server/utils.py

@@ -11,9 +11,11 @@
 from tests.model_serving.model_server.private_endpoint.utils import (
     InvalidStorageArgument,
 )
+from tenacity import retry, wait_exponential, stop_after_attempt
 from utilities.constants import KServeDeploymentType
 from utilities.inference_utils import UserInference
 from utilities.infra import wait_for_inference_deployment_replicas
+from utilities.plugins.openai_plugin import MAX_RETRIES
 
 LOGGER = get_logger(name=__name__)
 
@@ -90,7 +92,7 @@ def create_isvc(
         if deployment_mode == KServeDeploymentType.SERVERLESS:
             annotations["security.opendatahub.io/enable-auth"] = "true"
         elif deployment_mode == KServeDeploymentType.RAW_DEPLOYMENT:
-            labels["security.openshift.io/enable-authentication"] = "true"
+            labels["security.opendatahub.io/enable-auth"] = "true"
 
     # default to True if deployment_mode is Serverless (default behavior of Serverless) if was not provided by the user
     if external_route is None and deployment_mode == KServeDeploymentType.SERVERLESS:
@@ -133,6 +135,7 @@ def _check_storage_arguments(
     if (storage_uri and storage_path) or (not storage_uri and not storage_key) or (storage_key and not storage_path):
         raise InvalidStorageArgument(storage_uri, storage_key, storage_path)
 
+@retry(stop=stop_after_attempt(5), wait=wait_exponential(min=1, max=6))
 
 def verify_inference_response(
     inference_service: InferenceService,