refactor: Replace get_user_input with ask_for_user_consent and merge --quiet with --force. (#811)

* feat: Add --quiet flag

* Bump default Kueue version to 0.14.2

* feat: Prepare for Kueue upgrade

* Add user_input_test.py

* testing: CommandsTester updates

* refactor: Simplify kueue_manager_test with CommandsTester

* Add user_input docs

* Add unit tests

* Kueue v0.14.3

* Apply review feedback

* Fix import

* refactor: Replace `get_user_input` with `ask_for_user_consent` and merge --quiet with --force.

* nits

* not node_pools_to_update_WI

* Apply SikaGrr suggestions

Author: jamOne-

src/xpk/commands/cluster.py

@@ -71,9 +71,9 @@
 )
 from ..core.vertex import create_vertex_tensorboard
 from ..core.workload import get_workload_list
-from ..utils.console import get_user_input, xpk_exit, xpk_print
+from ..utils.console import ask_for_user_consent, xpk_exit, xpk_print
 from ..utils.file import write_tmp_file
-from ..utils.execution_context import is_dry_run
+from ..utils.execution_context import is_dry_run, is_quiet
 from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
 from . import cluster_gcluster
 from .common import set_cluster_command, validate_sub_slicing_system
@@ -1056,7 +1056,7 @@ def run_gke_cluster_delete_command(args) -> int:
   Returns:
     0 if successful and 1 otherwise.
   """
-  if not args.force:
+  if not is_quiet():
     xpk_print('Get the name of the workloads in the cluster.')
     args.filter_by_status = 'EVERYTHING'
     return_code, return_value = get_workload_list(args)
@@ -1067,10 +1067,9 @@ def run_gke_cluster_delete_command(args) -> int:
     # Ignore Column Names line.
     if len(return_value) > 1:
       workloads = [x.split(' ')[0] for x in return_value.splitlines()][1:]
-      if workloads and not get_user_input(
+      if workloads and not ask_for_user_consent(
           f'Planning to delete {len(workloads)} workloads in the cluster'
-          f' {args.cluster} including {workloads}. \nDo you wish to delete: y'
-          ' (yes) / n (no):\n'
+          f' {args.cluster} including {workloads}. \nDo you wish to delete?'
       ):
         xpk_print('Skipping delete command.')
         return 0

src/xpk/commands/storage.py

@@ -56,7 +56,7 @@
     list_storages,
     print_storages_for_cluster,
 )
-from ..utils.console import get_user_input, xpk_exit, xpk_print
+from ..utils.console import ask_for_user_consent, xpk_exit, xpk_print
 from ..utils.kubectl import apply_kubectl_manifest
 from ..utils.execution_context import is_dry_run
 from ..utils.validation import validate_dependencies_list, SystemDependency, should_validate_dependencies
@@ -133,15 +133,13 @@ def storage_delete(args: Namespace) -> None:
       if storage.bucket.startswith(filestore_instance_name)
   ]
 
-  if children and not args.force:
-    detach = get_user_input(
-        "Deleting a filestore storage will destroy your filestore instance and"
-        " all its data in all volumes will be lost. Do you wish to delete the"
-        f" filestore instance {filestore_instance_name}?\n y (yes) / n (no):\n'"
-    )
-    if not detach:
-      xpk_print("Deleting storage canceled.")
-      xpk_exit(0)
+  if children and not ask_for_user_consent(
+      "Deleting a filestore storage will destroy your filestore instance and"
+      " all its data in all volumes will be lost. Do you wish to delete the"
+      f" filestore instance {filestore_instance_name}?"
+  ):
+    xpk_print("Deleting storage canceled.")
+    xpk_exit(0)
 
   for child in children:
     delete_storage_resources(k8s_api_client, child)

src/xpk/commands/workload.py

@@ -94,7 +94,7 @@
     tcpx_decorator,
     tcpxo_decorator,
 )
-from ..utils.console import get_user_input, xpk_exit, xpk_print
+from ..utils.console import ask_for_user_consent, xpk_exit, xpk_print
 from packaging.version import Version
 from ..utils.file import write_tmp_file
 from ..utils.execution_context import is_dry_run
@@ -781,11 +781,10 @@ def workload_delete(args) -> None:
       xpk_exit(return_code)
     # Skip the header
     workloads = [x.split(' ')[0] for x in return_value.splitlines()][1:]
-    if workloads and not args.force:
-      will_delete = get_user_input(
+    if workloads:
+      will_delete = ask_for_user_consent(
           f'Planning to delete {len(workloads)} workloads in the cluster'
-          f' {args.cluster} including {workloads}. \nDo you wish to delete: y'
-          ' (yes) / n (no):\n'
+          f' {args.cluster} including {workloads}. \nDo you wish to delete?'
       )
   else:
     workloads = [args.workload]

src/xpk/core/kueue_manager.py

@@ -21,7 +21,6 @@
 import json
 from jinja2 import Environment, FileSystemLoader
 
-from ..utils.user_input import ask_for_user_consent
 from ..utils.execution_context import is_dry_run
 from ..utils.kueue import is_queued_cluster
 from kubernetes.utils import parse_quantity
@@ -41,7 +40,7 @@
     run_command_with_updates_retry,
 )
 from ..utils.file import write_tmp_file
-from ..utils.console import xpk_print, xpk_exit
+from ..utils.console import xpk_print, xpk_exit, ask_for_user_consent
 from ..utils.templates import TEMPLATE_PATH, get_templates_absolute_path
 from packaging.version import Version
 

src/xpk/core/nodepool.py

@@ -15,7 +15,7 @@
 """
 
 from typing import List
