Added unit test

Author: DannyLiCom

src/xpk/commands/cluster.py

@@ -1464,14 +1464,16 @@ def install_diagon_prerequisites():
   """
   deployment_name = 'kueue-controller-manager'
   namespace_name = 'kueue-system'
+  cert_webhook_deployment_name = 'cert-manager-webhook'
+  cert_webhook_namespace_name = 'cert-manager'
   # is_running = wait_for_cluster_running(args)
   is_running = wait_for_deployment_ready(deployment_name, namespace_name)
   if is_running:
     return_code = install_cert_manager()
     if return_code != 0:
       return return_code
 
-    cert_webhook_ready = check_cert_manager_webhook_status()
+    cert_webhook_ready = wait_for_deployment_ready(cert_webhook_deployment_name, cert_webhook_namespace_name)
     if cert_webhook_ready:    
 
       webhook_package = "mldiagnostics-injection-webhook"
@@ -1554,96 +1556,3 @@ def wait_for_deployment_ready(deployment_name: str, namespace: str, timeout_seco
     xpk_print(f"\nUnexpected API request error while checking deployment status: {e}")
     time.sleep(10)
     return False
-
-
-def wait_for_cluster_running(args, timeout_minutes: int = 30) -> bool:
-  """
-  Polls the GKE Cluster status using gcloud CLI until it enters the RUNNING state.
-
-  Args:
-      args: user provided arguments for running the command.
-      timeout_minutes: Timeout duration in minutes.
-
-  Returns:
-      bool: True if the Cluster successfully enters the RUNNING state, False otherwise.
-  """
-  timeout_seconds = timeout_minutes * 60
-  start_time = time.time()
-  
-  # Construct gcloud command to describe the cluster status
-  command = (
-      'gcloud container clusters describe'
-      f' {args.cluster} --region={zone_to_region(args.zone)} --project={args.project}'
-      " --format='value(status)'"
-  )
-  
-  print(f"Waiting for cluster {args.cluster} ({args.zone}) to enter RUNNING state (using gcloud CLI)...")
-
-  while time.time() - start_time < timeout_seconds:
-    try:
-      # Execute the gcloud command
-      return_code, return_output = run_command_for_value(
-          command, f'Get the status of cluster...'
-      )
-      # Check if gcloud command itself returned an error
-      if return_code != 0:
-        # If Not found error, the cluster does not exist
-        if "Not found" in return_output:
-          xpk_print(f"\nError: Cluster {args.cluster} does not exist in {args.zone}.")
-          return False
-        
-        # Other execution errors, wait and retry
-        xpk_print(f"\nError: gcloud command failed. {return_output}")
-        time.sleep(10)
-        continue
-
-      # Check cluster status returned by gcloud
-      if "RUNNING" in return_output:
-        xpk_print(f"Success: Cluster {args.cluster} status is RUNNING.")
-        return True
-      
-      elif "ERROR" in return_output or "DEGRADED" in return_output:
-        xpk_print(f"Error: Cluster status is {return_output}, creation failed.")
-        return False                
-      else:
-        elapsed_time = int(time.time() - start_time)
-        xpk_print(f"Current status: {return_output}. Elapsed time: {elapsed_time} seconds. Checking again...")
-    except Exception as e:
-      xpk_print(f"\nUnexpected API request error: {e}")
-      time.sleep(10) # Wait longer on unexpected errors
-    
-    # Poll interval
-    time.sleep(30) 
-
-  xpk_print(f"\nTimeout Error: Cluster did not reach RUNNING state within {timeout_minutes} minutes.")
-  return False
-
-def check_cert_manager_webhook_status(timeout_seconds: int = 300) -> bool:
-  """
-  Runs and checks the exit code of kubectl rollout status for a specific deployment.
-
-  Args:
-      timeout_seconds (int): The maximum time to wait for the rollout to complete.
-
-  Returns:
-      bool: True if the rollout status is successful (exit code 0), False otherwise.
-  """
-  # Build the kubectl command
-  kubectl_command = (
-      'kubectl rollout status deployment/cert-manager-webhook -n cert-manager'
-      f' --timeout={timeout_seconds}s'
-  )
-  
-  xpk_print(f"Running command to check deployment status: {kubectl_command}")
-
-  try:
-    return_code, return_output = run_command_for_value(
-        kubectl_command, f'check cert manager...'
-    )
-    if "successfully rolled out" in return_output:
-      xpk_print(f"SUCCESS: Deployment cert-manager-webhook rollout completed.")
-      return True
-
-  except Exception as e:
-    xpk_print(f"\nUnexpected error during kubectl execution: {e}")
-    return False

src/xpk/commands/cluster_test.py

@@ -20,7 +20,7 @@
 from unittest.mock import MagicMock, patch
 import pytest
 
-from xpk.commands.cluster import _install_kueue, _validate_cluster_create_args, run_gke_cluster_create_command
+from xpk.commands.cluster import _install_kueue, _validate_cluster_create_args, run_gke_cluster_create_command, install_diagon_prerequisites
 from xpk.core.system_characteristics import SystemCharacteristics, UserFacingNameToSystemCharacteristics
 from xpk.core.testing.commands_tester import CommandsTester
 from xpk.utils.feature_flags import FeatureFlags
@@ -87,6 +87,7 @@ def construct_args(**kwargs: Any) -> Namespace:
       memory_limit='100Gi',
       cpu_limit=100,
       cluster_cpu_machine_type='',
+      managed_mldiagnostics=False,
   )
   args_dict.update(kwargs)
   return Namespace(**args_dict)
@@ -247,3 +248,13 @@ def test_run_gke_cluster_create_command_with_gke_version_has_no_autoupgrade_flag
   mocks.commands_tester.assert_command_run(
       'clusters create', ' --no-enable-autoupgrade'
   )
+
+
+def test_cluster_create_args_for_install_diagon_prerequisites(
+    mocks: _Mocks,
+):
+
+  install_diagon_prerequisites()
+
+  assert mocks.common_print_mock.call_count == 0
+  

src/xpk/parser/cluster_test.py

@@ -64,3 +64,14 @@ def test_cluster_create_sub_slicing_can_be_set():
   )
 
   assert args.sub_slicing is True
+
+
+def test_cluster_create_managed_mldiagnostics():
+  parser = argparse.ArgumentParser()
+
+  set_cluster_create_parser(parser)
+  args = parser.parse_args(
+      ["--cluster", "test-cluster", "--tpu-type", "v5p-8", "--managed-mldiagnostics"]
+  )
+
+  assert args.managed_mldiagnostics is True