Wait for newer pod to be available after deployment deletion

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_pods_by_labels, wait_for_container_status
-
+from utilities.general import wait_for_container_status
+from tests.model_registry.utils import wait_for_new_running_mr_pods
 
 LOGGER = get_logger(name=__name__)
 
@@ -36,30 +36,32 @@ def test_db_migration_negative(
         admin_client: DynamicClient,
         model_registry_db_instance_pod: Pod,
         set_mr_db_dirty: int,
+        model_registry_pod: Pod,
         delete_mr_deployment: None,
     ):
         """
         RHOAIENG-27505: This test is to check the migration error when the database is dirty.
         The test will:
         1. Set the dirty flag to 1 for the latest migration version
         2. Delete the model registry deployment
-        3. Check the logs for the expected error
+        3. Wait for the old pods to be terminated
+        4. Check the logs for the expected error
         """
-        mr_pods = wait_for_pods_by_labels(
+        LOGGER.info(f"Model registry pod: {model_registry_pod.name}")
+        mr_pods = wait_for_new_running_mr_pods(
             admin_client=admin_client,
+            orig_pods=[model_registry_pod],
             namespace=py_config["model_registry_namespace"],
-            label_selector=f"app={MR_INSTANCE_NAME}",
-            expected_num_pods=1,
+            instance_name=MR_INSTANCE_NAME,
         )
         mr_pod = mr_pods[0]
-        LOGGER.info("Waiting for model registry pod to crash")
+        LOGGER.info(f"Pod that should contains the container in CrashLoopBackOff state: {mr_pod.name}")
         assert wait_for_container_status(mr_pod, "rest-container", Pod.Status.CRASH_LOOPBACK_OFF)
 
         LOGGER.info("Checking the logs for the expected error")
-
         log_output = mr_pod.log(container="rest-container")
         expected_error = (
             f"Error: {{{{ALERT}}}} error connecting to datastore: Dirty database version {set_mr_db_dirty}. "
             "Fix and force version."
         )
-        assert expected_error in log_output, "Expected error message not found in logs!"
+        assert expected_error in log_output, f"Expected error message not found in logs!\n{log_output}"

tests/model_registry/utils.py

@@ -235,6 +235,54 @@ 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")
+    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=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 orig_pods_names and pod.status == pod.Status.RUNNING for pod in pods):
+                    return pods
+
+    except TimeoutError:
+        LOGGER.error(f"Timeout waiting for pods {orig_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]