From 7a88bc3581f40827bdf848897c7241699f2c8089 Mon Sep 17 00:00:00 2001
From: Ganesh Vernekar <ganeshvern@gmail.com>
Date: Mon, 9 Jan 2023 18:35:04 +0530
Subject: [PATCH] Test federation with native histograms

Signed-off-by: Ganesh Vernekar <ganeshvern@gmail.com>
---
 web/federate_test.go | 115 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 115 insertions(+)

diff --git a/web/federate_test.go b/web/federate_test.go
index f39daedf1..b0def6f79 100644
--- a/web/federate_test.go
+++ b/web/federate_test.go
@@ -16,6 +16,8 @@ package web
 import (
 	"bytes"
 	"context"
+	"fmt"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"sort"
@@ -28,7 +30,9 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/prometheus/prometheus/config"
+	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/model/textparse"
 	"github.com/prometheus/prometheus/promql"
 	"github.com/prometheus/prometheus/storage"
 	"github.com/prometheus/prometheus/tsdb"
@@ -299,3 +303,114 @@ func normalizeBody(body *bytes.Buffer) string {
 	}
 	return strings.Join(lines, "")
 }
+
+func TestFederationWithNativeHistograms(t *testing.T) {
+	suite, err := promql.NewTest(t, "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer suite.Close()
+
+	if err := suite.Run(); err != nil {
+		t.Fatal(err)
+	}
+
+	var expVec promql.Vector
+
+	db := suite.TSDB()
+	hist := &histogram.Histogram{
+		Count:         10,
+		ZeroCount:     2,
+		ZeroThreshold: 0.001,
+		Sum:           39.4,
+		Schema:        1,
+		PositiveSpans: []histogram.Span{
+			{Offset: 0, Length: 2},
+			{Offset: 1, Length: 2},
+		},
+		PositiveBuckets: []int64{1, 1, -1, 0},
+		NegativeSpans: []histogram.Span{
+			{Offset: 0, Length: 2},
+			{Offset: 1, Length: 2},
+		},
+		NegativeBuckets: []int64{1, 1, -1, 0},
+	}
+	app := db.Appender(context.Background())
+	for i := 0; i < 6; i++ {
+		l := labels.FromStrings("__name__", "test_metric", "foo", fmt.Sprintf("%d", i))
+		expL := labels.FromStrings("__name__", "test_metric", "instance", "", "foo", fmt.Sprintf("%d", i))
+		if i%3 == 0 {
+			_, err = app.Append(0, l, 101*60*1000, float64(i*100))
+			expVec = append(expVec, promql.Sample{
+				Point:  promql.Point{T: 101 * 60 * 1000, V: float64(i * 100)},
+				Metric: expL,
+			})
+		} else {
+			hist.ZeroCount++
+			_, err = app.AppendHistogram(0, l, 101*60*1000, hist.Copy(), nil)
+			expVec = append(expVec, promql.Sample{
+				Point:  promql.Point{T: 101 * 60 * 1000, H: hist.ToFloat()},
+				Metric: expL,
+			})
+		}
+		require.NoError(t, err)
+	}
+	require.NoError(t, app.Commit())
+
+	h := &Handler{
+		localStorage:  &dbAdapter{suite.TSDB()},
+		lookbackDelta: 5 * time.Minute,
+		now:           func() model.Time { return 101 * 60 * 1000 }, // 101min after epoch.
+		config: &config.Config{
+			GlobalConfig: config.GlobalConfig{},
+		},
+	}
+
+	req := httptest.NewRequest("GET", "http://example.org/federate?match[]=test_metric", nil)
+	req.Header.Add("Accept", `application/vnd.google.protobuf;proto=io.prometheus.client.MetricFamily;encoding=delimited,application/openmetrics-text;version=1.0.0;q=0.8,application/openmetrics-text;version=0.0.1;q=0.75,text/plain;version=0.0.4;q=0.5,*/*;q=0.1`)
+	res := httptest.NewRecorder()
+
+	h.federation(res, req)
+
+	require.Equal(t, http.StatusOK, res.Code)
+	body, err := io.ReadAll(res.Body)
+	require.NoError(t, err)
+
+	p := textparse.NewProtobufParser(body)
+	var actVec promql.Vector
+	metricFamilies := 0
+	for {
+		et, err := p.Next()
+		if err == io.EOF {
+			break
+		}
+		require.NoError(t, err)
+		if et == textparse.EntryHelp {
+			metricFamilies++
+		}
+		if et == textparse.EntryHistogram || et == textparse.EntrySeries {
+			l := labels.Labels{}
+			p.Metric(&l)
+			actVec = append(actVec, promql.Sample{Metric: l})
+		}
+		if et == textparse.EntryHistogram {
+			_, parsedTimestamp, h, fh := p.Histogram()
+			require.Nil(t, h)
+			actVec[len(actVec)-1].Point = promql.Point{
+				T: *parsedTimestamp,
+				H: fh,
+			}
+		} else if et == textparse.EntrySeries {
+			_, parsedTimestamp, v := p.Series()
+			actVec[len(actVec)-1].Point = promql.Point{
+				T: *parsedTimestamp,
+				V: v,
+			}
+		}
+	}
+
+	// TODO(codesome): Once PromQL is able to set the CounterResetHint on histograms,
+	// test it with switching histogram types for metric families.
+	require.Equal(t, 4, metricFamilies)
+	require.Equal(t, expVec, actVec)
+}