def setUp(self):
pytest.importorskip('datadog')
from datadog import DogStatsd
self.dogstatsd_client = Mock(speck=DogStatsd)
self.dogstats = SafeDogStatsdLogger(
self.dogstatsd_client, AllowListValidator("stats_one, stats_two"))
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()
def setUp(self):
self.dogstatsd_client = Mock()
self.dogstats = SafeDogStatsdLogger(self.dogstatsd_client, AllowListValidator("stats_one, stats_two"))
def setUp(self):
self.dogstatsd_client = Mock()
self.dogstatsd = SafeDogStatsdLogger(self.dogstatsd_client)
def setUp(self):
pytest.importorskip('datadog')
from datadog import DogStatsd
self.dogstatsd_client = Mock(spec=DogStatsd)
self.dogstatsd = SafeDogStatsdLogger(self.dogstatsd_client)
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)
