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

xpk/core/nap.py

@@ -85,7 +85,7 @@ def enable_autoprovisioning_on_cluster(
   # TODO(@vbarr): Support timeout period for idle np before they are deleted.
   # TODO(@vbarr): Support for hot idle configuration (timeout period is infinity).
   return_code = 0
-  if system.accelerator_type == AcceleratorType['CPU']:
+  if system.accelerator_type == AcceleratorType.CPU:
     xpk_print("Error: XPK NAP doesn't support Accelerators of Types: CPUs.")
     return None, 1
 

xpk/core/nodepool.py

@@ -80,7 +80,7 @@ def run_gke_node_pool_create_command(
     if return_code > 0:
       xpk_print('Listing all reservations failed!')
     return_code = 1
-  if system.accelerator_type == AcceleratorType['TPU']:
+  if system.accelerator_type == AcceleratorType.TPU:
     max_nodes = system.vms_per_slice
   else:
     max_nodes = 1000
@@ -92,16 +92,14 @@ def run_gke_node_pool_create_command(
     return return_code
 
   desired_node_pool_count = (
-      1
-      if system.accelerator_type == AcceleratorType['GPU']
-      else args.num_slices
+      1 if system.accelerator_type == AcceleratorType.GPU else args.num_slices
   )
   message = (
       (
           f'Creating 1 node pool with {args.num_nodes} nodes of'
           f' {system.device_type}\nUnderlyingly, we assume that means: {system}'
       )
-      if system.accelerator_type == AcceleratorType['GPU']
+      if system.accelerator_type == AcceleratorType.GPU
       else (
           f'Creating {args.num_slices} node pool or pools of'
           f' {system.device_type}\nUnderlyingly, we assume that means: {system}'
@@ -267,7 +265,9 @@ def run_gke_node_pool_create_command(
 
   placement_args = ''
   if system.requires_workload_policy and is_topology_valid(system.topology):
-    placement_policy = f'{args.cluster}-placement-policy'
+    placement_policy = (
+        f'{system.device_type}-{system.topology}-placement-policy'
+    )
     ensure_resource_policy_exists(placement_policy, args, system.topology)
     placement_args = f' --placement-policy={placement_policy}'
 
@@ -288,7 +288,7 @@ def run_gke_node_pool_create_command(
         f'{placement_args}'
         ' --enable-gvnic'
     )
-    if system.accelerator_type == AcceleratorType['TPU']:
+    if system.accelerator_type == AcceleratorType.TPU:
       command += f' --node-version={gke_node_pool_version}'
       topology_product = get_topology_product(system.topology)
       if capacity_type == CapacityType.FLEX_START:
@@ -308,7 +308,7 @@ def run_gke_node_pool_create_command(
           command += f' --tpu-topology={system.topology}'
         command += ' --max-pods-per-node 15'
         command += f' {args.custom_tpu_nodepool_arguments}'
-    elif system.accelerator_type == AcceleratorType['GPU']:
+    elif system.accelerator_type == AcceleratorType.GPU:
       subnet_prefix = (
           f'{args.cluster}-{get_cluster_location(args.project, args.cluster, args.zone)}'
       )
@@ -328,7 +328,7 @@ def run_gke_node_pool_create_command(
               f' network={args.cluster}-net-{i},subnetwork={subnet_prefix}-sub-{i}'
           )
         command += ' --max-pods-per-node=32'
-    elif system.accelerator_type == AcceleratorType['CPU']:
+    elif system.accelerator_type == AcceleratorType.CPU:
       if capacity_type == CapacityType.FLEX_START:
         command += ' --num-nodes=0'
       else:

xpk/core/nodepool_test.py

@@ -174,7 +174,7 @@ def test_placement_policy_created_for_gpu_with_valid_topology(
       gke_accelerator="nvidia-h100-80gb",
       gce_machine_type="a3-highgpu-8g",
       chips_per_vm=8,
-      accelerator_type=AcceleratorType["GPU"],
+      accelerator_type=AcceleratorType.GPU,
       device_type="h100-80gb-8",
       supports_sub_slicing=False,
   )
@@ -204,7 +204,7 @@ def test_placement_policy_not_created_for_gpu_with_invalid_topology(
       gke_accelerator="nvidia-h100-80gb",
       gce_machine_type="a3-highgpu-8g",
       chips_per_vm=8,
-      accelerator_type=AcceleratorType["GPU"],
+      accelerator_type=AcceleratorType.GPU,
       device_type="h100-80gb-8",
       supports_sub_slicing=False,
   )
@@ -236,7 +236,7 @@ def test_placement_policy_created_for_tpu7x_with_valid_topology(
       gke_accelerator="tpu7x",
       gce_machine_type="tpu7x-standard-4t",
       chips_per_vm=4,
-      accelerator_type=AcceleratorType["TPU"],
+      accelerator_type=AcceleratorType.TPU,
       device_type="tpu7x-8",
       requires_workload_policy=True,
       supports_sub_slicing=False,
@@ -269,7 +269,7 @@ def test_placement_policy_not_created_for_non7x_tpu(
       gke_accelerator="v6e",
       gce_machine_type="tpu-v6e-slice",
       chips_per_vm=4,
-      accelerator_type=AcceleratorType["TPU"],
+      accelerator_type=AcceleratorType.TPU,
       device_type="v6e-4",
       supports_sub_slicing=True,
   )

xpk/core/pathways.py

@@ -91,7 +91,7 @@ def ensure_pathways_workload_prerequisites(args, system) -> bool:
     xpk_exit(1)
 
   # Ensure device type is TPUs - currently Pathways supports TPUs only.
-  if system.accelerator_type != AcceleratorType['TPU']:
+  if system.accelerator_type != AcceleratorType.TPU:
     xpk_print('Currently, Pathways workloads can only be run on TPUs.')
     xpk_exit(1)
 

xpk/core/resources.py

@@ -109,7 +109,7 @@ def create_cluster_configmaps(
 
   # ConfigMap to store resources available in the cluster.
   device_type = system.device_type
-  if system.accelerator_type == AcceleratorType['GPU']:
+  if system.accelerator_type == AcceleratorType.GPU:
     resources_data = f'{device_type}: "{int(args.num_nodes)}"'
   elif args.enable_autoprovisioning and autoprovisioning_config:
     resources_data = (

xpk/core/scheduling.py

@@ -97,7 +97,7 @@ def check_if_workload_can_schedule(args, system: SystemCharacteristics) -> bool:
   else:
     # Check if the size of the workload will fit in the cluster.
     max_vm_in_cluster = int(cluster_config_map[device_type])
-    if system.accelerator_type == AcceleratorType['GPU']:
+    if system.accelerator_type == AcceleratorType.GPU:
       vm_required_by_workload = args.num_nodes
     else:
       vm_required_by_workload = args.num_slices * system.vms_per_slice
@@ -125,7 +125,7 @@ def get_total_chips_requested_from_args(
   Returns:
     num of chips for the current request.
   """
-  if system.accelerator_type == AcceleratorType['GPU']:
+  if system.accelerator_type == AcceleratorType.GPU:
     num_chips = system.vms_per_slice * system.chips_per_vm * args.num_nodes
   else:
     num_chips = system.vms_per_slice * system.chips_per_vm * args.num_slices
@@ -152,7 +152,7 @@ def get_cpu_affinity(accelerator_type) -> str:
                         values:
                         - default-pool
 """
-  if accelerator_type == AcceleratorType['CPU']:
+  if accelerator_type == AcceleratorType.CPU:
     return yaml
   return ''
 
@@ -225,7 +225,7 @@ def create_accelerator_label(accelerator_type, system) -> str:
   Returns:
     The accelerator label.
   """
-  if accelerator_type == AcceleratorType['CPU']:
+  if accelerator_type == AcceleratorType.CPU:
     return ''
   return (
       f'{AcceleratorTypeToAcceleratorCharacteristics[accelerator_type].accelerator_label}:'
@@ -243,7 +243,7 @@ def create_tpu_machine_type(accelerator_type, system) -> str:
   Returns:
     The accelerator label.
   """
-  if accelerator_type == AcceleratorType['TPU']:
+  if accelerator_type == AcceleratorType.TPU:
     return f'{system.gce_machine_type}'
   return ''
 
@@ -261,10 +261,7 @@ def create_machine_label(
   Returns:
     The machine label.
   """
-  if (
-      accelerator_type == AcceleratorType['TPU']
-      and not autoprovisioning_enabled
-  ):
+  if accelerator_type == AcceleratorType.TPU and not autoprovisioning_enabled:
     return (
         f'{AcceleratorTypeToAcceleratorCharacteristics[accelerator_type].machine_label}:'
         f' {system.topology}'
@@ -285,10 +282,7 @@ def create_tpu_topology(
   Returns:
     The machine label.
   """
-  if (
-      accelerator_type == AcceleratorType['TPU']
-      and not autoprovisioning_enabled
-  ):
+  if accelerator_type == AcceleratorType.TPU and not autoprovisioning_enabled:
     return f'{system.topology}'
   return ''
 

xpk/core/system_characteristics.py

@@ -16,9 +16,16 @@ limitations under the License.
 
 from dataclasses import dataclass
 from ..utils.topology import get_topology_product
+from enum import Enum
 
 
-AcceleratorType = {'TPU': 1, 'GPU': 2, 'CPU': 3}
+class AcceleratorType(Enum):
+  TPU = 1
+  GPU = 2
+  CPU = 3
+
+  def __repr__(self):
+    return self._name_
 
 
 @dataclass
@@ -29,17 +36,17 @@ class AcceleratorCharacteristics:
 
 
 AcceleratorTypeToAcceleratorCharacteristics = {
-    AcceleratorType['TPU']: AcceleratorCharacteristics(
+    AcceleratorType.TPU: AcceleratorCharacteristics(
         resource_type='google.com/tpu',
         accelerator_label='cloud.google.com/gke-tpu-accelerator',
         machine_label='cloud.google.com/gke-tpu-topology',
     ),
-    AcceleratorType['GPU']: AcceleratorCharacteristics(
+    AcceleratorType.GPU: AcceleratorCharacteristics(
         resource_type='nvidia.com/gpu',
         accelerator_label='cloud.google.com/gke-accelerator',
         machine_label='cloud.google.com/gce-machine-type',
     ),
-    AcceleratorType['CPU']: AcceleratorCharacteristics(
+    AcceleratorType.CPU: AcceleratorCharacteristics(
         resource_type='cpu',
         accelerator_label='',
         machine_label='cloud.google.com/gke-nodepool',
@@ -80,13 +87,13 @@ class SystemCharacteristics:
   gke_accelerator: str
   gce_machine_type: str
   chips_per_vm: int
-  accelerator_type: int  # TODO: use enums
+  accelerator_type: AcceleratorType
   device_type: str
   supports_sub_slicing: bool
   requires_workload_policy: bool = False
 
   def __post_init__(self):
-    if self.accelerator_type == AcceleratorType['GPU']:
+    if self.accelerator_type == AcceleratorType.GPU:
       self.requires_workload_policy = True
 
 
@@ -144,7 +151,7 @@ def get_tpu_system_characteristics_map(
         gke_accelerator=gke_accelerator,
         gce_machine_type=machine_type,
         chips_per_vm=chips_per_vm,
-        accelerator_type=AcceleratorType['TPU'],
+        accelerator_type=AcceleratorType.TPU,
         device_type=f'{prefix}-{num_tensorcores}',
         requires_workload_policy=requires_workload_policy,
         supports_sub_slicing=supports_sub_slicing,
@@ -183,7 +190,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='nvidia-l4',
         gce_machine_type='g2-standard-12',
         chips_per_vm=1,
-        accelerator_type=AcceleratorType['GPU'],
+        accelerator_type=AcceleratorType.GPU,
         device_type='l4-1',
         supports_sub_slicing=False,
     ),
@@ -193,7 +200,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='nvidia-l4',
         gce_machine_type='g2-standard-24',
         chips_per_vm=2,
-        accelerator_type=AcceleratorType['GPU'],
+        accelerator_type=AcceleratorType.GPU,
         device_type='l4-2',
         supports_sub_slicing=False,
     ),
@@ -203,7 +210,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='nvidia-l4',
         gce_machine_type='g2-standard-48',
         chips_per_vm=4,
-        accelerator_type=AcceleratorType['GPU'],
+        accelerator_type=AcceleratorType.GPU,
         device_type='l4-4',
         supports_sub_slicing=False,
     ),
@@ -213,7 +220,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='nvidia-l4',
         gce_machine_type='g2-standard-96',
         chips_per_vm=8,
-        accelerator_type=AcceleratorType['GPU'],
+        accelerator_type=AcceleratorType.GPU,
         device_type='l4-8',
         supports_sub_slicing=False,
     ),
@@ -224,7 +231,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='nvidia-tesla-a100',
         gce_machine_type='a2-highgpu-1g',
         chips_per_vm=1,
-        accelerator_type=AcceleratorType['GPU'],
+        accelerator_type=AcceleratorType.GPU,
         device_type='a100-40gb-1',
         supports_sub_slicing=False,
     ),
@@ -234,7 +241,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='nvidia-tesla-a100',
         gce_machine_type='a2-highgpu-2g',
         chips_per_vm=2,
-        accelerator_type=AcceleratorType['GPU'],
+        accelerator_type=AcceleratorType.GPU,
         device_type='a100-40gb-2',
         supports_sub_slicing=False,
     ),
@@ -244,7 +251,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='nvidia-tesla-a100',
         gce_machine_type='a2-highgpu-4g',
         chips_per_vm=4,
-        accelerator_type=AcceleratorType['GPU'],
+        accelerator_type=AcceleratorType.GPU,
         device_type='a100-40gb-4',
         supports_sub_slicing=False,
     ),
@@ -254,7 +261,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='nvidia-tesla-a100',
         gce_machine_type='a2-highgpu-8g',
         chips_per_vm=8,
-        accelerator_type=AcceleratorType['GPU'],
+        accelerator_type=AcceleratorType.GPU,
         device_type='a100-40gb-8',
         supports_sub_slicing=False,
     ),
@@ -264,7 +271,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='nvidia-gb200',
         gce_machine_type='a4x-highgpu-4g',
         chips_per_vm=4,
-        accelerator_type=AcceleratorType['GPU'],
+        accelerator_type=AcceleratorType.GPU,
         device_type='gb200-4',
         supports_sub_slicing=False,
     ),
@@ -274,7 +281,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='nvidia-gb200',
         gce_machine_type='a4x-highgpu-4g-nolssd',
         chips_per_vm=4,
-        accelerator_type=AcceleratorType['GPU'],
+        accelerator_type=AcceleratorType.GPU,
         device_type='gb200-4',
         supports_sub_slicing=False,
     ),
@@ -284,7 +291,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='nvidia-b200',
         gce_machine_type='a4-highgpu-8g',
         chips_per_vm=8,
-        accelerator_type=AcceleratorType['GPU'],
+        accelerator_type=AcceleratorType.GPU,
         device_type='b200-8',
         supports_sub_slicing=False,
     ),
@@ -294,7 +301,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='nvidia-h200-141gb',
         gce_machine_type='a3-ultragpu-8g',
         chips_per_vm=8,
-        accelerator_type=AcceleratorType['GPU'],
+        accelerator_type=AcceleratorType.GPU,
         device_type='h200-141gb-8',
         supports_sub_slicing=False,
     ),
@@ -305,7 +312,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='nvidia-h100-80gb',
         gce_machine_type='a3-highgpu-8g',
         chips_per_vm=8,
-        accelerator_type=AcceleratorType['GPU'],
+        accelerator_type=AcceleratorType.GPU,
         device_type='h100-80gb-8',
         supports_sub_slicing=False,
     ),
@@ -316,7 +323,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='nvidia-h100-mega-80gb',
         gce_machine_type='a3-megagpu-8g',
         chips_per_vm=8,
-        accelerator_type=AcceleratorType['GPU'],
+        accelerator_type=AcceleratorType.GPU,
         device_type='h100-mega-80gb-8',
         supports_sub_slicing=False,
     ),
@@ -605,7 +612,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='N/A',
         gce_machine_type='m1-megamem-96',
         chips_per_vm=96,
-        accelerator_type=AcceleratorType['CPU'],
+        accelerator_type=AcceleratorType.CPU,
         device_type='m1-megamem-96-1',
         supports_sub_slicing=False,
     ),
@@ -616,7 +623,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='N/A',
         gce_machine_type='n2-standard-64',
         chips_per_vm=64,
-        accelerator_type=AcceleratorType['CPU'],
+        accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-64-1',
         supports_sub_slicing=False,
     ),
@@ -626,7 +633,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='N/A',
         gce_machine_type='n2-standard-32',
         chips_per_vm=32,
-        accelerator_type=AcceleratorType['CPU'],
+        accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-1',
         supports_sub_slicing=False,
     ),
@@ -636,7 +643,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='N/A',
         gce_machine_type='n2-standard-32',
         chips_per_vm=32,
-        accelerator_type=AcceleratorType['CPU'],
+        accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-2',
         supports_sub_slicing=False,
     ),
