github.com/thanos-io/thanos@v0.32.5/internal/cortex/querier/queryrange/roundtrip_test.go (about)

     1  // Copyright (c) The Cortex Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package queryrange
     5  
     6  import (
     7  	"context"
     8  	"io/ioutil"
     9  	"net/http"
    10  	"net/http/httptest"
    11  	"net/url"
    12  	"strconv"
    13  	"testing"
    14  	"time"
    15  
    16  	"github.com/go-kit/log"
    17  	"github.com/prometheus/prometheus/promql"
    18  	"github.com/stretchr/testify/require"
    19  	"github.com/weaveworks/common/middleware"
    20  	"github.com/weaveworks/common/user"
    21  )
    22  
    23  func TestRoundTrip(t *testing.T) {
    24  	s := httptest.NewServer(
    25  		middleware.AuthenticateUser.Wrap(
    26  			http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    27  				var err error
    28  				if r.RequestURI == query {
    29  					_, err = w.Write([]byte(responseBody))
    30  				} else {
    31  					_, err = w.Write([]byte("bar"))
    32  				}
    33  				if err != nil {
    34  					t.Fatal(err)
    35  				}
    36  			}),
    37  		),
    38  	)
    39  	defer s.Close()
    40  
    41  	u, err := url.Parse(s.URL)
    42  	require.NoError(t, err)
    43  
    44  	downstream := singleHostRoundTripper{
    45  		host: u.Host,
    46  		next: http.DefaultTransport,
    47  	}
    48  
    49  	tw, _, err := NewTripperware(Config{},
    50  		log.NewNopLogger(),
    51  		mockLimits{},
    52  		PrometheusCodec,
    53  		nil,
    54  		promql.EngineOpts{
    55  			Logger:     log.NewNopLogger(),
    56  			Reg:        nil,
    57  			MaxSamples: 1000,
    58  			Timeout:    time.Minute,
    59  		},
    60  		0,
    61  		nil,
    62  		nil,
    63  	)
    64  
    65  	if err != nil {
    66  		t.Fatal(err)
    67  	}
    68  
    69  	for i, tc := range []struct {
    70  		path, expectedBody string
    71  	}{
    72  		{"/foo", "bar"},
    73  		{query, responseBody},
    74  	} {
    75  		t.Run(strconv.Itoa(i), func(t *testing.T) {
    76  			req, err := http.NewRequest("GET", tc.path, http.NoBody)
    77  			require.NoError(t, err)
    78  
    79  			// query-frontend doesn't actually authenticate requests, we rely on
    80  			// the queriers to do this.  Hence we ensure the request doesn't have a
    81  			// org ID in the ctx, but does have the header.
    82  			ctx := user.InjectOrgID(context.Background(), "1")
    83  			req = req.WithContext(ctx)
    84  			err = user.InjectOrgIDIntoHTTPRequest(ctx, req)
    85  			require.NoError(t, err)
    86  
    87  			resp, err := tw(downstream).RoundTrip(req)
    88  			require.NoError(t, err)
    89  			require.Equal(t, 200, resp.StatusCode)
    90  
    91  			bs, err := ioutil.ReadAll(resp.Body)
    92  			require.NoError(t, err)
    93  			require.Equal(t, tc.expectedBody, string(bs))
    94  		})
    95  	}
    96  }
    97  
    98  type singleHostRoundTripper struct {
    99  	host string
   100  	next http.RoundTripper
   101  }
   102  
   103  func (s singleHostRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
   104  	r.URL.Scheme = "http"
   105  	r.URL.Host = s.host
   106  	return s.next.RoundTrip(r)
   107  }