Refactor ensure_resource_health to support statefulset

- Note that LocalExecutor use statefulset for scheduler instead of
  deployment

Author: jason810496

kubernetes_tests/test_base.py

@@ -23,6 +23,7 @@
 from datetime import datetime, timezone
 from pathlib import Path
 from subprocess import check_call, check_output
+from typing import Literal
 
 import pytest
 import re2
@@ -62,7 +63,7 @@ def base_tests_setup(self, request):
         self.session = self._get_session_with_retries()
 
         # Ensure the api-server deployment is healthy at kubernetes level before calling the any API
-        self.ensure_deployment_health("airflow-webserver")
+        self.ensure_resource_health("airflow-webserver")
         try:
             self._ensure_airflow_webserver_is_healthy()
             yield
@@ -195,12 +196,25 @@ def monitor_task(self, host, dag_run_id, dag_id, task_id, expected_final_state,
         assert state == expected_final_state
 
     @staticmethod
-    def ensure_deployment_health(deployment_name: str, namespace: str = "airflow"):
-        """Watch the deployment until it is healthy."""
-        deployment_rollout_status = check_output(
-            ["kubectl", "rollout", "status", "deployment", deployment_name, "-n", namespace, "--watch"]
+    def ensure_resource_health(
+        resource_name: str,
+        namespace: str = "airflow",
+        resource_type: Literal["deployment", "statefulset"] = "deployment",
+    ):
+        """Watch the resource until it is healthy.
+
+        Args:
+            resource_name (str): Name of the resource to check.
+            resource_type (str): Type of the resource (e.g., deployment, statefulset).
+            namespace (str): Kubernetes namespace where the resource is located.
+        """
+        rollout_status = check_output(
+            ["kubectl", "rollout", "status", f"{resource_type}/{resource_name}", "-n", namespace, "--watch"],
         ).decode()
-        assert "successfully rolled out" in deployment_rollout_status
+        if resource_type == "deployment":
+            assert "successfully rolled out" in rollout_status
+        else:
+            assert "roll out complete" in rollout_status
 
     def ensure_dag_expected_state(self, host, execution_date, dag_id, expected_final_state, timeout):
         tries = 0

kubernetes_tests/test_kubernetes_executor.py

@@ -55,7 +55,7 @@ def test_integration_run_dag_with_scheduler_failure(self):
 
         self._delete_airflow_pod("scheduler")
 
-        self.ensure_deployment_health("airflow-scheduler")
+        self.ensure_resource_health("airflow-scheduler")
 
         # Wait some time for the operator to complete
         self.monitor_task(

kubernetes_tests/test_other_executors.py

@@ -56,7 +56,13 @@ def test_integration_run_dag_with_scheduler_failure(self):
 
         self._delete_airflow_pod("scheduler")
 
-        self.ensure_deployment_health("airflow-scheduler")
+        if EXECUTOR == "CeleryExecutor":
+            scheduler_resource_type = "deployment"
+        elif EXECUTOR == "LocalExecutor":
+            scheduler_resource_type = "statefulset"
+        else:
+            raise ValueError(f"Unknown executor {EXECUTOR}")
+        self.ensure_resource_health("airflow-scheduler", resource_type=scheduler_resource_type)
 
         # Wait some time for the operator to complete
         self.monitor_task(