@@ -646,7 +653,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='N/A',
         gce_machine_type='n2-standard-32',
         chips_per_vm=32,
-        accelerator_type=AcceleratorType['CPU'],
+        accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-4',
         supports_sub_slicing=False,
     ),
@@ -656,7 +663,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='N/A',
         gce_machine_type='n2-standard-32',
         chips_per_vm=32,
-        accelerator_type=AcceleratorType['CPU'],
+        accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-8',
         supports_sub_slicing=False,
     ),
@@ -666,7 +673,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='N/A',
         gce_machine_type='n2-standard-32',
         chips_per_vm=32,
-        accelerator_type=AcceleratorType['CPU'],
+        accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-16',
         supports_sub_slicing=False,
     ),
@@ -676,7 +683,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='N/A',
         gce_machine_type='n2-standard-32',
         chips_per_vm=32,
-        accelerator_type=AcceleratorType['CPU'],
+        accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-32',
         supports_sub_slicing=False,
     ),
@@ -686,7 +693,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='N/A',
         gce_machine_type='n2-standard-32',
         chips_per_vm=32,
-        accelerator_type=AcceleratorType['CPU'],
+        accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-64',
         supports_sub_slicing=False,
     ),
@@ -696,7 +703,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='N/A',
         gce_machine_type='n2-standard-32',
         chips_per_vm=32,
