Address review feedback: harden worker-health-check command

- Replace os.EX_UNAVAILABLE with SystemExit(message) for portability
- Add timeout=5.0 to inspect() calls for defense in depth
- Add try/except around inspect calls for clean broker error handling
- Add tests for broker unreachable scenarios (ping and active_queues)
- Move import socket to module-level

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: antonio-mello-ai

providers/celery/src/airflow/providers/celery/cli/celery_command.py

@@ -20,6 +20,7 @@
 from __future__ import annotations
 
 import logging
+import socket
 import sys
 import time
 from contextlib import contextmanager, suppress
@@ -487,42 +488,53 @@ def worker_health_check(args):
 
     See: https://github.com/apache/airflow/issues/63580
     """
-    import os
-    import socket
-
     # This needs to be imported locally to not trigger Providers Manager initialization
     from airflow.providers.celery.executors.celery_executor import app as celery_app
 
     hostname = args.celery_hostname or f"celery@{socket.gethostname()}"
 
-    inspect = celery_app.control.inspect(destination=[hostname])
+    inspect = celery_app.control.inspect(destination=[hostname], timeout=5.0)
+
+    try:
+        # Step 1: Ping check — worker must be reachable
+        log.info("Checking worker liveness via ping: %s", hostname)
+        ping_result = inspect.ping()
+    except Exception:
+        log.exception("Worker health check failed: unable to reach broker while pinging worker %s", hostname)
+        raise SystemExit("Worker health check failed: broker connection error during ping")
 
-    # Step 1: Ping check — worker must be reachable
-    log.info("Checking worker liveness via ping: %s", hostname)
-    ping_result = inspect.ping()
     if not ping_result or hostname not in ping_result:
         log.error(
             "Worker health check failed: worker %s did not respond to ping. The worker process may be down.",
             hostname,
         )
-        raise SystemExit(os.EX_UNAVAILABLE)
+        raise SystemExit(f"Worker health check failed: worker {hostname} did not respond to ping")
 
     log.info("Ping check passed for worker: %s", hostname)
 
-    # Step 2: Active queues check — worker must have registered consumers.
-    # A worker in catatonic state (e.g. after Redis broker restart) responds to
-    # ping but has silently lost its queue consumer registrations, so
-    # active_queues() returns None or an empty list for that worker.
-    log.info("Checking active queue consumers for worker: %s", hostname)
-    active_queues_result = inspect.active_queues()
+    try:
+        # Step 2: Active queues check — worker must have registered consumers.
+        # A worker in catatonic state (e.g. after Redis broker restart) responds to
+        # ping but has silently lost its queue consumer registrations, so
+        # active_queues() returns None or an empty list for that worker.
+        log.info("Checking active queue consumers for worker: %s", hostname)
+        active_queues_result = inspect.active_queues()
+    except Exception:
+        log.exception(
+            "Worker health check failed: unable to reach broker while checking queues for %s", hostname
+        )
+        raise SystemExit("Worker health check failed: broker connection error during active_queues check")
+
     if not active_queues_result or hostname not in active_queues_result:
         log.error(
             "Worker health check failed: worker %s is alive but has no registered queue consumers. "
             "This may indicate a catatonic state caused by a broker restart (e.g. Redis). "
             "The worker will not pick up new tasks until restarted.",
             hostname,
         )
-        raise SystemExit(os.EX_UNAVAILABLE)
+        raise SystemExit(
+            f"Worker health check failed: worker {hostname} has no active queue consumers (catatonic state)"
+        )
 
     worker_queues = active_queues_result[hostname]
     if not worker_queues:
@@ -532,7 +544,9 @@ def worker_health_check(args):
             "The worker will not pick up new tasks until restarted.",
             hostname,
         )
-        raise SystemExit(os.EX_UNAVAILABLE)
+        raise SystemExit(
+            f"Worker health check failed: worker {hostname} has empty queue consumer list (catatonic state)"
+        )
 
     queue_names = [queue["name"] for queue in worker_queues if "name" in queue]
     log.info(

providers/celery/tests/unit/celery/cli/test_celery_command.py

@@ -688,7 +688,7 @@ def test_health_check_fails_when_ping_returns_none(self, mock_inspect):
         with pytest.raises(SystemExit) as exc_info:
             celery_command.worker_health_check(args)
 
-        assert exc_info.value.code != 0
+        assert "did not respond to ping" in str(exc_info.value)
         # active_queues must NOT be called if ping already failed
         mock_instance.active_queues.assert_not_called()
 
@@ -707,9 +707,43 @@ def test_health_check_fails_when_worker_not_in_ping_result(self, mock_inspect):
         with pytest.raises(SystemExit) as exc_info:
             celery_command.worker_health_check(args)
 
-        assert exc_info.value.code != 0
+        assert "did not respond to ping" in str(exc_info.value)
         mock_instance.active_queues.assert_not_called()
 
+    @pytest.mark.db_test
+    @mock.patch("airflow.providers.celery.executors.celery_executor.app.control.inspect")
+    def test_health_check_fails_when_broker_unreachable(self, mock_inspect):
+        """Broker is down — inspect.ping() raises an exception. Health check must fail cleanly."""
+        hostname = "celery@disconnected-host"
+        args = self.parser.parse_args(["celery", "worker-health-check", "-H", hostname])
+
+        mock_instance = MagicMock()
+        mock_instance.ping.side_effect = ConnectionError("Connection refused")
+        mock_inspect.return_value = mock_instance
+
+        with pytest.raises(SystemExit) as exc_info:
+            celery_command.worker_health_check(args)
+
+        assert "broker connection error" in str(exc_info.value)
+        mock_instance.active_queues.assert_not_called()
+
+    @pytest.mark.db_test
+    @mock.patch("airflow.providers.celery.executors.celery_executor.app.control.inspect")
+    def test_health_check_fails_when_active_queues_raises(self, mock_inspect):
+        """Broker fails during active_queues check — health check must fail cleanly."""
+        hostname = "celery@flaky-host"
+        args = self.parser.parse_args(["celery", "worker-health-check", "-H", hostname])
+
+        mock_instance = MagicMock()
+        mock_instance.ping.return_value = {hostname: {"ok": "pong"}}
+        mock_instance.active_queues.side_effect = ConnectionError("Connection lost")
+        mock_inspect.return_value = mock_instance
+
+        with pytest.raises(SystemExit) as exc_info:
+            celery_command.worker_health_check(args)
+
+        assert "broker connection error" in str(exc_info.value)
+
     @pytest.mark.db_test
     @mock.patch("airflow.providers.celery.executors.celery_executor.app.control.inspect")
     def test_health_check_fails_when_active_queues_returns_none(self, mock_inspect):
@@ -735,7 +769,7 @@ def test_health_check_fails_when_active_queues_returns_none(self, mock_inspect):
         with pytest.raises(SystemExit) as exc_info:
             celery_command.worker_health_check(args)
 
-        assert exc_info.value.code != 0
+        assert "catatonic state" in str(exc_info.value)
         mock_instance.ping.assert_called_once()
         mock_instance.active_queues.assert_called_once()
 
@@ -760,7 +794,7 @@ def test_health_check_fails_when_worker_not_in_active_queues(self, mock_inspect)
         with pytest.raises(SystemExit) as exc_info:
             celery_command.worker_health_check(args)
 
-        assert exc_info.value.code != 0
+        assert "catatonic state" in str(exc_info.value)
 
     @pytest.mark.db_test
     @mock.patch("airflow.providers.celery.executors.celery_executor.app.control.inspect")
@@ -783,7 +817,7 @@ def test_health_check_fails_when_active_queues_list_is_empty(self, mock_inspect)
         with pytest.raises(SystemExit) as exc_info:
             celery_command.worker_health_check(args)
 
-        assert exc_info.value.code != 0
+        assert "catatonic state" in str(exc_info.value)
 
     @pytest.mark.db_test
     @mock.patch("socket.gethostname")
@@ -805,7 +839,7 @@ def test_health_check_defaults_hostname_from_socket(self, mock_inspect, mock_get
         celery_command.worker_health_check(args)
 
         # Verify inspect was targeted at the auto-resolved hostname
-        mock_inspect.assert_called_once_with(destination=[expected_hostname])
+        mock_inspect.assert_called_once_with(destination=[expected_hostname], timeout=5.0)
 
 
 class TestLoggerSetupHandler: