Add unit test for bucket limit appender

Refactors textparser test to use a common test utility to create
protobuf representation from MetricFamily

Signed-off-by: György Krajcsovits <gyorgy.krajcsovits@grafana.com>

model/textparse/protobufparse_test.go

@@ -15,7 +15,6 @@ package textparse
 
 import (
 	"bytes"
-	"encoding/binary"
 	"errors"
 	"io"
 	"testing"
@@ -26,8 +25,9 @@ import (
 	"github.com/prometheus/prometheus/model/exemplar"
 	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/util/testutil"
 
-	dto "github.com/prometheus/prometheus/prompb/io/prometheus/client"
+	dto "github.com/prometheus/client_model/go"
 )
 
 func TestProtobufParse(t *testing.T) {
@@ -449,7 +449,6 @@ metric: <
 `,
 	}
 
-	varintBuf := make([]byte, binary.MaxVarintLen32)
 	inputBuf := &bytes.Buffer{}
 
 	for _, tmf := range textMetricFamilies {
@@ -457,13 +456,8 @@ metric: <
 		// From text to proto message.
 		require.NoError(t, proto.UnmarshalText(tmf, pb))
 		// From proto message to binary protobuf.
-		protoBuf, err := proto.Marshal(pb)
+		err := testutil.AddMetricFamilyToProtobuf(inputBuf, pb)
 		require.NoError(t, err)
-
-		// Write first length, then binary protobuf.
-		varintLength := binary.PutUvarint(varintBuf, uint64(len(protoBuf)))
-		inputBuf.Write(varintBuf[:varintLength])
-		inputBuf.Write(protoBuf)
 	}
 
 	exp := []struct {

scrape/scrape_test.go

@@ -30,6 +30,7 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/pkg/errors"
+	"github.com/prometheus/client_golang/prometheus"
 	dto "github.com/prometheus/client_model/go"
 	config_util "github.com/prometheus/common/config"
 	"github.com/prometheus/common/model"
@@ -1709,6 +1710,95 @@ func TestScrapeLoopAppendSampleLimit(t *testing.T) {
 	require.Equal(t, 0, seriesAdded)
 }
 
+func TestScrapeLoop_HistogramBucketLimit(t *testing.T) {
+	resApp := &collectResultAppender{}
+	app := &bucketLimitAppender{Appender: resApp, limit: 2}
+
+	sl := newScrapeLoop(context.Background(),
+		nil, nil, nil,
+		func(l labels.Labels) labels.Labels {
+			if l.Has("deleteme") {
+				return labels.EmptyLabels()
+			}
+			return l
+		},
+		nopMutator,
+		func(ctx context.Context) storage.Appender { return app },
+		nil,
+		0,
+		true,
+		app.limit, 0,
+		nil,
+		0,
+		0,
+		false,
+		false,
+		nil,
+		false,
+	)
+
+	metric := dto.Metric{}
+	err := targetScrapeNativeHistogramBucketLimit.Write(&metric)
+	require.NoError(t, err)
+	beforeMetricValue := metric.GetCounter().GetValue()
+
+	nativeHistogram := prometheus.NewHistogram(prometheus.HistogramOpts{
+		Namespace:                      "testing",
+		Name:                           "example_native_histogram",
+		Help:                           "This is used for testing",
+		ConstLabels:                    map[string]string{"some": "value"},
+		NativeHistogramBucketFactor:    1.1, // 10% increase from bucket to bucket
+		NativeHistogramMaxBucketNumber: 100, // intentionally higher than the limit we'll use in the scraper
+	})
+	registry := prometheus.NewRegistry()
+	registry.Register(nativeHistogram)
+	nativeHistogram.Observe(1.0)
+	nativeHistogram.Observe(10.0) // in different bucket since > 1*1.1
+
+	gathered, err := registry.Gather()
+	require.NoError(t, err)
+	require.NotEmpty(t, gathered)
+
+	histogramMetricFamily := gathered[0]
+	msg, err := testutil.MetricFamilyToProtobuf(histogramMetricFamily)
+	require.NoError(t, err)
+
+	now := time.Now()
+	total, added, seriesAdded, err := sl.append(app, msg, "application/vnd.google.protobuf", now)
+	require.NoError(t, err)
+	require.Equal(t, 1, total)
+	require.Equal(t, 1, added)
+	require.Equal(t, 1, seriesAdded)
+
+	err = targetScrapeNativeHistogramBucketLimit.Write(&metric)
+	require.NoError(t, err)
+	metricValue := metric.GetCounter().GetValue()
+	require.Equal(t, beforeMetricValue, metricValue)
+	beforeMetricValue = metricValue
+
+	nativeHistogram.Observe(100.0) // in different bucket since > 10*1.1
+
+	gathered, err = registry.Gather()
+	require.NoError(t, err)
+	require.NotEmpty(t, gathered)
+
+	histogramMetricFamily = gathered[0]
+	msg, err = testutil.MetricFamilyToProtobuf(histogramMetricFamily)
+	require.NoError(t, err)
+
+	now = time.Now()
+	total, added, seriesAdded, err = sl.append(app, msg, "application/vnd.google.protobuf", now)
+	require.NoError(t, err)
+	require.Equal(t, 1, total)
+	require.Equal(t, 1, added)
+	require.Equal(t, 0, seriesAdded)
+
+	err = targetScrapeNativeHistogramBucketLimit.Write(&metric)
+	require.NoError(t, err)
+	metricValue = metric.GetCounter().GetValue()
+	require.Equal(t, beforeMetricValue+1, metricValue)
+}
+
 func TestScrapeLoop_ChangingMetricString(t *testing.T) {
 	// This is a regression test for the scrape loop cache not properly maintaining
 	// IDs when the string representation of a metric changes across a scrape. Thus

util/testutil/protobuf.go

@@ -0,0 +1,50 @@
+// Copyright 2023 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package testutil
+
+import (
+	"bytes"
+	"encoding/binary"
+
+	"github.com/gogo/protobuf/proto"
+	dto "github.com/prometheus/client_model/go"
+)
+
+// Write a MetricFamily into a protobuf
+func MetricFamilyToProtobuf(metricFamily *dto.MetricFamily) ([]byte, error) {
+	buffer := &bytes.Buffer{}
+	err := AddMetricFamilyToProtobuf(buffer, metricFamily)
+	if err != nil {
+		return nil, err
+	}
+	return buffer.Bytes(), nil
+}
+
+// Append a MetricFamily protobuf representation to a buffer
+func AddMetricFamilyToProtobuf(buffer *bytes.Buffer, metricFamily *dto.MetricFamily) error {
+	protoBuf, err := proto.Marshal(metricFamily)
+	if err != nil {
+		return err
+	}
+
+	varintBuf := make([]byte, binary.MaxVarintLen32)
+	varintLength := binary.PutUvarint(varintBuf, uint64(len(protoBuf)))
+
+	_, err = buffer.Write(varintBuf[:varintLength])
+	if err != nil {
+		return err
+	}
+	_, err = buffer.Write(protoBuf)
+	return err
+}