From 5baf4d2bd8aedf9ff2ebd075fbda566760792ae7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Irving=20Mondrag=C3=B3n?= <irving_mondragon@epam.com>
Date: Tue, 19 Nov 2024 17:38:55 +0100
Subject: [PATCH] Allow to switch local cluster

---
 src/xpk/commands/batch.py | 44 ++++++++++++++++++++++++++++++++++++---
 1 file changed, 41 insertions(+), 3 deletions(-)

diff --git a/src/xpk/commands/batch.py b/src/xpk/commands/batch.py
index d82be7e2..9269aa26 100644
--- a/src/xpk/commands/batch.py
+++ b/src/xpk/commands/batch.py
@@ -23,6 +23,7 @@
 from ..core.app_profile import APP_PROFILE_TEMPLATE_DEFAULT_NAME
 from ..core.commands import (
     run_command_for_value,
+    run_command_with_updates,
 )
 
 
@@ -34,11 +35,14 @@ def batch(args: Namespace) -> None:
   Returns:
     None
   """
-  if not args.local_test:
+  if not args.kind_cluster:
     add_zone_and_project(args)
     set_cluster_command_code = set_cluster_command(args)
-    if set_cluster_command_code != 0:
-      xpk_exit(set_cluster_command_code)
+  else:
+    set_cluster_command_code = set_local_cluster_command(args)
+
+  if set_cluster_command_code != 0:
+    xpk_exit(set_cluster_command_code)
 
   create_job_template_instance(args)
   create_app_profile_instance(args)
@@ -57,3 +61,37 @@ def submit_job(args: Namespace) -> None:
   if return_code != 0:
     xpk_print(f'Running batch job returned ERROR {return_code}')
     xpk_exit(return_code)
+
+
+def set_local_cluster_command(args) -> int:
+  """Run local cluster configuration command to set the kubectl config.
+
+  Args:
+    args: user provided arguments for running the command.
+
+  Returns:
+    0 if successful and 1 otherwise.
+  """
+  if not args.cluster:
+    command = 'kubectl config current-context'
+    return_code, current_context = run_command_for_value(
+        command, 'get current-context', args
+    )
+    xpk_print(
+        'No local cluster name specified. Using current-context'
+        f' `{current_context.strip()}`'
+    )
+    return return_code
+
+  command = (
+      f'kubectl config use-context kind-{args.cluster} --namespace=default'
+  )
+  task = f'switch to cluster {args.cluster}'
+  return_code = run_command_with_updates(
+      command,
+      task,
+      args,
+  )
+  if return_code != 0:
+    xpk_print(f'{task} returned ERROR {return_code}')
+  return return_code