Prompts user to install gke auth plugin if missing

Author: JyotinderSingh

keras_remote/cli/prerequisites_check.py

@@ -41,6 +41,25 @@ def check_docker():
     )
 
 
+def check_gke_auth_plugin():
+  """Verify gke-gcloud-auth-plugin is installed."""
+  if not shutil.which("gke-gcloud-auth-plugin"):
+    warning("gke-gcloud-auth-plugin not found.")
+    if click.confirm(
+      "Install it via `gcloud components install gke-gcloud-auth-plugin`?",
+      default=True,
+    ):
+      subprocess.run(
+        ["gcloud", "components", "install", "gke-gcloud-auth-plugin"],
+        check=True,
+      )
+    else:
+      raise click.ClickException(
+        "gke-gcloud-auth-plugin is required. "
+        "Install with: gcloud components install gke-gcloud-auth-plugin"
+      )
+
+
 def check_gcloud_auth():
   """Check if gcloud Application Default Credentials are configured."""
   result = subprocess.run(
@@ -61,5 +80,6 @@ def check_all():
   check_gcloud()
   check_pulumi()
   check_kubectl()
+  check_gke_auth_plugin()
   check_docker()
   check_gcloud_auth()

keras_remote/cli/prerequisites_check_test.py

@@ -9,6 +9,7 @@
   check_docker,
   check_gcloud,
   check_gcloud_auth,
+  check_gke_auth_plugin,
   check_kubectl,
   check_pulumi,
 )
@@ -71,6 +72,40 @@ def test_missing(self, check_fn, error_match):
       check_fn()
 
 
+class TestCheckGkeAuthPlugin(absltest.TestCase):
+  def test_present(self):
+    with mock.patch(
+      "shutil.which", return_value="/usr/bin/gke-gcloud-auth-plugin"
+    ):
+      check_gke_auth_plugin()
+
+  def test_missing_user_accepts_install(self):
+    with (
+      mock.patch("shutil.which", return_value=None),
+      mock.patch("keras_remote.cli.prerequisites_check.warning"),
+      mock.patch("click.confirm", return_value=True),
+      mock.patch(
+        "keras_remote.cli.prerequisites_check.subprocess.run",
+      ) as mock_run,
+    ):
+      check_gke_auth_plugin()
+      mock_run.assert_called_once_with(
+        ["gcloud", "components", "install", "gke-gcloud-auth-plugin"],
+        check=True,
+      )
+
+  def test_missing_user_declines_install(self):
+    with (
+      mock.patch("shutil.which", return_value=None),
+      mock.patch("keras_remote.cli.prerequisites_check.warning"),
+      mock.patch("click.confirm", return_value=False),
+      self.assertRaisesRegex(
+        click.ClickException, "gke-gcloud-auth-plugin is required"
+      ),
+    ):
+      check_gke_auth_plugin()
+
+
 class TestCheckGcloudAuth(absltest.TestCase):
   def test_token_success(self):
     """When print-access-token succeeds, no login is triggered."""