diff --git a/datadog/dogstatsd/base.py b/datadog/dogstatsd/base.py index 1a2fa7772..25d509c49 100644 --- 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 index fdba1d838..13992ca65 100644 --- 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()