github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/util/server/middleware_test.go (about)

     1  package server
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"net/http"
     7  	"net/http/httptest"
     8  	"net/url"
     9  	"testing"
    10  
    11  	"github.com/stretchr/testify/require"
    12  )
    13  
    14  func TestPrepopulate(t *testing.T) {
    15  	success := http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
    16  		_, err := w.Write([]byte("ok"))
    17  		require.Nil(t, err)
    18  	})
    19  
    20  	for _, tc := range []struct {
    21  		desc     string
    22  		method   string
    23  		qs       string
    24  		body     io.Reader
    25  		expected url.Values
    26  		error    bool
    27  	}{
    28  		{
    29  			desc:   "passthrough GET w/ querystring",
    30  			method: "GET",
    31  			qs:     "?" + url.Values{"foo": {"bar"}}.Encode(),
    32  			body:   nil,
    33  			expected: url.Values{
    34  				"foo": {"bar"},
    35  			},
    36  		},
    37  		{
    38  			desc:   "passthrough POST w/ querystring",
    39  			method: "POST",
    40  			qs:     "?" + url.Values{"foo": {"bar"}}.Encode(),
    41  			body:   nil,
    42  			expected: url.Values{
    43  				"foo": {"bar"},
    44  			},
    45  		},
    46  		{
    47  			desc:   "parse form body",
    48  			method: "POST",
    49  			qs:     "",
    50  			body: bytes.NewBuffer([]byte(url.Values{
    51  				"match": {"up", "down"},
    52  			}.Encode())),
    53  			expected: url.Values{
    54  				"match": {"up", "down"},
    55  			},
    56  		},
    57  		{
    58  			desc:   "querystring extends form body",
    59  			method: "POST",
    60  			qs: "?" + url.Values{
    61  				"match": {"sideways"},
    62  				"foo":   {"bar"},
    63  			}.Encode(),
    64  			body: bytes.NewBuffer([]byte(url.Values{
    65  				"match": {"up", "down"},
    66  			}.Encode())),
    67  			expected: url.Values{
    68  				"match": {"up", "down", "sideways"},
    69  				"foo":   {"bar"},
    70  			},
    71  		},
    72  		{
    73  			desc:   "nil body",
    74  			method: "POST",
    75  			body:   nil,
    76  			error:  true,
    77  		},
    78  	} {
    79  		t.Run(tc.desc, func(t *testing.T) {
    80  			req := httptest.NewRequest(tc.method, "http://testing"+tc.qs, tc.body)
    81  
    82  			// For some reason nil bodies aren't maintained after passed to httptest.NewRequest,
    83  			// but are a failure condition for parsing the form data.
    84  			// Therefore set to nil if we're passing a nil body to force an error.
    85  			if tc.body == nil {
    86  				req.Body = nil
    87  			}
    88  
    89  			if tc.method == "POST" {
    90  				req.Header["Content-Type"] = []string{"application/x-www-form-urlencoded"}
    91  			}
    92  
    93  			w := httptest.NewRecorder()
    94  			mware := NewPrepopulateMiddleware().Wrap(success)
    95  
    96  			mware.ServeHTTP(w, req)
    97  
    98  			if tc.error {
    99  				require.Equal(t, http.StatusBadRequest, w.Result().StatusCode)
   100  			} else {
   101  				require.Equal(t, tc.expected, req.Form)
   102  			}
   103  
   104  		})
   105  	}
   106  }