xpk 0.14.3__py3-none-any.whl → 0.15.0__py3-none-any.whl

xpk/utils/console.py

@@ -16,6 +16,9 @@ limitations under the License.
 
 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,48 @@ 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.
+  """
+  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.")
+
+
+def exit_code_to_int(exit_code: str | int | None) -> int:
+  """
+  Converts sys._ExitCode to an int value that is used to exit the program.
+  See more: https://github.com/python/typeshed/issues/8513#issue-1333671093
   """
-  user_input = input(input_msg)
-  return user_input in ('y', 'yes')
+  if isinstance(exit_code, int):
+    return int(exit_code)
+  if exit_code is None:
+    return 0
+  return 1

xpk/utils/console_test.py

@@ -0,0 +1,106 @@
+"""
+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 .console import exit_code_to_int, 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()
+
+
+@pytest.mark.parametrize(
+    "exit_code,expected",
+    [
+        (0, 0),
+        (1, 1),
+        ("Error", 1),
+        ({"foo": "bar"}, 1),
+        (None, 0),
+    ],
+)
+def test_exit_code_to_int_returns_correct_value(exit_code, expected):
+  assert exit_code_to_int(exit_code) == expected

xpk/utils/feature_flags.py

@@ -18,11 +18,17 @@ import os
 
 
 def _get_boolean_flag(flag: str, default: bool) -> bool:
-  return os.getenv(flag, str(default)).lower() == "true"
+  experiment_value = os.getenv(flag, "").lower()
+  if experiment_value in ["true", "false"]:
+    return experiment_value == "true"
+
+  xpk_tester = os.getenv("XPK_TESTER", "").lower() == "true"
+  return xpk_tester or default
 
 
 class _FeatureFlags:
   SUB_SLICING_ENABLED = _get_boolean_flag("SUB_SLICING_ENABLED", default=False)
+  TELEMETRY_ENABLED = _get_boolean_flag("TELEMETRY_ENABLED", default=False)
 
 
 FeatureFlags = _FeatureFlags()

xpk/utils/file.py

@@ -18,6 +18,7 @@ import tempfile
 import os
 import hashlib
 from .execution_context import is_dry_run
+from .console import xpk_print
 
 
 def make_tmp_files(per_command_name: list[str]) -> list[str]:
@@ -51,7 +52,9 @@ def write_tmp_file(payload: str) -> str:
     A file object that was written to.
   """
   if is_dry_run():
-    return _hash_filename(payload)
+    name = _hash_filename(payload)
+    xpk_print(f'Temp file ({name}) content: \n{payload}')
+    return name
 
   with tempfile.NamedTemporaryFile(delete=False) as tmp:
     with open(file=tmp.name, mode='w', encoding='utf=8') as f:

xpk/utils/topology.py

@@ -44,3 +44,7 @@ def is_topology_contained(contained: str, container: str) -> bool:
       contained <= container
       for contained, container in zip(contained_parsed, container_parsed)
   )
+
+
+def get_slice_topology_level(topology: str) -> str:
+  return f"cloud.google.com/gke-tpu-slice-{topology}-id"

xpk/utils/user_agent.py

@@ -0,0 +1,35 @@
+"""
+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.
+"""
+
+import platform
+from ..core.config import __version__ as xpk_version
+
+
+def get_user_agent() -> str:
+  return f'XPK/{xpk_version} ({_get_user_agent_platform()})'
+
+
+def _get_user_agent_platform() -> str:
+  system = platform.system().lower()
+  if system == 'windows':
+    return f'Windows NT {platform.version()}'
+  elif system == 'linux':
+    return f'Linux; {platform.machine()}'
+  elif system == 'darwin':
+    version, _, arch = platform.mac_ver()
+    return f'Macintosh; {arch} Mac OS X {version}'
+  else:
+    return ''

xpk/utils/user_agent_test.py

@@ -0,0 +1,44 @@
+"""
+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 .user_agent import get_user_agent
+
+
+def test_get_user_agent_returns_correct_value_for_windows(mocker):
+  mocker.patch('xpk.utils.user_agent.xpk_version', 'v1.0.0')
+  mocker.patch('platform.system', return_value='Windows')
+  mocker.patch('platform.version', return_value='10.0')
+  assert get_user_agent() == 'XPK/v1.0.0 (Windows NT 10.0)'
+
+
+def test_get_user_agent_returns_correct_value_for_linux(mocker):
+  mocker.patch('xpk.utils.user_agent.xpk_version', 'v1.0.0')
+  mocker.patch('platform.system', return_value='Linux')
+  mocker.patch('platform.machine', return_value='x86_64')
+  assert get_user_agent() == 'XPK/v1.0.0 (Linux; x86_64)'
+
+
+def test_get_user_agent_returns_correct_value_for_darwin(mocker):
+  mocker.patch('xpk.utils.user_agent.xpk_version', 'v1.0.0')
+  mocker.patch('platform.system', return_value='Darwin')
+  mocker.patch('platform.mac_ver', return_value=('10.15', '', 'x86_64'))
+  assert get_user_agent() == 'XPK/v1.0.0 (Macintosh; x86_64 Mac OS X 10.15)'
+
+
+def test_get_user_agent_returns_correct_value_for_unknown(mocker):
+  mocker.patch('xpk.utils.user_agent.xpk_version', 'v1.0.0')
+  mocker.patch('platform.system', return_value='Unknown')
+  assert get_user_agent() == 'XPK/v1.0.0 ()'

