From a461ff4b3d26938b3f2b8cc07f700b501be74efa Mon Sep 17 00:00:00 2001
From: ronakgarg <rdgarg58@gmail.com>
Date: Wed, 27 Mar 2024 17:21:04 +0530
Subject: [PATCH] Honorning buckets passed as part of PrometheusMetrics
 initialization in Histogram

---
 prometheus_flask_exporter/__init__.py |  4 +++
 tests/test_defaults.py                | 37 +++++++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/prometheus_flask_exporter/__init__.py b/prometheus_flask_exporter/__init__.py
index f5ca7d0..a433bb0 100644
--- 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
index 2f633a7..a236999 100644
--- 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()