-        accelerator_type=AcceleratorType['CPU'],
+        accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-128',
         supports_sub_slicing=False,
     ),
@@ -706,7 +713,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='N/A',
         gce_machine_type='n2-standard-32',
         chips_per_vm=32,
-        accelerator_type=AcceleratorType['CPU'],
+        accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-256',
         supports_sub_slicing=False,
     ),
@@ -716,7 +723,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='N/A',
         gce_machine_type='n2-standard-32',
         chips_per_vm=32,
-        accelerator_type=AcceleratorType['CPU'],
+        accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-512',
         supports_sub_slicing=False,
     ),
@@ -726,7 +733,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='N/A',
         gce_machine_type='n2-standard-32',
         chips_per_vm=32,
-        accelerator_type=AcceleratorType['CPU'],
+        accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-1024',
         supports_sub_slicing=False,
     ),
@@ -736,7 +743,7 @@ UserFacingNameToSystemCharacteristics = {
         gke_accelerator='N/A',
         gce_machine_type='n2-standard-32',
         chips_per_vm=32,
-        accelerator_type=AcceleratorType['CPU'],
+        accelerator_type=AcceleratorType.CPU,
         device_type='n2-standard-32-2048',
         supports_sub_slicing=False,
     ),

xpk/core/system_characteristics_test.py

