You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by as...@apache.org on 2021/04/01 13:47:26 UTC
[airflow] 11/11: Fix bug in airflow.stats timing that broke
dogstatsd mode (#15132)
This is an automated email from the ASF dual-hosted git repository.
ash pushed a commit to branch v2-0-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 663985d8972a7ddb66684a413f17a29278ee509d
Author: Ash Berlin-Taylor <as...@firemirror.com>
AuthorDate: Thu Apr 1 14:46:32 2021 +0100
Fix bug in airflow.stats timing that broke dogstatsd mode (#15132)
The fix for this was very easy -- just a `timer` -> `timed` typo.
However it turns out that the tests for airflow.stats were insufficient
and didn't catch this, so I have extended the tests in two ways:
1. Test all the other stat methods than just incr (guage, timer, timing,
decr)
2. Use "auto-specing" feature of Mock to ensure that we can't make up
methods to call on a mock object.
> Autospeccing is based on the existing spec feature of mock.
> It limits the api of mocks to the api of an original object (the
> spec), but it is recursive (implemented lazily) so that attributes of
> mocks only have the same api as the attributes of the spec. In
> addition mocked functions / methods have the same call signature as
> the original so they raise a TypeError if they are called
> incorrectly.
(cherry picked from commit b7cd2df056ac3ab113d77c5f6b65f02a77337907)
---
airflow/stats.py | 2 +-
tests/core/test_stats.py | 193 +++++++++++++++++++++++++----------------------
2 files changed, 104 insertions(+), 91 deletions(-)
diff --git a/airflow/stats.py b/airflow/stats.py
index 8207220..34677da 100644
--- a/airflow/stats.py
+++ b/airflow/stats.py
@@ -343,7 +343,7 @@ class SafeDogStatsdLogger:
"""Timer metric that can be cancelled"""
if stat and self.allow_list_validator.test(stat):
tags = tags or []
- return Timer(self.dogstatsd.timer(stat, *args, tags=tags, **kwargs))
+ return Timer(self.dogstatsd.timed(stat, *args, tags=tags, **kwargs))
return Timer()
diff --git a/tests/core/test_stats.py b/tests/core/test_stats.py
index b635e62..83169e2 100644
--- a/tests/core/test_stats.py
+++ b/tests/core/test_stats.py
@@ -31,17 +31,7 @@ from tests.test_utils.config import conf_vars
class CustomStatsd(statsd.StatsClient):
- incr_calls = 0
-
- def __init__(self, host=None, port=None, prefix=None):
- super().__init__()
-
- def incr(self, stat, count=1, rate=1): # pylint: disable=unused-argument
- CustomStatsd.incr_calls += 1
-
- @classmethod
- def _reset(cls):
- cls.incr_calls = 0
+ pass
class InvalidCustomStatsd:
@@ -50,25 +40,14 @@ class InvalidCustomStatsd:
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.statsd_client = Mock(spec=statsd.StatsClient)
self.stats = SafeStatsdLogger(self.statsd_client)
- CustomStatsd._reset()
- InvalidCustomStatsd._reset()
def test_increment_counter_with_valid_name(self):
self.stats.incr('test_stats_run')
@@ -86,49 +65,56 @@ class TestStats(unittest.TestCase):
self.stats.incr('test/$tats')
self.statsd_client.assert_not_called()
- @conf_vars({('metrics', '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)
+ def test_timer(self):
+ with self.stats.timer("dummy_timer"):
+ pass
+ self.statsd_client.timer.assert_called_once_with('dummy_timer')
- @conf_vars({('metrics', '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 test_empty_timer(self):
+ with self.stats.timer():
+ pass
+ self.statsd_client.timer.assert_not_called()
- @conf_vars(
- {
- ("metrics", "statsd_on"): "True",
- ("metrics", "statsd_custom_client_path"): "tests.core.test_stats.CustomStatsd",
- }
- )
- def test_load_custom_statsd_client(self):
+ 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)
- assert 'CustomStatsd' == type(airflow.stats.Stats.statsd).__name__ # noqa: E721
- @conf_vars(
- {
- ("metrics", "statsd_on"): "True",
- ("metrics", "statsd_custom_client_path"): "tests.core.test_stats.CustomStatsd",
- }
- )
- def test_does_use_custom_statsd_client(self):
+ 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)
- airflow.stats.Stats.incr("dummy_key")
- assert airflow.stats.Stats.statsd.incr_calls == 1
- @conf_vars(
- {
- ("metrics", "statsd_on"): "True",
- ("metrics", "statsd_custom_client_path"): "tests.core.test_stats.InvalidCustomStatsd",
- }
- )
def test_load_invalid_custom_stats_client(self):
- with pytest.raises(
+ 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.'
@@ -137,15 +123,15 @@ class TestStats(unittest.TestCase):
):
importlib.reload(airflow.stats)
airflow.stats.Stats.incr("dummy_key")
-
- def tearDown(self) -> None:
- # To avoid side-effect
importlib.reload(airflow.stats)
class TestDogStats(unittest.TestCase):
def setUp(self):
- self.dogstatsd_client = Mock()
+ 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):
@@ -166,48 +152,72 @@ class TestDogStats(unittest.TestCase):
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(
+ 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
)
- @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(
+ 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
)
- @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(
+ 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
)
- @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):
+ 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)
- airflow.stats.Stats.incr("dummy_key")
- mock_statsd.return_value.assert_not_called()
- def tearDown(self) -> None:
- # To avoid side-effect
+ 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)
class TestStatsWithAllowList(unittest.TestCase):
def setUp(self):
- self.statsd_client = Mock()
+ self.statsd_client = Mock(spec=statsd.StatsClient)
self.stats = SafeStatsdLogger(self.statsd_client, AllowListValidator("stats_one, stats_two"))
def test_increment_counter_with_allowed_key(self):
@@ -225,7 +235,10 @@ class TestStatsWithAllowList(unittest.TestCase):
class TestDogStatsWithAllowList(unittest.TestCase):
def setUp(self):
- self.dogstatsd_client = Mock()
+ pytest.importorskip('datadog')
+ from datadog import DogStatsd
+
+ self.dogstatsd_client = Mock(speck=DogStatsd)
self.dogstats = SafeDogStatsdLogger(self.dogstatsd_client, AllowListValidator("stats_one, stats_two"))
def test_increment_counter_with_allowed_key(self):