[AIRFLOW-6530]: Add tests for loading custom statsd client

airflow/stats.py

@@ -200,6 +200,7 @@ def __init__(self, *args, **kwargs):
                     self.__class__.instance = DummyStatsLogger()
             except (socket.gaierror, ImportError) as e:
                 log.warning("Could not configure StatsClient: %s, using DummyStatsLogger instead.", e)
+                self.__class__.instance = DummyStatsLogger()
 
     def get_statsd_logger(self):
         if conf.getboolean('scheduler', 'statsd_on'):
@@ -217,6 +218,9 @@ def get_statsd_logger(self):
                         raise Exception(
                             """Your custom Statsd client must extend the statsd.StatsClient in order to ensure backwards
                             compatibility.""")
+                    else:
+                        log.info("Successfully loaded custom Statsd client "
+                                 f"from {custom_statsd_module_path}")
 
                 except Exception as err:
                     raise ImportError('Unable to load custom Statsd client from '

tests/test_stats.py

@@ -20,17 +20,53 @@
 from unittest import mock
 from unittest.mock import Mock
 
+import statsd
+
 import airflow
 from airflow.exceptions import InvalidStatsNameException
 from airflow.stats import AllowListValidator, SafeDogStatsdLogger, SafeStatsdLogger
 from tests.test_utils.config import conf_vars
 
 
+class CustomStatsd(statsd.StatsClient):
+    incr_calls = 0
+
+    def __init__(self, host=None, port=None, prefix=None):
+        pass
+
+    def incr(self, stat, count=1, rate=1):  # pylint: disable=unused-argument
+        CustomStatsd.incr_calls += 1
+
+    @classmethod
+    def _reset(cls):
+        cls.incr_calls = 0
+
+
+class InvalidCustomStatsd:
+    """
+    This custom Statsd class is invalid because it does not subclass
+    statsd.StatsClient.
+    """
+    incr_calls = 0
+
+    def __init__(self, host=None, port=None, prefix=None):
+        pass
+
+    def incr(self, stat, count=1, rate=1):  # pylint: disable=unused-argument
+        InvalidCustomStatsd.incr_calls += 1
+
+    @classmethod
+    def _reset(cls):
+        cls.incr_calls = 0
+
+
 class TestStats(unittest.TestCase):
 
     def setUp(self):
         self.statsd_client = Mock()
         self.stats = SafeStatsdLogger(self.statsd_client)
+        CustomStatsd._reset()
+        InvalidCustomStatsd._reset()
 
     def test_increment_counter_with_valid_name(self):
         self.stats.incr('test_stats_run')
@@ -66,6 +102,32 @@ def test_does_not_send_stats_using_dogstatsd(self, mock_dogstatsd):
         airflow.stats.Stats.incr("dummy_key")
         mock_dogstatsd.return_value.assert_not_called()
 
+    @conf_vars({
+        ("scheduler", "statsd_on"): "True",
+        ("scheduler", "statsd_custom_client_path"): "tests.test_stats.CustomStatsd",
+    })
+    def test_load_custom_statsd_client(self):
+        importlib.reload(airflow.stats)
+        assert isinstance(airflow.stats.Stats.statsd, CustomStatsd)
+
+    @conf_vars({
+        ("scheduler", "statsd_on"): "True",
+        ("scheduler", "statsd_custom_client_path"): "tests.test_stats.CustomStatsd",
+    })
+    def test_does_use_custom_statsd_client(self):
+        importlib.reload(airflow.stats)
+        airflow.stats.Stats.incr("dummy_key")
+        assert CustomStatsd.incr_calls == 1
+
+    @conf_vars({
+        ("scheduler", "statsd_on"): "True",
+        ("scheduler", "statsd_custom_client_path"): "tests.test_stats.InvalidCustomStatsd",
+    })
+    def test_load_invalid_custom_stats_client(self):
+        importlib.reload(airflow.stats)
+        airflow.stats.Stats.incr("dummy_key")
+        assert InvalidCustomStatsd.incr_calls == 0
+
     def tearDown(self) -> None:
         # To avoid side-effect
         importlib.reload(airflow.stats)