Merge pull request #147 from ross/timed-context-self

_TimedContextManagerDecorator enter -> self
DataDog · Aug 24, 2016 · be01e58 · be01e58
2 parents 393f804 + 188ee82
commit be01e58
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/datadog/dogstatsd/base.py b/datadog/dogstatsd/base.py
@@ -209,6 +209,7 @@ def __enter__(self):
             if not self.metric:
                 raise TypeError("Cannot used timed without a metric!")
             self.start = time()
+            return self
 
         def __exit__(self, type, value, traceback):
             # Report the elapsed time of the context manager.

diff --git a/tests/unit/dogstatsd/test_statsd.py b/tests/unit/dogstatsd/test_statsd.py
@@ -313,7 +313,8 @@ def test_timed_context(self):
         Measure the distribution of a context's run time.
         """
         # In seconds
-        with self.statsd.timed('timed_context.test'):
+        with self.statsd.timed('timed_context.test') as timer:
+            t.assert_is_instance(timer, DogStatsd._TimedContextManagerDecorator)
             time.sleep(0.5)
 
         packet = self.recv()