-from ..utils.console import get_user_input, xpk_print
+from ..utils.console import ask_for_user_consent, xpk_print
 from ..utils.topology import get_topology_product, is_topology_valid
 from .capacity import (
     AUTOPROVISIONING_CONFIG_VALUE,
@@ -110,6 +110,7 @@ def run_gke_node_pool_create_command(
       existing_node_pool_names, args.cluster, desired_node_pool_count
   )
 
+  node_pools_to_delete = []
   node_pools_to_remain = []
   delete_commands = []
   delete_task_names = []
@@ -186,14 +187,10 @@ def run_gke_node_pool_create_command(
   # when cluster is getting updated from 'x' device_type/gke_accelerator to 'y' device_type/gke_accelerator.
   # In that case, '{args.cluster}-np-i' nodepool will be re-created for 'y' device_type/gke_accelerator.
   if delete_commands:
-    will_delete = True
-    if node_pools_to_delete and not args.force:
-      will_delete = get_user_input(
-          f'Planning to delete {len(node_pools_to_delete)} node pools including'
-          f' {node_pools_to_delete}. \nDo you wish to delete: y (yes) / n'
-          ' (no):\n'
-      )
-    if not will_delete:
+    if node_pools_to_delete and not ask_for_user_consent(
+        f'Planning to delete {len(node_pools_to_delete)} node pools including'
+        f' {node_pools_to_delete}. \nDo you wish to delete?'
+    ):
       xpk_print(
           'You have requested to not delete the existing nodepools in the'
           ' cluster. There will be no change to the cluster.'
@@ -215,18 +212,15 @@ def run_gke_node_pool_create_command(
 
   # Enable Workload Identity on existing Nodepools
   if update_WI_commands:
-    will_update_WI = True
-    if node_pools_to_update_WI and not args.force:
-      will_update_WI = get_user_input(
-          'Planning to enable Workload Identity Federation on'
-          f' {len(node_pools_to_update_WI)} existing node pools including'
-          f' {node_pools_to_update_WI}.This immediately enables Workload'
-          ' Identity Federation for GKE for any workloads running in the node'
-          ' pool. Also, xpk does not support disabling Workload Identity on'
-          ' clusters that have it enabled already \nDo you wish to update: y'
-          ' (yes) / n (no):\n'
-      )
-    if not will_update_WI:
+    will_update_WI = not node_pools_to_update_WI or ask_for_user_consent(
+        'Planning to enable Workload Identity Federation on'
+        f' {len(node_pools_to_update_WI)} existing node pools including'
+        f' {node_pools_to_update_WI}. This immediately enables Workload'
+        ' Identity Federation for GKE for any workloads running in the node'
+        ' pool. Also, xpk does not support disabling Workload Identity on'
+        ' clusters that have it enabled already \nDo you wish to update?'
+    )
+    if will_update_WI:
       for i, command in enumerate(update_WI_commands):
         xpk_print(
             f'To complete {update_WI_task_names[i]} we are executing {command}'

src/xpk/core/nodepool_test.py

@@ -145,7 +145,7 @@ def mock_nodepool_dependencies(mocker):
       "xpk.core.nodepool.get_cluster_location", return_value="us-central1"
   )
   mocker.patch("xpk.core.nodepool.run_commands", return_value=0)
-  mocker.patch("xpk.core.nodepool.get_user_input", return_value=True)
+  mocker.patch("xpk.core.nodepool.ask_for_user_consent", return_value=True)
   mock_is_topology_valid = mocker.patch("xpk.core.nodepool.is_topology_valid")
   mock_ensure_resource_policy = mocker.patch(
       "xpk.core.nodepool.ensure_resource_policy_exists"

src/xpk/main.py

@@ -70,7 +70,10 @@ def main() -> None:
   main_args.enable_ray_cluster = False
   set_context(
       dry_run_value='dry_run' in main_args and main_args.dry_run,
-      quiet_value='quiet' in main_args and main_args.quiet,
+      quiet_value=(
+          ('quiet' in main_args and main_args.quiet)
+          or ('force' in main_args and main_args.force)
+      ),
   )
   generate_client_id()
   print_xpk_hello()

src/xpk/utils/console.py

@@ -16,6 +16,9 @@
 
 import sys
 from typing import NoReturn
+from typing import Literal
+
+from .execution_context import is_quiet
 
 
 def xpk_print(*args, **kwargs):
@@ -25,7 +28,7 @@ def xpk_print(*args, **kwargs):
     *args: user provided print args.
     **kwargs: user provided print args.
   """
-  sys.stdout.write('[XPK] ')
+  sys.stdout.write("[XPK] ")
   print(*args, **kwargs)
   sys.stdout.flush()
 
@@ -37,20 +40,36 @@ def xpk_exit(error_code) -> NoReturn:
     error_code: If the code provided is zero, then no issues occurred.
   """
   if error_code == 0:
-    xpk_print('Exiting XPK cleanly')
+    xpk_print("Exiting XPK cleanly")
     sys.exit(0)
   else:
-    xpk_print(f'XPK failed, error code {error_code}')
+    xpk_print(f"XPK failed, error code {error_code}")
     sys.exit(error_code)
 
 
-def get_user_input(input_msg):
-  """Function to get the user input for a prompt.
+def ask_for_user_consent(
+    question: str, default_option: Literal["Y", "N"] = "N"
+) -> bool:
+  """Prompts user with the given question, asking for a yes/no answer and returns a relevant boolean.
+  Important: immediatelly returns `True` in quiet mode!
+
+  Example prompt for `question='Continue?'`: `[XPK] Continue? (y/N): `.
 
   Args:
-    input_msg: message to be displayed by the prompt.
-  Returns:
-    True if user enter y or yes at the prompt, False otherwise.
+    question: The question to ask the user.
+    default_option: Option to use when user response is empty.
   """
-  user_input = input(input_msg)
-  return user_input in ('y', 'yes')
+  if is_quiet():
+    return True
+
+  options = "y/N" if default_option == "N" else "Y/n"
+  prompt = f"[XPK] {question} ({options}): "
+
+  while True:
+    user_input = input(prompt) or default_option
+    if user_input.lower() in ["yes", "y"]:
+      return True
+    elif user_input.lower() in ["no", "n"]:
+      return False
+    else:
+      xpk_print("Invalid input. Please enter: yes/no/y/n.")

src/xpk/utils/console_test.py

@@ -0,0 +1,92 @@
+"""
+Copyright 2025 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from unittest.mock import MagicMock, patch
+import pytest
+from pytest_mock import MockerFixture
+
+from xpk.utils.console import ask_for_user_consent
+
+
+@pytest.fixture(autouse=True)
+def mock_is_quiet(mocker: MockerFixture):
+  return mocker.patch("xpk.utils.console.is_quiet", return_value=False)
+
+
+@pytest.mark.parametrize(
+    "user_input,expected",
+    [
+        ("yes", True),
+        ("y", True),
+        ("Y", True),
+        ("Yes", True),
+        ("YES", True),
+        ("no", False),
+        ("n", False),
+        ("N", False),
+        ("No", False),
+        ("NO", False),
+    ],
+)
+@patch("xpk.utils.console.input")
+def test_ask_for_user_consent(mock_input: MagicMock, user_input, expected):
+  mock_input.return_value = user_input
+
+  assert ask_for_user_consent("Test question?") is expected
+
+
+def fake_input_factory(user_inputs: list[str]):
+  def fake_input(prompt: str) -> str:
+    return user_inputs.pop(0)
+
+  return fake_input
+
+
+@patch("xpk.utils.console.input", wraps=fake_input_factory(["invalid", "y"]))
+def test_ask_for_user_consent_invalid_input(mock_input: MagicMock):
+  agreed = ask_for_user_consent("Test question?")
+
+  assert agreed is True
+  assert mock_input.call_count == 2
+
+
+@patch("xpk.utils.console.input", return_value="")
+def test_ask_for_user_consent_default_No(mock_input: MagicMock):
+  agreed = ask_for_user_consent("Test question?", default_option="N")
+
+  assert agreed is False
+  mock_input.assert_called_once_with("[XPK] Test question? (y/N): ")
+
+
+@patch("xpk.utils.console.input", return_value="")
+def test_ask_for_user_consent_default_Yes(mock_input: MagicMock):
+  agreed = ask_for_user_consent("Test question?", default_option="Y")
+
+  assert agreed is True
+  mock_input.assert_called_once_with("[XPK] Test question? (Y/n): ")
+
+
+@patch("xpk.utils.console.input")
+def test_ask_for_user_consent_with_quiet_mode_always_agrees(
+    mock_input: MagicMock,
+    mock_is_quiet: MagicMock,
+):
+  mock_is_quiet.return_value = True
+
+  agreed = ask_for_user_consent("Test question?", default_option="N")
+
+  assert agreed is True
+  mock_input.assert_not_called()