Function TestHandler_MetricCollection

pkg/util/push/push_test.go:1071–1120

(t *testing.T)

}

func TestHandler_MetricCollection(t *testing.T) {
	var limits validation.Limits
	flagext.DefaultValues(&limits)
	overrides := validation.NewOverrides(limits, nil)

	counter := prometheus.NewCounterVec(prometheus.CounterOpts{
		Name: "test_counter",
		Help: "test help",
	}, []string{"type"})

	handler := Handler(true, false, 100000, overrides, nil, verifyWriteRequestHandler(t, cortexpb.API), counter)

	t.Run("counts v1 requests", func(t *testing.T) {
		req := createRequest(t, createPrometheusRemoteWriteProtobuf(t), false)
		resp := httptest.NewRecorder()
		handler.ServeHTTP(resp, req)
		assert.Equal(t, http.StatusOK, resp.Code)

		val := testutil.ToFloat64(counter.WithLabelValues("prw1"))
		assert.Equal(t, 1.0, val)
	})

	t.Run("counts v2 requests", func(t *testing.T) {
		ctx := context.Background()
		ctx = user.InjectOrgID(ctx, "user-1")

		req := createRequest(t, createPrometheusRemoteWriteV2Protobuf(t), true)
		req = req.WithContext(ctx)
		resp := httptest.NewRecorder()
		handler.ServeHTTP(resp, req)
		assert.Equal(t, http.StatusNoContent, resp.Code)

		val := testutil.ToFloat64(counter.WithLabelValues("prw2"))
		assert.Equal(t, 1.0, val)
	})

	t.Run("counts unknown or invalid content-type as unknown when acceptUnknownRemoteWriteContentType is true", func(t *testing.T) {
		handlerWithUnknown := Handler(true, true, 100000, overrides, nil, verifyWriteRequestHandler(t, cortexpb.API), counter)
		req := createRequestWithHeaders(t, map[string]string{
			"Content-Type":     "yolo",
			"Content-Encoding": "snappy",
		}, createCortexWriteRequestProtobuf(t, false, cortexpb.API))
		resp := httptest.NewRecorder()
		handlerWithUnknown.ServeHTTP(resp, req)
		assert.Equal(t, http.StatusOK, resp.Code)

		val := testutil.ToFloat64(counter.WithLabelValues("unknown"))
		assert.Equal(t, 1.0, val)
	})
}

func verifyWriteRequestHandler(t *testing.T, expectSource cortexpb.SourceEnum) func(ctx context.Context, request *cortexpb.WriteRequest) (response *cortexpb.WriteResponse, err error) {
	t.Helper()