github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/frontend/transport/handler_test.go (about)

     1  package transport
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"net/http"
     7  	"net/http/httptest"
     8  	"strings"
     9  	"testing"
    10  
    11  	"github.com/go-kit/log"
    12  	"github.com/pkg/errors"
    13  	"github.com/prometheus/client_golang/prometheus"
    14  	promtest "github.com/prometheus/client_golang/prometheus/testutil"
    15  	"github.com/stretchr/testify/assert"
    16  	"github.com/stretchr/testify/require"
    17  	"github.com/weaveworks/common/httpgrpc"
    18  	"github.com/weaveworks/common/user"
    19  )
    20  
    21  type roundTripperFunc func(*http.Request) (*http.Response, error)
    22  
    23  func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) {
    24  	return f(r)
    25  }
    26  
    27  func TestWriteError(t *testing.T) {
    28  	for _, test := range []struct {
    29  		status int
    30  		err    error
    31  	}{
    32  		{http.StatusInternalServerError, errors.New("unknown")},
    33  		{http.StatusGatewayTimeout, context.DeadlineExceeded},
    34  		{StatusClientClosedRequest, context.Canceled},
    35  		{http.StatusBadRequest, httpgrpc.Errorf(http.StatusBadRequest, "")},
    36  	} {
    37  		t.Run(test.err.Error(), func(t *testing.T) {
    38  			w := httptest.NewRecorder()
    39  			writeError(w, test.err)
    40  			require.Equal(t, test.status, w.Result().StatusCode)
    41  		})
    42  	}
    43  }
    44  
    45  func TestHandler_ServeHTTP(t *testing.T) {
    46  	for _, tt := range []struct {
    47  		name            string
    48  		cfg             HandlerConfig
    49  		expectedMetrics int
    50  	}{
    51  		{
    52  			name:            "test handler with stats enabled",
    53  			cfg:             HandlerConfig{QueryStatsEnabled: true},
    54  			expectedMetrics: 3,
    55  		},
    56  		{
    57  			name:            "test handler with stats disabled",
    58  			cfg:             HandlerConfig{QueryStatsEnabled: false},
    59  			expectedMetrics: 0,
    60  		},
    61  	} {
    62  		t.Run(tt.name, func(t *testing.T) {
    63  			roundTripper := roundTripperFunc(func(req *http.Request) (*http.Response, error) {
    64  				return &http.Response{
    65  					StatusCode: http.StatusOK,
    66  					Body:       io.NopCloser(strings.NewReader("{}")),
    67  				}, nil
    68  			})
    69  
    70  			reg := prometheus.NewPedanticRegistry()
    71  			handler := NewHandler(tt.cfg, roundTripper, log.NewNopLogger(), reg)
    72  
    73  			ctx := user.InjectOrgID(context.Background(), "12345")
    74  			req := httptest.NewRequest("GET", "/", nil)
    75  			req = req.WithContext(ctx)
    76  			resp := httptest.NewRecorder()
    77  
    78  			handler.ServeHTTP(resp, req)
    79  			_, _ = io.ReadAll(resp.Body)
    80  			require.Equal(t, resp.Code, http.StatusOK)
    81  
    82  			count, err := promtest.GatherAndCount(
    83  				reg,
    84  				"cortex_query_seconds_total",
    85  				"cortex_query_fetched_series_total",
    86  				"cortex_query_fetched_chunks_bytes_total",
    87  			)
    88  
    89  			assert.NoError(t, err)
    90  			assert.Equal(t, tt.expectedMetrics, count)
    91  		})
    92  	}
    93  }