class TestDogStats(unittest.TestCase): def setUp(self): self.dogstatsd_client = Mock() self.dogstatsd = SafeDogStatsdLogger(self.dogstatsd_client) def test_increment_counter_with_valid_name_with_dogstatsd(self): self.dogstatsd.incr('test_stats_run') self.dogstatsd_client.increment.assert_called_once_with( metric='test_stats_run', sample_rate=1, tags=[], value=1 ) def test_stat_name_must_be_a_string_with_dogstatsd(self): self.dogstatsd.incr([]) self.dogstatsd_client.assert_not_called() def test_stat_name_must_not_exceed_max_length_with_dogstatsd(self): self.dogstatsd.incr('X' * 300) self.dogstatsd_client.assert_not_called() def test_stat_name_must_only_include_allowed_characters_with_dogstatsd(self): self.dogstatsd.incr('test/$tats') self.dogstatsd_client.assert_not_called() @conf_vars({('metrics', 'statsd_datadog_enabled'): 'True'}) @mock.patch("datadog.DogStatsd") def test_does_send_stats_using_dogstatsd_when_dogstatsd_on(self, mock_dogstatsd): importlib.reload(airflow.stats) airflow.stats.Stats.incr("dummy_key") mock_dogstatsd.return_value.increment.assert_called_once_with( metric='dummy_key', sample_rate=1, tags=[], value=1 ) @conf_vars({('metrics', 'statsd_datadog_enabled'): 'True'}) @mock.patch("datadog.DogStatsd") def test_does_send_stats_using_dogstatsd_with_tags(self, mock_dogstatsd): importlib.reload(airflow.stats) airflow.stats.Stats.incr("dummy_key", 1, 1, ['key1:value1', 'key2:value2']) mock_dogstatsd.return_value.increment.assert_called_once_with( metric='dummy_key', sample_rate=1, tags=['key1:value1', 'key2:value2'], value=1 ) @conf_vars({('metrics', 'statsd_on'): 'True', ('metrics', 'statsd_datadog_enabled'): 'True'}) @mock.patch("datadog.DogStatsd") def test_does_send_stats_using_dogstatsd_when_statsd_and_dogstatsd_both_on(self, mock_dogstatsd): importlib.reload(airflow.stats) airflow.stats.Stats.incr("dummy_key") mock_dogstatsd.return_value.increment.assert_called_once_with( metric='dummy_key', sample_rate=1, tags=[], value=1 ) @conf_vars({('metrics', 'statsd_on'): 'True', ('metrics', 'statsd_datadog_enabled'): 'True'}) @mock.patch("statsd.StatsClient") def test_does_not_send_stats_using_statsd_when_statsd_and_dogstatsd_both_on(self, mock_statsd): importlib.reload(airflow.stats) airflow.stats.Stats.incr("dummy_key") mock_statsd.return_value.assert_not_called() def tearDown(self) -> None: # To avoid side-effect importlib.reload(airflow.stats)
class TestDogStatsWithAllowList(unittest.TestCase): def setUp(self): self.dogstatsd_client = Mock() self.dogstats = SafeDogStatsdLogger( self.dogstatsd_client, AllowListValidator("stats_one, stats_two")) def test_increment_counter_with_allowed_key(self): self.dogstats.incr('stats_one') self.dogstatsd_client.increment.assert_called_once_with( metric='stats_one', sample_rate=1, tags=[], value=1) def test_increment_counter_with_allowed_prefix(self): self.dogstats.incr('stats_two.bla') self.dogstatsd_client.increment.assert_called_once_with( metric='stats_two.bla', sample_rate=1, tags=[], value=1) def test_not_increment_counter_if_not_allowed(self): self.dogstats.incr('stats_three') self.dogstatsd_client.assert_not_called()
class TestDogStats(unittest.TestCase): def setUp(self): pytest.importorskip('datadog') from datadog import DogStatsd self.dogstatsd_client = Mock(spec=DogStatsd) self.dogstatsd = SafeDogStatsdLogger(self.dogstatsd_client) def test_increment_counter_with_valid_name_with_dogstatsd(self): self.dogstatsd.incr('test_stats_run') self.dogstatsd_client.increment.assert_called_once_with( metric='test_stats_run', sample_rate=1, tags=[], value=1) def test_stat_name_must_be_a_string_with_dogstatsd(self): self.dogstatsd.incr([]) self.dogstatsd_client.assert_not_called() def test_stat_name_must_not_exceed_max_length_with_dogstatsd(self): self.dogstatsd.incr('X' * 300) self.dogstatsd_client.assert_not_called() def test_stat_name_must_only_include_allowed_characters_with_dogstatsd( self): self.dogstatsd.incr('test/$tats') self.dogstatsd_client.assert_not_called() def test_does_send_stats_using_dogstatsd_when_dogstatsd_on(self): self.dogstatsd.incr("dummy_key") self.dogstatsd_client.increment.assert_called_once_with( metric='dummy_key', sample_rate=1, tags=[], value=1) def test_does_send_stats_using_dogstatsd_with_tags(self): self.dogstatsd.incr("dummy_key", 1, 1, ['key1:value1', 'key2:value2']) self.dogstatsd_client.increment.assert_called_once_with( metric='dummy_key', sample_rate=1, tags=['key1:value1', 'key2:value2'], value=1) def test_does_send_stats_using_dogstatsd_when_statsd_and_dogstatsd_both_on( self): self.dogstatsd.incr("dummy_key") self.dogstatsd_client.increment.assert_called_once_with( metric='dummy_key', sample_rate=1, tags=[], value=1) def test_timer(self): with self.dogstatsd.timer("dummy_timer"): pass self.dogstatsd_client.timed.assert_called_once_with('dummy_timer', tags=[]) def test_empty_timer(self): with self.dogstatsd.timer(): pass self.dogstatsd_client.timed.assert_not_called() def test_timing(self): self.dogstatsd.timing("dummy_timer", 123) self.dogstatsd_client.timing.assert_called_once_with( metric='dummy_timer', value=123, tags=[]) def test_gauge(self): self.dogstatsd.gauge("dummy", 123) self.dogstatsd_client.gauge.assert_called_once_with(metric='dummy', sample_rate=1, value=123, tags=[]) def test_decr(self): self.dogstatsd.decr("dummy") self.dogstatsd_client.decrement.assert_called_once_with(metric='dummy', sample_rate=1, value=1, tags=[]) def test_enabled_by_config(self): """Test that enabling this sets the right instance properties""" from datadog import DogStatsd with conf_vars({('metrics', 'statsd_datadog_enabled'): 'True'}): importlib.reload(airflow.stats) assert isinstance(airflow.stats.Stats.dogstatsd, DogStatsd) assert not hasattr(airflow.stats.Stats, 'statsd') # Avoid side-effects importlib.reload(airflow.stats) def test_does_not_send_stats_using_statsd_when_statsd_and_dogstatsd_both_on( self): from datadog import DogStatsd with conf_vars({ ('metrics', 'statsd_on'): 'True', ('metrics', 'statsd_datadog_enabled'): 'True' }): importlib.reload(airflow.stats) assert isinstance(airflow.stats.Stats.dogstatsd, DogStatsd) assert not hasattr(airflow.stats.Stats, 'statsd') importlib.reload(airflow.stats)