@@ -14,7 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-from .system_characteristics import get_tpu_system_characteristics_map, SystemCharacteristics
+from .system_characteristics import get_tpu_system_characteristics_map, SystemCharacteristics, AcceleratorType
 
 
 def test_get_tpu_system_characteristics_map_returns_correct_values_for_1x1_topology():
@@ -34,7 +34,7 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_1x1_topol
       gke_accelerator="test",
       gce_machine_type="test",
       chips_per_vm=1,
-      accelerator_type=1,
+      accelerator_type=AcceleratorType.TPU,
       device_type="test-1",
       supports_sub_slicing=False,
       requires_workload_policy=True,
@@ -62,7 +62,7 @@ def test_get_tpu_system_characteristics_map_returns_correct_values_for_2x2_topol
       gke_accelerator="test",
       gce_machine_type="test",
       chips_per_vm=4,
-      accelerator_type=1,
+      accelerator_type=AcceleratorType.TPU,
       device_type="test-8",
       supports_sub_slicing=False,
       requires_workload_policy=True,

xpk/core/testing/__init__.py

@@ -0,0 +1,15 @@
+"""
+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.
+"""

xpk/core/testing/commands_tester.py

@@ -0,0 +1,131 @@
+"""
+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 re
+from pytest_mock import MockerFixture
+
+
+class CommandsTester:
+  """Tester class useful for mocking and asserting command runs."""
+
+  def __init__(
+      self,
+      mocker: MockerFixture,
+      run_command_for_value_path: str | None = None,
+      run_command_with_updates_path: str | None = None,
+      run_command_with_updates_retry_path: str | None = None,
+  ):
+    self.__results: dict[re.Pattern, tuple[int, str]] = {}
+    self.commands_history: list[str] = []
+    if run_command_for_value_path:
+      mocker.patch(
+          run_command_for_value_path,
+          wraps=self.__fake_run_command_for_value,
+      )
+    if run_command_with_updates_path:
+      mocker.patch(
+          run_command_with_updates_path,
+          wraps=self.__fake_run_command_with_updates,
+      )
+    if run_command_with_updates_retry_path:
+      mocker.patch(
+          run_command_with_updates_retry_path,
+          wraps=self.__fake_run_command_with_updates_retry,
+      )
+
+  def set_result_for_command(
+      self, result: tuple[int, str], *command_parts: str
+  ):
+    """Sets the result for the given command parts.
+    The command parts will be joined with '.*' during comparison with the actual commands.
+    """
+    pattern = self.__get_pattern_for_command_parts(*command_parts)
+    self.__results[pattern] = result
+
+  def assert_command_run(self, *command_parts: str, times: int = 1):
+    """Asserts the command composed from the command parts (joined with '.*') was run exactly `times` times."""
+    matching = self.get_matching_commands(*command_parts)
+    if not matching:
+      raise AssertionError(
+          f"{command_parts} was not found in {self.commands_history}"
+      )
+    elif len(matching) != times:
+      raise AssertionError(
+          f"{command_parts} was expected to be run {times} times in"
+          f" {self.commands_history}"
+      )
+
+  def assert_command_not_run(self, *command_parts: str):
+    """Asserts the command composed from the command parts (joined with '.*') was never run."""
+    if self.get_matching_commands(*command_parts):
+      raise AssertionError(
+          f"{command_parts} was found in {self.commands_history}"
+      )
+
+  def get_matching_commands(self, *command_parts: str) -> list[str]:
+    """Returns list of already run commands matching the command parts (joined with '.*')."""
+    pattern = self.__get_pattern_for_command_parts(*command_parts)
+    return [c for c in self.commands_history if pattern.match(c)]
+
+  # Unused arguments, but the signature has to match the original one:
+  # pylint: disable=unused-argument
+  def __fake_run_command_with_updates(
+      self,
+      command: str,
+      task: str,
+      verbose=True,
+  ) -> int:
+    return self.__common_fake_run_command(command, (0, ""))[0]
+
+  def __fake_run_command_with_updates_retry(
+      self,
+      command: str,
+      task: str,
+      verbose=True,
+      num_retry_attempts=5,
+      wait_seconds=10,
+  ) -> int:
+    return self.__common_fake_run_command(command, (0, ""))[0]
+
+  def __fake_run_command_for_value(
+      self,
+      command: str,
+      task: str,
+      dry_run_return_val="0",
+      print_timer=False,
+      hide_error=False,
+      quiet=False,
+  ) -> tuple[int, str]:
+    return self.__common_fake_run_command(command, (0, dry_run_return_val))
+
+  # pylint: enable=unused-argument
+
+  def __common_fake_run_command(
+      self,
+      command: str,
+      default_result: tuple[int, str],
+  ) -> tuple[int, str]:
+    self.commands_history.append(command)
+    matching_results = [
+        result
+        for pattern, result in self.__results.items()
+        if pattern.match(command)
+    ]
+    return len(matching_results) > 0 and matching_results[0] or default_result
+
+  def __get_pattern_for_command_parts(self, *command_parts: str) -> re.Pattern:
+    pattern_s = ".*" + ".*".join(map(re.escape, command_parts)) + ".*"
+    return re.compile(pattern_s)