class TestStats(unittest.TestCase):
def setUp(self):
self.statsd_client = Mock()
self.stats = SafeStatsdLogger(self.statsd_client)
def test_increment_counter_with_valid_name(self):
self.stats.incr('test_stats_run')
self.statsd_client.incr.assert_called_once_with('test_stats_run', 1, 1)
def test_stat_name_must_be_a_string(self):
self.stats.incr(list())
self.statsd_client.assert_not_called()
def test_stat_name_must_not_exceed_max_length(self):
self.stats.incr('X' * 300)
self.statsd_client.assert_not_called()
def test_stat_name_must_only_include_whitelisted_characters(self):
self.stats.incr('test/$tats')
self.statsd_client.assert_not_called()
class TestStats(unittest.TestCase):
def setUp(self):
self.statsd_client = Mock()
self.stats = SafeStatsdLogger(self.statsd_client)
def test_increment_counter_with_valid_name(self):
self.stats.incr('test_stats_run')
self.statsd_client.incr.assert_called_once_with('test_stats_run', 1, 1)
def test_stat_name_must_be_a_string(self):
self.stats.incr(list())
self.statsd_client.assert_not_called()
def test_stat_name_must_not_exceed_max_length(self):
self.stats.incr('X' * 300)
self.statsd_client.assert_not_called()
def test_stat_name_must_only_include_whitelisted_characters(self):
self.stats.incr('test/$tats')
self.statsd_client.assert_not_called()
class TestStats(unittest.TestCase):
def setUp(self):
self.statsd_client = Mock()
self.stats = SafeStatsdLogger(self.statsd_client)
def test_increment_counter_with_valid_name(self):
self.stats.incr('test_stats_run')
self.statsd_client.incr.assert_called_once_with('test_stats_run', 1, 1)
def test_stat_name_must_be_a_string(self):
self.stats.incr(list())
self.statsd_client.assert_not_called()
def test_stat_name_must_not_exceed_max_length(self):
self.stats.incr('X' * 300)
self.statsd_client.assert_not_called()
def test_stat_name_must_only_include_whitelisted_characters(self):
self.stats.incr('test/$tats')
self.statsd_client.assert_not_called()
@conf_vars({('scheduler', 'statsd_on'): 'True'})
@mock.patch("statsd.StatsClient")
def test_does_send_stats_using_statsd(self, mock_statsd):
importlib.reload(airflow.stats)
airflow.stats.Stats.incr("dummy_key")
mock_statsd.return_value.incr.assert_called_once_with(
'dummy_key', 1, 1)
@conf_vars({('scheduler', 'statsd_on'): 'True'})
@mock.patch("datadog.DogStatsd")
def test_does_not_send_stats_using_dogstatsd(self, mock_dogstatsd):
importlib.reload(airflow.stats)
airflow.stats.Stats.incr("dummy_key")
mock_dogstatsd.return_value.assert_not_called()
def tearDown(self) -> None:
# To avoid side-effect
importlib.reload(airflow.stats)
class TestStatsWithAllowList(unittest.TestCase):
def setUp(self):
self.statsd_client = Mock()
self.stats = SafeStatsdLogger(self.statsd_client, AllowListValidator("stats_one, stats_two"))
def test_increment_counter_with_allowed_key(self):
self.stats.incr('stats_one')
self.statsd_client.incr.assert_called_once_with('stats_one', 1, 1)
def test_increment_counter_with_allowed_prefix(self):
self.stats.incr('stats_two.bla')
self.statsd_client.incr.assert_called_once_with('stats_two.bla', 1, 1)
def test_not_increment_counter_if_not_allowed(self):
self.stats.incr('stats_three')
self.statsd_client.assert_not_called()
def setUp(self):
self.statsd_client = Mock()
self.stats = SafeStatsdLogger(self.statsd_client, AllowListValidator("stats_one, stats_two"))
def setUp(self):
self.statsd_client = Mock()
self.stats = SafeStatsdLogger(self.statsd_client)
def setUp(self):
self.statsd_client = Mock()
self.stats = SafeStatsdLogger(self.statsd_client)
CustomStatsd._reset()
InvalidCustomStatsd._reset()
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')
self.statsd_client.incr.assert_called_once_with('test_stats_run', 1, 1)
def test_stat_name_must_be_a_string(self):
self.stats.incr([])
self.statsd_client.assert_not_called()
def test_stat_name_must_not_exceed_max_length(self):
self.stats.incr('X' * 300)
self.statsd_client.assert_not_called()
def test_stat_name_must_only_include_allowed_characters(self):
self.stats.incr('test/$tats')
self.statsd_client.assert_not_called()
@conf_vars({
('scheduler', 'statsd_on'): 'True'
})
@mock.patch("statsd.StatsClient")
def test_does_send_stats_using_statsd(self, mock_statsd):
importlib.reload(airflow.stats)
airflow.stats.Stats.incr("dummy_key")
mock_statsd.return_value.incr.assert_called_once_with('dummy_key', 1, 1)
@conf_vars({
('scheduler', 'statsd_on'): 'True'
})
@mock.patch("datadog.DogStatsd")
def test_does_not_send_stats_using_dogstatsd(self, mock_dogstatsd):
importlib.reload(airflow.stats)
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.core.test_stats.CustomStatsd",
})
def test_load_custom_statsd_client(self):
importlib.reload(airflow.stats)
self.assertEqual('CustomStatsd', type(airflow.stats.Stats.statsd).__name__)
@conf_vars({
("scheduler", "statsd_on"): "True",
("scheduler", "statsd_custom_client_path"): "tests.core.test_stats.CustomStatsd",
})
def test_does_use_custom_statsd_client(self):
importlib.reload(airflow.stats)
airflow.stats.Stats.incr("dummy_key")
assert airflow.stats.Stats.statsd.incr_calls == 1
@conf_vars({
("scheduler", "statsd_on"): "True",
("scheduler", "statsd_custom_client_path"): "tests.core.test_stats.InvalidCustomStatsd",
})
def test_load_invalid_custom_stats_client(self):
with self.assertRaisesRegex(
AirflowConfigException,
re.escape(
'Your custom Statsd client must extend the statsd.'
'StatsClient in order to ensure backwards compatibility.'
)
):
importlib.reload(airflow.stats)
def tearDown(self) -> None:
# To avoid side-effect
importlib.reload(airflow.stats)
def setUp(self):
self.statsd_client = Mock()
self.stats = SafeStatsdLogger(self.statsd_client)
def setUp(self):
self.statsd_client = Mock(spec=statsd.StatsClient)
self.stats = SafeStatsdLogger(self.statsd_client)
class TestStats(unittest.TestCase):
def setUp(self):
self.statsd_client = Mock(spec=statsd.StatsClient)
self.stats = SafeStatsdLogger(self.statsd_client)
def test_increment_counter_with_valid_name(self):
self.stats.incr('test_stats_run')
self.statsd_client.incr.assert_called_once_with('test_stats_run', 1, 1)
def test_stat_name_must_be_a_string(self):
self.stats.incr([])
self.statsd_client.assert_not_called()
def test_stat_name_must_not_exceed_max_length(self):
self.stats.incr('X' * 300)
self.statsd_client.assert_not_called()
def test_stat_name_must_only_include_allowed_characters(self):
self.stats.incr('test/$tats')
self.statsd_client.assert_not_called()
def test_timer(self):
with self.stats.timer("dummy_timer"):
pass
self.statsd_client.timer.assert_called_once_with('dummy_timer')
def test_empty_timer(self):
with self.stats.timer():
pass
self.statsd_client.timer.assert_not_called()
def test_timing(self):
self.stats.timing("dummy_timer", 123)
self.statsd_client.timing.assert_called_once_with('dummy_timer', 123)
def test_gauge(self):
self.stats.gauge("dummy", 123)
self.statsd_client.gauge.assert_called_once_with(
'dummy', 123, 1, False)
def test_decr(self):
self.stats.decr("dummy")
self.statsd_client.decr.assert_called_once_with('dummy', 1, 1)
def test_enabled_by_config(self):
"""Test that enabling this sets the right instance properties"""
with conf_vars({('metrics', 'statsd_on'): 'True'}):
importlib.reload(airflow.stats)
assert isinstance(airflow.stats.Stats.statsd, statsd.StatsClient)
assert not hasattr(airflow.stats.Stats, 'dogstatsd')
# Avoid side-effects
importlib.reload(airflow.stats)
def test_load_custom_statsd_client(self):
with conf_vars({
("metrics", "statsd_on"):
"True",
("metrics", "statsd_custom_client_path"):
f"{__name__}.CustomStatsd",
}):
importlib.reload(airflow.stats)
assert isinstance(airflow.stats.Stats.statsd, CustomStatsd)
# Avoid side-effects
importlib.reload(airflow.stats)
def test_load_invalid_custom_stats_client(self):
with conf_vars({
("metrics", "statsd_on"):
"True",
("metrics", "statsd_custom_client_path"):
f"{__name__}.InvalidCustomStatsd",
}), pytest.raises(
AirflowConfigException,
match=re.escape(
'Your custom Statsd client must extend the statsd.'
'StatsClient in order to ensure backwards compatibility.'),
):
importlib.reload(airflow.stats)
airflow.stats.Stats.incr("dummy_key")
importlib.reload(airflow.stats)