Repository: apache/airflow
File: shared/observability/tests/observability/metrics/test_stats.py (Python)

Code before:
            tags={"key0,": "val0", "key1": "val1", "key2": "val2", "key3": "val3"},
        )
        self.statsd_client.incr.assert_called_once_with("test_stats_run.delay,key1=val1", 1, 1)


def always_invalid(stat_name):
    raise InvalidStatsNameException(f"Invalid name: {stat_name}")


def always_valid(stat_name):
    return stat_name


class TestCustomStatsName:
    @conf_vars(
        {
            ("metrics", "statsd_on"): "True",
            ("metrics", "stat_name_handler"): "unit.observability.metrics.test_stats.always_invalid",
        }
    )
    @mock.patch("statsd.StatsClient")
    def test_does_not_send_stats_using_statsd_when_the_name_is_not_valid(self, mock_statsd):
        importlib.reload(airflow_shared.observability.metrics.stats)
        airflow_shared.observability.metrics.stats.Stats.incr("empty_key")
        mock_statsd.return_value.assert_not_called()

    @skip_if_force_lowest_dependencies_marker
    @conf_vars(
        {
            ("metrics", "statsd_datadog_enabled"): "True",
            (
                "metrics",
                "stat_name_handler",
            ): "observability.metrics.test_stats.always_invalid",
        }
    )
    @mock.patch("datadog.DogStatsd")
    def test_does_not_send_stats_using_dogstatsd_when_the_name_is_not_valid(self, mock_dogstatsd):
        importlib.reload(airflow_shared.observability.metrics.stats)
        airflow_shared.observability.metrics.stats.Stats.incr("empty_key")
        mock_dogstatsd.return_value.assert_not_called()

    @mock.patch("statsd.StatsClient")
    def test_does_send_stats_using_statsd_when_the_name_is_valid(self, mock_statsd):
        importlib.reload(airflow_shared.observability.metrics.stats)
        airflow_shared.observability.metrics.stats.Stats.initialize(
            factory=get_statsd_logger_factory(
                stats_class=mock_statsd,
                stat_name_handler=always_valid,
            )
        )

Review comment (refactor):
Now that `get_dogstatsd_logger` is used more that once, it could be refactored into a method like `get_statsd_logger_factory`. WDYT?

Code after:
    raise InvalidStatsNameException(f"Invalid name: {stat_name}")


def always_valid(stat_name):
    return stat_name


class TestCustomStatsName:
    def test_does_not_send_stats_using_statsd_when_the_name_is_not_valid(self):
        with mock.patch("statsd.StatsClient") as mock_statsd:
            importlib.reload(airflow_shared.observability.metrics.stats)
            airflow_shared.observability.metrics.stats.Stats.initialize(
                factory=get_statsd_logger_factory(
                    stats_class=mock_statsd,
                    stat_name_handler=always_invalid,
                )
            )
            airflow_shared.observability.metrics.stats.Stats.incr("empty_key")
            mock_statsd.return_value.assert_not_called()

    @skip_if_force_lowest_dependencies_marker
    def test_does_not_send_stats_using_dogstatsd_when_the_name_is_not_valid(self):
        with mock.patch("datadog.DogStatsd") as mock_dogstatsd:
            importlib.reload(airflow_shared.observability.metrics.stats)
            airflow_shared.observability.metrics.stats.Stats.initialize(
                factory=lambda: datadog_logger.get_dogstatsd_logger(
                    mock_dogstatsd,
                    host="localhost",
                    port="1234",
                    namespace="airflow",
                    stat_name_handler=always_invalid,
                )
            )
            airflow_shared.observability.metrics.stats.Stats.incr("empty_key")
            mock_dogstatsd.return_value.assert_not_called()

    def test_does_send_stats_using_statsd_when_the_name_is_valid(self):
        with mock.patch("statsd.StatsClient") as mock_statsd:
            importlib.reload(airflow_shared.observability.metrics.stats)
            airflow_shared.observability.metrics.stats.Stats.initialize(
                factory=get_statsd_logger_factory(
                    stats_class=mock_statsd,
                    stat_name_handler=always_valid,
                )
            )
            airflow_shared.observability.metrics.stats.Stats.incr("empty_key")
            mock_statsd.return_value.incr.assert_called_once_with("empty_key", 1, 1)

    @skip_if_force_lowest_dependencies_marker
    def test_does_send_stats_using_dogstatsd_when_the_name_is_valid(self):
        with mock.patch("datadog.DogStatsd") as mock_dogstatsd: