From f97ba6d816d4a94dd22234390cdb3ee4821bc918 Mon Sep 17 00:00:00 2001
From: Aviram Hassan <aviram.hassan@biocatch.com>
Date: Wed, 9 Dec 2020 15:42:58 +0200
Subject: [PATCH] [feat] add threshold_ms for conditional timing metrics

---
 aiodogstatsd/client.py | 4 +++-
 examples/timeit.py     | 2 ++
 tests/test_client.py   | 6 ++++++
 3 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/aiodogstatsd/client.py b/aiodogstatsd/client.py
index 980bbaa..4c1e9fa 100644
--- a/aiodogstatsd/client.py
+++ b/aiodogstatsd/client.py
@@ -257,6 +257,7 @@ def timeit(
         *,
         tags: Optional[typedefs.MTags] = None,
         sample_rate: Optional[typedefs.MSampleRate] = None,
+        threshold_ms: typedefs.MValue = 0
     ) -> Iterator[None]:
         """
         Context manager for easily timing methods.
@@ -268,7 +269,8 @@ def timeit(
             yield
         finally:
             value = (loop.time() - started_at) * 1000
-            self.timing(name, value=int(value), tags=tags, sample_rate=sample_rate)
+            if value > threshold_ms:
+                self.timing(name, value=int(value), tags=tags, sample_rate=sample_rate)
 
 
 class DatagramProtocol(asyncio.DatagramProtocol):
diff --git a/examples/timeit.py b/examples/timeit.py
index 225943b..aa39dfb 100644
--- a/examples/timeit.py
+++ b/examples/timeit.py
@@ -9,10 +9,12 @@ async def main():
     )
     await client.connect()
 
+    # Use threshold_ms for setting a threshold for sending the timing metric.
     with client.timeit("fire"):
         # Do action we want to time
         pass
 
+
     await client.close()
 
 
diff --git a/tests/test_client.py b/tests/test_client.py
index f605aaf..28511ab 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -153,6 +153,12 @@ async def test_timeit(self, statsd_client, statsd_server, wait_for, mocker):
         loop.return_value.time.return_value = 1.0
         with statsd_client.timeit("test_timer", tags={"and": "robin"}):
             loop.return_value.time.return_value = 2.0
+
+        # This shouldn't be logged.
+        loop.return_value.time.return_value = 1.0
+        with statsd_client.timeit("test_timer", tags={"and": "robin"}, threshold_ms=3000.0):
+            loop.return_value.time.return_value = 2.0
+
         async with udp_server:
             await wait_for(collected)
         assert collected == [b"test_timer:1000|ms|#whoami:batman,and:robin"]