change: Create a general function and update tests

Author: fege

tests/model_registry/negative_tests/test_db_migration.py

@@ -7,8 +7,8 @@
 from utilities.constants import DscComponents
 from tests.model_registry.constants import MR_INSTANCE_NAME
 from kubernetes.dynamic.client import DynamicClient
-from utilities.general import wait_for_container_status
-from tests.model_registry.utils import wait_for_new_running_mr_pods
+from utilities.general import wait_for_container_status, wait_for_new_running_pods
+from utilities.constants import Timeout
 
 LOGGER = get_logger(name=__name__)
 
@@ -48,11 +48,17 @@ def test_db_migration_negative(
         4. Check the logs for the expected error
         """
         LOGGER.info(f"Model registry pod: {model_registry_pod.name}")
-        mr_pods = wait_for_new_running_mr_pods(
-            admin_client=admin_client,
+        mr_pods = wait_for_new_running_pods(
             orig_pods=[model_registry_pod],
-            namespace=py_config["model_registry_namespace"],
-            instance_name=MR_INSTANCE_NAME,
+            pod_fetcher_func=lambda: list(
+                Pod.get(
+                    dyn_client=admin_client,
+                    namespace=py_config["model_registry_namespace"],
+                    label_selector=f"app={MR_INSTANCE_NAME}",
+                )
+            ),
+            timeout=Timeout.TIMEOUT_2MIN,
+            return_pods=True,
         )
         mr_pod = mr_pods[0]
         LOGGER.info(f"Pod that should contains the container in CrashLoopBackOff state: {mr_pod.name}")

tests/model_registry/utils.py

@@ -235,54 +235,6 @@ def wait_for_pods_running(
     return None
 
 
-def wait_for_new_running_mr_pods(
-    admin_client: DynamicClient,
-    orig_pods: list[Pod],
-    namespace: str,
-    instance_name: str,
-    expected_num_pods: int | None = None,
-) -> list[Pod]:
-    """
-    Wait for the model registry pod to be replaced.
-
-    Args:
-        admin_client (DynamicClient): The admin client.
-        orig_pods (list): List of Pod objects.
-        expected_num_pods (int): Number of pods expected to be running.
-        If not provided, the number of pods is expected to be len(orig_pods)
-    Returns:
-        List of Pod objects.
-
-    Raises:
-        TimeoutError: If the pods are not replaced.
-
-    """
-    LOGGER.info("Waiting for pods to be replaced")
-    oring_pods_names = [pod.name for pod in orig_pods]
-
-    expected_num_pods = expected_num_pods or len(orig_pods)
-
-    try:
-        for pods in TimeoutSampler(
-            wait_timeout=180,
-            sleep=5,
-            func=lambda: list(
-                Pod.get(
-                    dyn_client=admin_client,
-                    namespace=namespace,
-                    label_selector=f"app={instance_name}",
-                )
-            ),
-        ):
-            if pods and len(pods) == expected_num_pods:
-                if all(pod.name not in oring_pods_names and pod.status == pod.Status.RUNNING for pod in pods):
-                    return pods
-
-    except TimeoutError:
-        LOGGER.error(f"Timeout waiting for pods {oring_pods_names} to be replaced")
-        raise
-
-
 def generate_namespace_name(file_path: str) -> str:
     return (file_path.removesuffix(".py").replace("/", "-").replace("_", "-"))[-63:].split("-", 1)[-1]
 

tests/model_serving/model_server/inference_service_configuration/utils.py

@@ -3,13 +3,12 @@
 
 from kubernetes.dynamic import DynamicClient
 from ocp_resources.inference_service import InferenceService
-from ocp_resources.pod import Pod
 from ocp_resources.resource import ResourceEditor
 from simple_logger.logger import get_logger
-from timeout_sampler import TimeoutSampler
 
 from utilities.constants import Timeout
 from utilities.infra import get_pods_by_isvc_label
+from utilities.general import wait_for_new_running_pods
 
 LOGGER = get_logger(name=__name__)
 
@@ -32,7 +31,11 @@ def update_inference_service(
 
     with ResourceEditor(patches={isvc: isvc_updated_dict}):
         if wait_for_new_pods:
-            wait_for_new_running_inference_pods(isvc=isvc, orig_pods=orig_pods)
+            wait_for_new_running_pods(
+                orig_pods=orig_pods,
+                pod_fetcher_func=lambda: get_pods_by_isvc_label(client=isvc.client, isvc=isvc),
+                timeout=Timeout.TIMEOUT_10MIN,
+            )
 
         yield isvc
 
@@ -72,44 +75,6 @@ def verify_env_vars_in_isvc_pods(isvc: InferenceService, env_vars: list[dict[str
         )
 
 
-def wait_for_new_running_inference_pods(
-    isvc: InferenceService, orig_pods: list[Pod], expected_num_pods: int | None = None
-) -> None:
-    """
-    Wait for the inference pod to be replaced.
-
-    Args:
-        isvc (InferenceService): InferenceService object.
-        orig_pods (list): List of Pod objects.
-        expected_num_pods (int): Number of pods expected to be running. I
-            f not provided, the number of pods is expected to be len(orig_pods)
-
-    Raises:
-        TimeoutError: If the pods are not replaced.
-
-    """
-    LOGGER.info("Waiting for pods to be replaced")
-    oring_pods_names = [pod.name for pod in orig_pods]
-
-    expected_num_pods = expected_num_pods or len(orig_pods)
-
-    try:
-        for pods in TimeoutSampler(
-            wait_timeout=Timeout.TIMEOUT_10MIN,
-            sleep=5,
-            func=get_pods_by_isvc_label,
-            client=isvc.client,
-            isvc=isvc,
-        ):
-            if pods and len(pods) == expected_num_pods:
-                if all(pod.name not in oring_pods_names and pod.status == pod.Status.RUNNING for pod in pods):
-                    return
-
-    except TimeoutError:
-        LOGGER.error(f"Timeout waiting for pods {oring_pods_names} to be replaced")
-        raise
-
-
 def verify_pull_secret(isvc: InferenceService, pull_secret: str, secret_exists: bool) -> None:
     """
     Verify that the ImagePullSecret in the InferenceService pods match the expected values.

utilities/general.py

@@ -1,6 +1,6 @@
 import base64
 import re
-from typing import List, Tuple
+from typing import List, Tuple, Callable
 import uuid
 
 from kubernetes.dynamic import DynamicClient
@@ -14,7 +14,7 @@
 from utilities.constants import Annotations, KServeDeploymentType, MODELMESH_SERVING, Timeout
 from utilities.exceptions import UnexpectedResourceCountError, ResourceValueMismatch
 from ocp_resources.resource import Resource
-from timeout_sampler import retry
+from timeout_sampler import retry, TimeoutSampler
 
 # Constants for image validation
 SHA256_DIGEST_PATTERN = r"@sha256:[a-f0-9]{64}$"
@@ -378,3 +378,42 @@ def wait_for_container_status(pod: Pod, container_name: str, expected_status: st
         LOGGER.info(f"Container {container_name} is in the expected status {expected_status}")
         return True
     raise ResourceValueMismatch(f"Container {container_name} is not in the expected status {container_status.state}")
+
+
+def wait_for_new_running_pods(  # type: ignore[return]
+    orig_pods: list[Pod],
+    pod_fetcher_func: Callable[[], list[Pod]],
+    expected_num_pods: int | None = None,
+    timeout: int = Timeout.TIMEOUT_10MIN,
+    return_pods: bool = False,
+) -> list[Pod] | None:
+    """
+    Wait for pods to be replaced with new running pods.
+
+    Args:
+        orig_pods: List of original Pod objects to be replaced
+        pod_fetcher_func: Function that returns list of current pods
+        expected_num_pods: Expected number of pods (defaults to len(orig_pods))
+        timeout: Timeout in seconds
+        return_pods: Whether to return the new pods list
+
+    Returns:
+        List of new pods if return_pods=True, otherwise None
+    """
+    LOGGER.info("Waiting for pods to be replaced")
+    orig_pods_names = [pod.name for pod in orig_pods]
+    expected_num_pods = expected_num_pods or len(orig_pods)
+
+    try:
+        for pods in TimeoutSampler(
+            wait_timeout=timeout,
+            sleep=5,
+            func=pod_fetcher_func,
+        ):
+            if pods and len(pods) == expected_num_pods:
+                if all(pod.name not in orig_pods_names and pod.status == pod.Status.RUNNING for pod in pods):
+                    return pods if return_pods else None
+
+    except TimeoutError:
+        LOGGER.error(f"Timeout waiting for pods {orig_pods_names} to be replaced")
+        raise