github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/util/push/push_test.go (about)

     1  package push
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"net/http"
     7  	"net/http/httptest"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/golang/snappy"
    12  	"github.com/prometheus/prometheus/prompb"
    13  	"github.com/stretchr/testify/assert"
    14  	"github.com/stretchr/testify/require"
    15  	"github.com/weaveworks/common/middleware"
    16  
    17  	"github.com/cortexproject/cortex/pkg/cortexpb"
    18  )
    19  
    20  func TestHandler_remoteWrite(t *testing.T) {
    21  	req := createRequest(t, createPrometheusRemoteWriteProtobuf(t))
    22  	resp := httptest.NewRecorder()
    23  	handler := Handler(100000, nil, verifyWriteRequestHandler(t, cortexpb.API))
    24  	handler.ServeHTTP(resp, req)
    25  	assert.Equal(t, 200, resp.Code)
    26  }
    27  
    28  func TestHandler_cortexWriteRequest(t *testing.T) {
    29  	req := createRequest(t, createCortexWriteRequestProtobuf(t, false))
    30  	resp := httptest.NewRecorder()
    31  	sourceIPs, _ := middleware.NewSourceIPs("SomeField", "(.*)")
    32  	handler := Handler(100000, sourceIPs, verifyWriteRequestHandler(t, cortexpb.RULE))
    33  	handler.ServeHTTP(resp, req)
    34  	assert.Equal(t, 200, resp.Code)
    35  }
    36  
    37  func TestHandler_ignoresSkipLabelNameValidationIfSet(t *testing.T) {
    38  	for _, req := range []*http.Request{
    39  		createRequest(t, createCortexWriteRequestProtobuf(t, true)),
    40  		createRequest(t, createCortexWriteRequestProtobuf(t, false)),
    41  	} {
    42  		resp := httptest.NewRecorder()
    43  		handler := Handler(100000, nil, verifyWriteRequestHandler(t, cortexpb.RULE))
    44  		handler.ServeHTTP(resp, req)
    45  		assert.Equal(t, 200, resp.Code)
    46  	}
    47  }
    48  
    49  func verifyWriteRequestHandler(t *testing.T, expectSource cortexpb.WriteRequest_SourceEnum) func(ctx context.Context, request *cortexpb.WriteRequest) (response *cortexpb.WriteResponse, err error) {
    50  	t.Helper()
    51  	return func(ctx context.Context, request *cortexpb.WriteRequest) (response *cortexpb.WriteResponse, err error) {
    52  		assert.Len(t, request.Timeseries, 1)
    53  		assert.Equal(t, "__name__", request.Timeseries[0].Labels[0].Name)
    54  		assert.Equal(t, "foo", request.Timeseries[0].Labels[0].Value)
    55  		assert.Equal(t, expectSource, request.Source)
    56  		assert.False(t, request.SkipLabelNameValidation)
    57  		return &cortexpb.WriteResponse{}, nil
    58  	}
    59  }
    60  
    61  func createRequest(t *testing.T, protobuf []byte) *http.Request {
    62  	t.Helper()
    63  	inoutBytes := snappy.Encode(nil, protobuf)
    64  	req, err := http.NewRequest("POST", "http://localhost/", bytes.NewReader(inoutBytes))
    65  	require.NoError(t, err)
    66  	req.Header.Add("Content-Encoding", "snappy")
    67  	req.Header.Set("Content-Type", "application/x-protobuf")
    68  	req.Header.Set("X-Prometheus-Remote-Write-Version", "0.1.0")
    69  	return req
    70  }
    71  
    72  func createPrometheusRemoteWriteProtobuf(t *testing.T) []byte {
    73  	t.Helper()
    74  	input := prompb.WriteRequest{
    75  		Timeseries: []prompb.TimeSeries{
    76  			{
    77  				Labels: []prompb.Label{
    78  					{Name: "__name__", Value: "foo"},
    79  				},
    80  				Samples: []prompb.Sample{
    81  					{Value: 1, Timestamp: time.Date(2020, 4, 1, 0, 0, 0, 0, time.UTC).UnixNano()},
    82  				},
    83  			},
    84  		},
    85  	}
    86  	inoutBytes, err := input.Marshal()
    87  	require.NoError(t, err)
    88  	return inoutBytes
    89  }
    90  func createCortexWriteRequestProtobuf(t *testing.T, skipLabelNameValidation bool) []byte {
    91  	t.Helper()
    92  	ts := cortexpb.PreallocTimeseries{
    93  		TimeSeries: &cortexpb.TimeSeries{
    94  			Labels: []cortexpb.LabelAdapter{
    95  				{Name: "__name__", Value: "foo"},
    96  			},
    97  			Samples: []cortexpb.Sample{
    98  				{Value: 1, TimestampMs: time.Date(2020, 4, 1, 0, 0, 0, 0, time.UTC).UnixNano()},
    99  			},
   100  		},
   101  	}
   102  	input := cortexpb.WriteRequest{
   103  		Timeseries:              []cortexpb.PreallocTimeseries{ts},
   104  		Source:                  cortexpb.RULE,
   105  		SkipLabelNameValidation: skipLabelNameValidation,
   106  	}
   107  	inoutBytes, err := input.Marshal()
   108  	require.NoError(t, err)
   109  	return inoutBytes
   110  }