Add noop async client

lirsacc · lirsacc · commit 75cbd34fe547 · 2024-06-23T23:21:30.000+01:00
diff --git a/requirements/test.in b/requirements/test.in
@@ -2,3 +2,4 @@ pytest
 pytest-cov
 pytest-xdist
 coverage
+pytest-asyncio
diff --git a/requirements/test.txt b/requirements/test.txt
@@ -17,8 +17,11 @@ pluggy==1.5.0
 pytest==8.2.2
     # via
     #   -r requirements/test.in
+    #   pytest-asyncio
     #   pytest-cov
     #   pytest-xdist
+pytest-asyncio==0.23.7
+    # via -r requirements/test.in
 pytest-cov==5.0.0
     # via -r requirements/test.in
 pytest-xdist==3.6.1
diff --git a/src/statsd/__init__.py b/src/statsd/__init__.py
@@ -4,6 +4,7 @@
 Version: v\ |version|.
 """
 
+from .async_client import BaseAsyncStatsdClient, DebugAsyncStatsdClient
 from .base import Sample
 from .client import (
     BaseStatsdClient,
@@ -15,7 +16,9 @@
 
 
 __all__ = (
+    "BaseAsyncStatsdClient",
     "BaseStatsdClient",
+    "DebugAsyncStatsdClient",
     "DebugStatsdClient",
     "Sample",
     "StatsdClient",
diff --git a/src/statsd/async_client.py b/src/statsd/async_client.py
@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+import abc
+import contextlib
+import functools
+import logging
+import time
+from typing import Any, AsyncIterator, Awaitable, Callable, Mapping, TypeVar
+from typing_extensions import ParamSpec
+
+from statsd.base import AbstractStatsdClient
+
+
+P = ParamSpec("P")
+T = TypeVar("T")
+U = TypeVar("U")
+
+logger = logging.getLogger("statsd")
+
+
+class BaseAsyncStatsdClient(AbstractStatsdClient[Awaitable[None]]):
+    """
+    Base async client.
+
+    This class exposes the public interface and takes care of packet formatting
+    as well as sampling. It does not actually send packets anywhere, which is
+    left to concrete subclasses implementing :meth:`_emit`.
+    """
+
+    @abc.abstractmethod
+    async def _emit(self, packets: list[str]) -> None:
+        """
+        Async send implementation.
+
+        This method is responsible for actually sending the formatted packets
+        and should be implemented by all subclasses.
+
+        It may batch or buffer packets but should not modify them in any way. It
+        should be agnostic to the Statsd format.
+        """
+        raise NotImplementedError()
+
+    def timed(
+        self,
+        name: str | None = None,
+        *,
+        tags: Mapping[str, str] | None = None,
+        sample_rate: float | None = None,
+        use_distribution: bool = False,
+    ) -> Callable[[Callable[P, Awaitable[U]]], Callable[P, Awaitable[U]]]:
+        """
+        Wrap a function to record its execution time.
+
+        This just wraps the function call with a :meth:`timer` context manager.
+
+        If a name is not provided, the function name will be used.
+
+        Passing ``use_distribution=True`` will report the value as a globally
+        aggregated :meth:`distribution` metric instead of a :meth:`timing`
+        metric.
+
+        >>> client = AsyncStatsdClient()
+        >>> @client.timed()
+        ... async def do_something():
+        ...     pass
+        """
+
+        def decorator(
+            fn: Callable[P, Awaitable[U]],
+        ) -> Callable[P, Awaitable[U]]:
+            # TODO: Should the fallback include the module? Class (for methods)?
+            # or func.__name__
+            metric_name = name or fn.__name__
+
+            @functools.wraps(fn)
+            async def wrapped(*args: P.args, **kwargs: P.kwargs) -> U:
+                async with self.timer(
+                    metric_name,
+                    tags=tags,
+                    use_distribution=use_distribution,
+                    sample_rate=sample_rate,
+                ):
+                    return await fn(*args, **kwargs)
+
+            return wrapped
+
+        return decorator
+
+    @contextlib.asynccontextmanager
+    async def timer(
+        self,
+        name: str,
+        *,
+        tags: Mapping[str, str] | None = None,
+        sample_rate: float | None = None,
+        use_distribution: bool = False,
+    ) -> AsyncIterator[None]:
+        """
+        Context manager to measure the execution time of an async block.
+
+        Passing ``use_distribution=True`` will report the value as a globally
+        aggregated :meth:`distribution` metric instead of a :meth:`timing`
+        metric.
+
+        >>> client = AsyncStatsdClient()
+        >>> async def operation():
+        ...     async with client.timer("download_duration"):
+        ...         pass
+        """
+        start = time.perf_counter()
+        try:
+            yield
+        finally:
+            duration_ms = int(1000 * (time.perf_counter() - start))
+            if use_distribution:
+                await self.distribution(
+                    name,
+                    duration_ms,
+                    tags=tags,
+                    sample_rate=sample_rate,
+                )
+            else:
+                await self.timing(
+                    name,
+                    duration_ms,
+                    tags=tags,
+                    sample_rate=sample_rate,
+                )
+
+
+class DebugAsyncStatsdClient(BaseAsyncStatsdClient):
+    """
+    Verbose client for development or debugging purposes.
+
+    All Statsd packets will be logged and optionally forwarded to a wrapped
+    client.
+    """
+
+    def __init__(
+        self,
+        level: int = logging.INFO,
+        logger: logging.Logger = logger,
+        inner: BaseAsyncStatsdClient | None = None,
+        **kwargs: Any,
+    ) -> None:
+        r"""
+        Initialize DebugStatsdClient.
+
+        :param level: Log level to use, defaults to ``INFO``.
+
+        :param logger: Logger instance to use, defaults to ``statsd``.
+
+        :param inner: Wrapped client.
+
+        :param \**kwargs: Extra arguments forwarded to :class:`BaseAsyncStatsdClient`.
+        """
+        super().__init__(**kwargs)
+        self.level = level
+        self.logger = logger
+        self.inner = inner
+
+    async def _emit(self, packets: list[str]) -> None:
+        for packet in packets:
+            self.logger.log(self.level, "> %s", packet)
+        if self.inner:
+            await self.inner._emit(packets)
+
+
+AsyncStatsdClient = DebugAsyncStatsdClient
diff --git a/tests/test_base_async_client.py b/tests/test_base_async_client.py
@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+import logging
+from typing import Any
+from unittest import mock
+
+import pytest
+
+from statsd import BaseAsyncStatsdClient, DebugAsyncStatsdClient
+
+
+class MockClient(BaseAsyncStatsdClient):
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.mock = mock.Mock()
+
+    async def _emit(self, packets: list[str]) -> None:
+        if packets:
+            self.mock(packets)
+
+    def assert_emitted(self, expected: list[str] | str) -> None:
+        self.mock.assert_called_once_with(
+            expected if isinstance(expected, list) else [expected],
+        )
+
+    def assert_did_not_emit(self) -> None:
+        self.mock.assert_not_called()
+
+
+@pytest.mark.asyncio()
+async def test_timed_decorator() -> None:
+    client = MockClient()
+
+    @client.timed("foo", tags={"foo": "1"})
+    async def fn() -> None:
+        pass
+
+    with mock.patch(
+        "time.perf_counter",
+        side_effect=[7.886838544, 20.181117592],
+    ):
+        await fn()
+
+    client.mock.assert_called_once_with(["foo:12294|ms|#foo:1"])
+
+
+@pytest.mark.asyncio()
+async def test_timed_decorator_use_distribution() -> None:
+    client = MockClient()
+
+    @client.timed("foo", tags={"foo": "1"}, use_distribution=True)
+    async def fn() -> None:
+        pass
+
+    with mock.patch(
+        "time.perf_counter",
+        side_effect=[7.886838544, 20.181117592],
+    ):
+        await fn()
+
+    client.mock.assert_called_once_with(["foo:12294|d|#foo:1"])
+
+
+SIMPLE_TEST_CASES: list[
+    tuple[str, tuple[Any, ...], dict[str, Any], list[str] | str]
+] = [
+    ("increment", ("foo",), {}, "foo:1|c"),
+    ("increment", ("foo", 10), {}, "foo:10|c"),
+]
+
+
+@pytest.mark.parametrize(
+    ("method", "args", "kwargs", "expected"),
+    SIMPLE_TEST_CASES,
+)
+@pytest.mark.asyncio()
+async def test_debug_client_no_inner(
+    method: str,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+    expected: list[str] | str,
+    caplog: Any,
+) -> None:
+    client = DebugAsyncStatsdClient()
+    with caplog.at_level(logging.INFO, logger="statsd"):
+        await getattr(client, method)(*args, **kwargs)
+
+    if isinstance(expected, list):
+        assert len(caplog.records) == len(expected)
+        for x in expected:
+            assert x in caplog.text
+    else:
+        assert len(caplog.records) == 1
+        assert expected in caplog.text
+
+
+@pytest.mark.parametrize(
+    ("method", "args", "kwargs", "expected"),
+    SIMPLE_TEST_CASES,
+)
+@pytest.mark.asyncio()
+async def test_debug_client(
+    method: str,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+    expected: str | list[str],
+) -> None:
+    inner = MockClient()
+    client = DebugAsyncStatsdClient(inner=inner)
+    await getattr(client, method)(*args, **kwargs)
+    inner.assert_emitted(expected)
+
+
+@pytest.mark.parametrize(
+    ("method", "args", "kwargs", "expected"),
+    SIMPLE_TEST_CASES,
+)
+@pytest.mark.asyncio()
+async def test_debug_client_custom_logger_and_level(
+    method: str,
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+    expected: list[str] | str,
+    caplog: Any,
+) -> None:
+    client = DebugAsyncStatsdClient(
+        logger=logging.getLogger("foo"),
+        level=logging.DEBUG,
+    )
+    with caplog.at_level(logging.DEBUG, logger="foo"):
+        await getattr(client, method)(*args, **kwargs)
+
+    if isinstance(expected, list):
+        assert len(caplog.records) == len(expected)
+        for x in expected:
+            assert x in caplog.text
+    else:
+        assert len(caplog.records) == 1
+        assert expected in caplog.text