xpk/utils/user_input.py

@@ -0,0 +1,48 @@
+"""
+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 typing import Literal
+
+from .console import xpk_print
+from .execution_context import is_quiet
+
+
+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:
+    question: The question to ask the user.
+    default_option: Option to use when user response is empty.
+  """
+  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.")

xpk/utils/user_input_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.user_input import ask_for_user_consent
+
+
+@pytest.fixture(autouse=True)
+def mock_is_quiet(mocker: MockerFixture):
+  return mocker.patch("xpk.utils.user_input.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.user_input.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.user_input.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.user_input.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.user_input.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.user_input.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()

xpk/utils/validation.py

@@ -15,10 +15,7 @@ limitations under the License.
 """
 
 from ..core.commands import run_command_for_value
-from ..core.config import __version__ as xpk_version
 from .console import xpk_exit, xpk_print
-from ..commands.config import xpk_cfg
-from ..core.config import DEPENDENCIES_KEY
 from enum import Enum
 from dataclasses import dataclass
 
@@ -80,14 +77,6 @@ def should_validate_dependencies(args):
   return not skip_validation and not dry_run
 
 
-def validate_dependencies():
-  """Validates all system dependencies if validation has not been done with current XPK version."""
-  deps_version = xpk_cfg.get(DEPENDENCIES_KEY)
-  if deps_version is None or deps_version != xpk_version:
-    validate_dependencies_list(list(SystemDependency))
-    xpk_cfg.set(DEPENDENCIES_KEY, xpk_version)
-
-
 def validate_dependencies_list(dependencies: list[SystemDependency]):
   """Validates a list of system dependencies and returns none or exits with error."""
   for dependency in dependencies:

xpk/utils/versions.py

@@ -0,0 +1,31 @@
+"""
+Copyright 2024 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.
+"""
+
+import enum
+
+
+class ReleaseChannel(enum.Enum):
+  """
+  Represents the GKE cluster release channels.
+
+  See: https://cloud.google.com/kubernetes-engine/docs/concepts/release-channels
+  """
+
+  RELEASE_CHANNEL_UNSPECIFIED = "RELEASE_CHANNEL_UNSPECIFIED"
+  RAPID = "RAPID"
+  REGULAR = "REGULAR"
+  STABLE = "STABLE"
+  EXTENDED = "EXTENDED"