From 75bb6217849be284106d05292e6cf34235ea1664 Mon Sep 17 00:00:00 2001
From: Dave Shawley <daves@aweber.com>
Date: Sun, 18 Jul 2021 10:10:41 -0400
Subject: [PATCH] Add Connector.timer.

This improves compatibility with python-statsd.
---
 CHANGELOG.rst              |  4 ++++
 docs/conf.py               |  1 +
 sprockets_statsd/statsd.py | 16 ++++++++++++++++
 tests/test_processor.py    | 32 +++++++++++++++++++++++++++++++-
 4 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 11a3e27..aabb5c1 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -1,3 +1,7 @@
+Next release
+------------
+- Added ``Connector.timer`` method (addresses :issue:`8`)
+
 :tag:`0.1.0 <0.0.1...0.1.0>` (10-May-2021)
 ------------------------------------------
 - Added :envvar:`STATSD_ENABLED` environment variable to disable the Tornado integration
diff --git a/docs/conf.py b/docs/conf.py
index 7f1cdb5..b0e18c7 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -19,6 +19,7 @@ intersphinx_mapping = {
 # https://www.sphinx-doc.org/en/master/usage/extensions/extlinks.html
 extensions.append('sphinx.ext.extlinks')
 extlinks = {
+    'issue': ("https://github.com/sprockets/sprockets-statsd/issues/%s", "#%s"),
     'tag': ("https://github.com/sprockets/sprockets-statsd/compare/%s", "%s"),
 }
 
diff --git a/sprockets_statsd/statsd.py b/sprockets_statsd/statsd.py
index 513803e..d8080e7 100644
--- a/sprockets_statsd/statsd.py
+++ b/sprockets_statsd/statsd.py
@@ -1,6 +1,8 @@
 import asyncio
+import contextlib
 import logging
 import socket
+import time
 import typing
 
 
@@ -108,6 +110,20 @@ class AbstractConnector:
         """
         self.inject_metric(f'timers.{path}', str(seconds * 1000.0), 'ms')
 
+    @contextlib.contextmanager
+    def timer(self, path):
+        """Send a timer metric using a context manager.
+
+        :param path: timer to append the measured time to
+
+        """
+        start = time.time()
+        try:
+            yield
+        finally:
+            fini = time.time()
+            self.timing(path, max(fini, start) - start)
+
 
 class Connector(AbstractConnector):
     """Sends metrics to a statsd server.
diff --git a/tests/test_processor.py b/tests/test_processor.py
index ab69c5e..e2f98d3 100644
--- a/tests/test_processor.py
+++ b/tests/test_processor.py
@@ -286,7 +286,13 @@ class ConnectorTests(ProcessorTestCase):
         recvd_path, _, rest = decoded.partition(':')
         recvd_value, _, recvd_code = rest.partition('|')
         self.assertEqual(path, recvd_path, 'metric path mismatch')
-        self.assertEqual(recvd_value, str(value), 'metric value mismatch')
+        if type_code == 'ms':
+            self.assertAlmostEqual(float(recvd_value),
+                                   value,
+                                   places=3,
+                                   msg='metric value mismatch')
+        else:
+            self.assertEqual(recvd_value, str(value), 'metric value mismatch')
         self.assertEqual(recvd_code, type_code, 'metric type mismatch')
 
     async def test_adjusting_counter(self):
@@ -331,6 +337,30 @@ class ConnectorTests(ProcessorTestCase):
         self.assert_metrics_equal(self.statsd_server.metrics[0],
                                   'timers.simple.timer', 12340.0, 'ms')
 
+    async def test_timing_context_manager(self):
+        with unittest.mock.patch(
+                'sprockets_statsd.statsd.time.time') as time_function:
+            time_function.side_effect = [10.0, 22.345]
+            with self.connector.timer('some.timer'):
+                pass
+            self.assertEqual(2, time_function.call_count)
+
+        await self.wait_for(self.statsd_server.message_received.acquire())
+        self.assert_metrics_equal(self.statsd_server.metrics[0],
+                                  'timers.some.timer', 12345.0, 'ms')
+
+    async def test_timer_is_monotonic(self):
+        with unittest.mock.patch(
+                'sprockets_statsd.statsd.time.time') as time_function:
+            time_function.side_effect = [10.001, 10.000]
+            with self.connector.timer('some.timer'):
+                pass
+            self.assertEqual(2, time_function.call_count)
+
+        await self.wait_for(self.statsd_server.message_received.acquire())
+        self.assert_metrics_equal(self.statsd_server.metrics[0],
+                                  'timers.some.timer', 0.0, 'ms')
+
     async def test_that_queued_metrics_are_drained(self):
         # The easiest way to test that the internal metrics queue
         # is drained when the processor is stopped is to monkey