From ff76dcb5d7ae9179537b5156f28b991bca79be6e Mon Sep 17 00:00:00 2001
From: Sivanantham Chinnaiyan <sivanantham.chinnaiyan@ideas2it.com>
Date: Thu, 15 Feb 2024 16:06:16 +0530
Subject: [PATCH] Use worksteal dist type for pytest

Signed-off-by: Sivanantham Chinnaiyan <sivanantham.chinnaiyan@ideas2it.com>
---
 .github/actions/minikube-setup/action.yml |  2 +-
 .github/workflows/e2e-test.yml            |  4 ++--
 test/e2e/qpext/test_qpext.py              | 15 ++++++++++++---
 test/scripts/gh-actions/run-e2e-tests.sh  |  2 +-
 4 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/.github/actions/minikube-setup/action.yml b/.github/actions/minikube-setup/action.yml
index 17dc295074d..f8ed073551e 100644
--- a/.github/actions/minikube-setup/action.yml
+++ b/.github/actions/minikube-setup/action.yml
@@ -5,7 +5,7 @@ runs:
   using: "composite"
   steps:
     - name: Install kubectl
-      uses: azure/setup-kubectl@v4
+      uses: azure/setup-kubectl@v4.0.0
       with:
         version: 'v1.27.4'
 
diff --git a/.github/workflows/e2e-test.yml b/.github/workflows/e2e-test.yml
index d3204fa981a..538a01fbcbb 100644
--- a/.github/workflows/e2e-test.yml
+++ b/.github/workflows/e2e-test.yml
@@ -716,7 +716,7 @@ jobs:
         uses: ./.github/actions/minikube-setup
 
       - name: KServe dependency setup
-        uses: ./.github/actions/kserve-dep-setup\
+        uses: ./.github/actions/kserve-dep-setup
         with:
           deployment-mode: "raw"
 
@@ -790,7 +790,7 @@ jobs:
         uses: ./.github/actions/minikube-setup
 
       - name: KServe dependency setup
-        uses: ./.github/actions/kserve-dep-setup\
+        uses: ./.github/actions/kserve-dep-setup
         with:
           network-layer: "kourier"
 
diff --git a/test/e2e/qpext/test_qpext.py b/test/e2e/qpext/test_qpext.py
index 18e5acff570..5dcc7fd6eec 100644
--- a/test/e2e/qpext/test_qpext.py
+++ b/test/e2e/qpext/test_qpext.py
@@ -26,7 +26,7 @@
     V1beta1SKLearnSpec,
 )
 from kubernetes.client import V1ResourceRequirements
-from ..common.utils import KSERVE_TEST_NAMESPACE
+from ..common.utils import KSERVE_TEST_NAMESPACE, get_cluster_ip
 from ..common.utils import predict
 
 logging.basicConfig(level=logging.INFO)
@@ -39,11 +39,12 @@
 def test_qpext_kserve():
     # test the qpext using the sklearn predictor
     service_name = "sklearn-v2-metrics"
+    protocol_version = "v2"
     predictor = V1beta1PredictorSpec(
         min_replicas=1,
         sklearn=V1beta1SKLearnSpec(
             storage_uri="gs://seldon-models/sklearn/mms/lr_model",
-            protocol_version="v2",
+            protocol_version=protocol_version,
             resources=V1ResourceRequirements(
                 requests={"cpu": "50m", "memory": "128Mi"},
                 limits={"cpu": "100m", "memory": "512Mi"},
@@ -65,8 +66,16 @@ def test_qpext_kserve():
     kserve_client = KServeClient(config_file=os.environ.get("KUBECONFIG", "~/.kube/config"))
     kserve_client.create(isvc)
     kserve_client.wait_isvc_ready(service_name, namespace=KSERVE_TEST_NAMESPACE)
+    kserve_client.wait_model_ready(
+        service_name,
+        model_name=service_name,
+        isvc_namespace=KSERVE_TEST_NAMESPACE,
+        isvc_version=constants.KSERVE_V1BETA1_VERSION,
+        protocol_version=protocol_version,
+        cluster_ip=get_cluster_ip(),
+    )
 
-    res = predict(service_name, "./data/iris_input_v2.json", protocol_version="v2")
+    res = predict(service_name, "./data/iris_input_v2.json", protocol_version=protocol_version)
     assert res["outputs"][0]["data"] == [1, 1]
 
     send_metrics_request(kserve_client, service_name)
diff --git a/test/scripts/gh-actions/run-e2e-tests.sh b/test/scripts/gh-actions/run-e2e-tests.sh
index 17443f12a40..292a301d4f1 100755
--- a/test/scripts/gh-actions/run-e2e-tests.sh
+++ b/test/scripts/gh-actions/run-e2e-tests.sh
@@ -30,5 +30,5 @@ fi
 PARALLELISM="${2:-1}"
 source python/kserve/.venv/bin/activate
 pushd test/e2e >/dev/null
-  pytest -m "$1" --ignore=qpext --log-level=INFO -o log_cli=true -n $PARALLELISM
+  pytest -m "$1" --ignore=qpext --log-level=INFO -o log_cli=true -n $PARALLELISM --dist worksteal
 popd