From b6cbdecbf899c49a87eeb776350ad75cabdea04d Mon Sep 17 00:00:00 2001
From: Himanshu Sachdeva <hisachdeva@google.com>
Date: Thu, 14 Mar 2024 20:07:25 +0100
Subject: [PATCH] fix jupyter test script args

---
 cloudbuild-ap.yaml                             | 17 +++++++++--------
 modules/jupyter/tests/change_jupyter_config.py |  2 +-
 modules/jupyter/tests/test_hub.py              |  2 +-
 3 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/cloudbuild-ap.yaml b/cloudbuild-ap.yaml
index 2d46d9d1e..5b0af464f 100644
--- a/cloudbuild-ap.yaml
+++ b/cloudbuild-ap.yaml
@@ -136,14 +136,14 @@ steps:
         set -e
 
         cd /workspace/modules/jupyter/tests
-        python3 change_jupyter_config.py autopilot
+        python3 change_jupyter_config.py $_AUTOPILOT_CLUSTER
 
         cd /workspace/applications/jupyter
         terraform apply \
         -var-file=workloads-without-iap.example.tfvars \
         -var=project_id=$PROJECT_ID \
         -var=cluster_name=ml-$SHORT_SHA-$_PR_NUMBER-$_SUFFIX-cluster \
-        -var=namespace=ml-$SHORT_SHA-$_SUFFIX \
+        -var=kubernetes_namespace=ml-$SHORT_SHA-$_SUFFIX \
         -var=workload_identity_service_account=jupyter-sa-$SHORT_SHA-$_SUFFIX \
         -var=gcs_bucket=gke-aieco-jupyter-$SHORT_SHA-$_SUFFIX \
         -auto-approve -no-color -lock=false
@@ -155,7 +155,7 @@ steps:
         echo "HOST URL is " $(cat /workspace/jupyterhub_host_url.txt)
 
         cd /workspace/modules/jupyter/tests
-        python3 test_hub.py $(cat /workspace/jupyterhub_host_url.txt) autopilot
+        python3 test_hub.py $(cat /workspace/jupyterhub_host_url.txt) $_AUTOPILOT_CLUSTER
         echo "pass" > /workspace/jupyterhub_test_result.txt
     allowFailure: true
 
@@ -172,10 +172,10 @@ steps:
         -var-file=workloads-without-iap.example.tfvars \
         -var=project_id=$PROJECT_ID \
         -var=cluster_name=ml-$SHORT_SHA-$_PR_NUMBER-$_SUFFIX-cluster \
-        -var=namespace=ml-$SHORT_SHA-$_SUFFIX \
+        -var=kubernetes_namespace=ml-$SHORT_SHA-$_SUFFIX \
         -var=workload_identity_service_account=jupyter-sa-$SHORT_SHA-$_SUFFIX \
         -var=gcs_bucket=gke-aieco-jupyter-$SHORT_SHA-$_SUFFIX \
-        -auto-approve -no-color
+        -auto-approve -no-color -lock=false
 
     allowFailure: true
     waitFor: ['test jupyterhub']
@@ -195,7 +195,7 @@ steps:
         --project $PROJECT_ID 
 
         cd /workspace/modules/jupyter/tests
-        python3 change_jupyter_config.py autopilot
+        python3 change_jupyter_config.py $_AUTOPILOT_CLUSTER
 
         cd /workspace/applications/rag/
         terraform apply \
@@ -231,7 +231,7 @@ steps:
 
         # Validate Jupyterhub: Test Hub
         cd /workspace/modules/jupyter/tests
-        python3 test_hub.py $(cat /workspace/rag_jupyterhub_host_url.txt) autopilot
+        python3 test_hub.py $(cat /workspace/rag_jupyterhub_host_url.txt) $_AUTOPILOT_CLUSTER
         echo "pass" > /workspace/rag_jupyterhub_test_result.txt
 
         # Validate RAG: Test rag frontend
@@ -280,6 +280,7 @@ steps:
         cd /workspace/infrastructure
         terraform destroy -var-file=tfvars_tests/standard-gke-public.platform.tfvars -var=project_id=$PROJECT_ID \
         -var=cluster_name=ml-$SHORT_SHA-$_PR_NUMBER-$_SUFFIX-cluster \
+        -var=autopilot_cluster=$_AUTOPILOT_CLUSTER \
         -var=cluster_location=$_REGION -auto-approve -no-color
 
     allowFailure: true
@@ -331,7 +332,7 @@ steps:
           exit 1
         fi
 
-    waitFor: ['cleanup rag']
+    waitFor: ['cleanup gke cluster']
         
 substitutions:
   _REGION: us-central1
diff --git a/modules/jupyter/tests/change_jupyter_config.py b/modules/jupyter/tests/change_jupyter_config.py
index 302d24763..b84a30030 100644
--- a/modules/jupyter/tests/change_jupyter_config.py
+++ b/modules/jupyter/tests/change_jupyter_config.py
@@ -3,7 +3,7 @@
 
 config_file = "../jupyter_config/config-selfauth.yaml"
 if len(sys.argv) == 2:
-    autopilot = (sys.argv[1] == "autopilot")
+    autopilot = (sys.argv[1] == "true")
     if autopilot:
         config_file = "../jupyter_config/config-selfauth-autopilot.yaml"
 
diff --git a/modules/jupyter/tests/test_hub.py b/modules/jupyter/tests/test_hub.py
index 23f23739e..4c06daba6 100644
--- a/modules/jupyter/tests/test_hub.py
+++ b/modules/jupyter/tests/test_hub.py
@@ -6,7 +6,7 @@
 
 config_file = "../jupyter_config/config-selfauth.yaml"
 if len(sys.argv) == 3:
-    autopilot = (sys.argv[2] == "autopilot")
+    autopilot = (sys.argv[2] == "true")
     if autopilot:
         config_file = "../jupyter_config/config-selfauth-autopilot.yaml"