Fix behavior when adding a custom Gauge or Counter with a custom Regi…

…stry (#39)

* fix: ensure AddCustomGauge and AddCustomCounter register the Gauge and Value in the correct registry

Before that, the AddCustomGauge and AddCustomCounter would both register the respective Gauge and Counter in the default prometheus Registry, even if a custom one was set.
Also, since this commit introduced a new simple mustRegister method on *Prometheus, it seemed logical to use it in the register function.

* feat: add tests to check that custom counters and gauges are registered in the correct registry

prom.go

@@ -153,7 +153,7 @@ func (p *Prometheus) AddCustomGauge(name, help string, labels []string) {
 	},
 		labels)
 	p.customGauges.values[name] = *g
-	prometheus.MustRegister(g)
+	p.mustRegister(g)
 }
 
 // IncrementCounterValue increments a custom counter.
@@ -193,7 +193,12 @@ func (p *Prometheus) AddCustomCounter(name, help string, labels []string) {
 		Help:      help,
 	}, labels)
 	p.customCounters.values[name] = *g
-	prometheus.MustRegister(g)
+	p.mustRegister(g)
+}
+
+func (p *Prometheus) mustRegister(c ...prometheus.Collector) {
+	registerer, _ := p.getRegistererAndGatherer()
+	registerer.MustRegister(c...)
 }
 
 // Path is an option allowing to set the metrics path when initializing with New.
@@ -350,7 +355,6 @@ func (p *Prometheus) getRegistererAndGatherer() (prometheus.Registerer, promethe
 }
 
 func (p *Prometheus) register() {
-	registerer, _ := p.getRegistererAndGatherer()
 	p.reqCnt = prometheus.NewCounterVec(
 		prometheus.CounterOpts{
 			Namespace: p.Namespace,
@@ -360,7 +364,7 @@ func (p *Prometheus) register() {
 		},
 		[]string{"code", "method", "handler", "host", "path"},
 	)
-	registerer.MustRegister(p.reqCnt)
+	p.mustRegister(p.reqCnt)
 
 	p.reqDur = prometheus.NewHistogramVec(prometheus.HistogramOpts{
 		Namespace: p.Namespace,
@@ -369,7 +373,7 @@ func (p *Prometheus) register() {
 		Name:      p.RequestDurationMetricName,
 		Help:      "The HTTP request latency bucket.",
 	}, []string{"method", "path", "host"})
-	registerer.MustRegister(p.reqDur)
+	p.mustRegister(p.reqDur)
 
 	p.reqSz = prometheus.NewSummary(
 		prometheus.SummaryOpts{
@@ -379,7 +383,7 @@ func (p *Prometheus) register() {
 			Help:      "The HTTP request sizes in bytes.",
 		},
 	)
-	registerer.MustRegister(p.reqSz)
+	p.mustRegister(p.reqSz)
 
 	p.resSz = prometheus.NewSummary(
 		prometheus.SummaryOpts{
@@ -389,7 +393,7 @@ func (p *Prometheus) register() {
 			Help:      "The HTTP response sizes in bytes.",
 		},
 	)
-	registerer.MustRegister(p.resSz)
+	p.mustRegister(p.resSz)
 }
 
 func (p *Prometheus) isIgnored(path string) bool {

prom_test.go

@@ -553,3 +553,49 @@ func TestMultipleGinWithDifferentRegistry(t *testing.T) {
 	p2 := New(Engine(r2), Registry(prometheus.NewRegistry()))
 	r2.Use(p2.Instrument())
 }
+
+func TestCustomGaugeCorrectRegistry(t *testing.T) {
+	reg := prometheus.NewRegistry()
+	p := New(Registry(reg))
+
+	p.AddCustomGauge("some_gauge", "", nil)
+	// increment the gauge value so it is reported by Gather
+	err := p.IncrementGaugeValue("some_gauge", nil)
+	assert.Nil(t, err)
+
+	fams, err := reg.Gather()
+	assert.Nil(t, err)
+	assert.Len(t, fams, 3)
+
+	assert.Condition(t, func() (success bool) {
+		for _, fam := range fams {
+			if fam.GetName() == fmt.Sprintf("%s_%s_some_gauge", p.Namespace, p.Subsystem) {
+				return true
+			}
+		}
+		return false
+	})
+}
+
+func TestCustomCounterCorrectRegistry(t *testing.T) {
+	reg := prometheus.NewRegistry()
+	p := New(Registry(reg))
+
+	p.AddCustomCounter("some_counter", "", nil)
+	// increment the counter value so it is reported by Gather
+	err := p.IncrementCounterValue("some_counter", nil)
+	assert.Nil(t, err)
+
+	fams, err := reg.Gather()
+	assert.Nil(t, err)
+	assert.Len(t, fams, 3)
+
+	assert.Condition(t, func() (success bool) {
+		for _, fam := range fams {
+			if fam.GetName() == fmt.Sprintf("%s_%s_some_counter", p.Namespace, p.Subsystem) {
+				return true
+			}
+		}
+		return false
+	})
+}