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)