Ejemplo n.º 1
0
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()
Ejemplo n.º 2
0
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()
Ejemplo n.º 3
0
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)
Ejemplo n.º 4
0
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()
Ejemplo n.º 5
0
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)
Ejemplo n.º 6
0
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)