diff --git a/prometheus_client/context_managers.py b/prometheus_client/context_managers.py index bb2ae847..eba60bcb 100644 --- a/prometheus_client/context_managers.py +++ b/prometheus_client/context_managers.py @@ -55,6 +55,8 @@ def __enter__(self): def __exit__(self, typ, value, traceback): # Time can go backwards. duration = max(default_timer() - self._start, 0) + instance = self._callback.__self__ + instance._raise_if_not_observable() self._callback(duration) def labels(self, *args, **kw): diff --git a/prometheus_client/metrics.py b/prometheus_client/metrics.py index ea797e9d..252b859e 100644 --- a/prometheus_client/metrics.py +++ b/prometheus_client/metrics.py @@ -396,7 +396,6 @@ def time(self): Can be used as a function decorator or context manager. """ - self._raise_if_not_observable() return Timer(self.set) def set_function(self, f): @@ -475,7 +474,6 @@ def time(self): Can be used as a function decorator or context manager. """ - self._raise_if_not_observable() return Timer(self.observe) def _child_samples(self): diff --git a/tests/test_core.py b/tests/test_core.py index 63b536b6..38bc20c4 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -222,13 +222,25 @@ def test_time_block_decorator(self): time.sleep(.001) self.assertNotEqual(0, self.registry.get_sample_value('g')) + def test_time_block_decorator_with_label(self): + value = self.registry.get_sample_value + self.assertEqual(None, value('g2', {'label1': 'foo'})) + with self.gauge_with_label.time() as metric: + metric.labels('foo') + self.assertLess(0, value('g2', {'label1': 'foo'})) + def test_track_in_progress_not_observable(self): g = Gauge('test', 'help', labelnames=('label',), registry=self.registry) assert_not_observable(g.track_inprogress) def test_timer_not_observable(self): g = Gauge('test', 'help', labelnames=('label',), registry=self.registry) - assert_not_observable(g.time) + + def manager(): + with g.time(): + pass + + assert_not_observable(manager) class TestSummary(unittest.TestCase): @@ -318,10 +330,21 @@ def test_block_decorator(self): pass self.assertEqual(1, self.registry.get_sample_value('s_count')) + def test_block_decorator_with_label(self): + value = self.registry.get_sample_value + self.assertEqual(None, value('s_with_labels_count', {'label1': 'foo'})) + with self.summary_with_labels.time() as metric: + metric.labels('foo') + self.assertEqual(1, value('s_with_labels_count', {'label1': 'foo'})) + def test_timer_not_observable(self): s = Summary('test', 'help', labelnames=('label',), registry=self.registry) - assert_not_observable(s.time) + def manager(): + with s.time(): + pass + + assert_not_observable(manager) class TestHistogram(unittest.TestCase):