Honorning buckets passed as part of PrometheusMetrics initialization …

…in Histogram
rycus86 · Mar 27, 2024 · a461ff4 · a461ff4
1 parent 9e20f3f
commit a461ff4
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 0 deletions.
diff --git a/prometheus_flask_exporter/__init__.py b/prometheus_flask_exporter/__init__.py
@@ -679,6 +679,10 @@ def _track(self, metric_type, metric_call, metric_kwargs, name, description, lab
 
         labels = self._get_combined_labels(labels)
 
+        if metric_type is Histogram:
+            if metric_kwargs.get('buckets') is None and self.buckets is not None:
+                metric_kwargs['buckets'] = self.buckets
+
         parent_metric = metric_type(
             name, description, labelnames=labels.keys(), registry=registry,
             **metric_kwargs

diff --git a/tests/test_defaults.py b/tests/test_defaults.py
@@ -247,6 +247,7 @@ def decorate_metrics(f):
             def decorated(*args):
                 invocations.append('metrics')
                 return f(*args)
+
             return decorated
 
         self.metrics(metrics_decorator=decorate_metrics)
@@ -639,6 +640,42 @@ def test():
             ('le', '5.0'), ('method', 'GET'), ('path', '/test'), ('status', 200)
         )
 
+    def test_custom_buckets_honored_in_histogram(self):
+        metrics = self.metrics(buckets=(0.2, 2, 4))
+
+        @self.app.route('/test')
+        @metrics.histogram('hist_1', 'Histogram 1')
+        def test1():
+            return 'OK'
+
+        self.client.get('/test')
+        self.client.get('/test')
+
+        self.assertMetric('hist_1_bucket', '2.0', ('le', '0.2'))
+
+        self.assertMetric('hist_1_bucket', '2.0', ('le', '2.0'))
+
+        self.assertMetric('hist_1_bucket', '2.0', ('le', '4.0'))
+
+        self.assertMetric('hist_1_bucket', '2.0', ('le', '+Inf'))
+
+        self.assertAbsent(
+            'hist_1_bucket',
+            ('le', '0.1'), ('method', 'GET'), ('path', '/test'), ('status', 200)
+        )
+        self.assertAbsent(
+            'hist_1',
+            ('le', '0.3'), ('method', 'GET'), ('path', '/test'), ('status', 200)
+        )
+        self.assertAbsent(
+            'hist_1',
+            ('le', '1.2'), ('method', 'GET'), ('path', '/test'), ('status', 200)
+        )
+        self.assertAbsent(
+            'hist_1',
+            ('le', '5.0'), ('method', 'GET'), ('path', '/test'), ('status', 200)
+        )
+
     def test_invalid_labels(self):
         metrics = self.metrics()