github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/util/http_test.go (about) 1 package util_test 2 3 import ( 4 "bytes" 5 "context" 6 "html/template" 7 "io/ioutil" 8 "math/rand" 9 "net/http" 10 "net/http/httptest" 11 "strconv" 12 "testing" 13 14 "github.com/stretchr/testify/assert" 15 "github.com/stretchr/testify/require" 16 "gopkg.in/yaml.v2" 17 18 "github.com/grafana/loki/pkg/logproto" 19 "github.com/grafana/loki/pkg/util" 20 util_log "github.com/grafana/loki/pkg/util/log" 21 ) 22 23 func TestRenderHTTPResponse(t *testing.T) { 24 type testStruct struct { 25 Name string `json:"name"` 26 Value int `json:"value"` 27 } 28 29 tests := []struct { 30 name string 31 headers map[string]string 32 tmpl string 33 expectedOutput string 34 expectedContentType string 35 value testStruct 36 }{ 37 { 38 name: "Test Renders json", 39 headers: map[string]string{ 40 "Accept": "application/json", 41 }, 42 tmpl: "<html></html>", 43 expectedOutput: `{"name":"testName","value":42}`, 44 expectedContentType: "application/json", 45 value: testStruct{ 46 Name: "testName", 47 Value: 42, 48 }, 49 }, 50 { 51 name: "Test Renders html", 52 headers: map[string]string{}, 53 tmpl: "<html>{{ .Name }}</html>", 54 expectedOutput: "<html>testName</html>", 55 expectedContentType: "text/html; charset=utf-8", 56 value: testStruct{ 57 Name: "testName", 58 Value: 42, 59 }, 60 }, 61 } 62 63 for _, tt := range tests { 64 t.Run(tt.name, func(t *testing.T) { 65 tmpl := template.Must(template.New("webpage").Parse(tt.tmpl)) 66 writer := httptest.NewRecorder() 67 request := httptest.NewRequest("GET", "/", nil) 68 69 for k, v := range tt.headers { 70 request.Header.Add(k, v) 71 } 72 73 util.RenderHTTPResponse(writer, tt.value, tmpl, request) 74 75 assert.Equal(t, tt.expectedContentType, writer.Header().Get("Content-Type")) 76 assert.Equal(t, 200, writer.Code) 77 assert.Equal(t, tt.expectedOutput, writer.Body.String()) 78 }) 79 } 80 } 81 82 func TestWriteTextResponse(t *testing.T) { 83 w := httptest.NewRecorder() 84 85 util.WriteTextResponse(w, "hello world") 86 87 assert.Equal(t, 200, w.Code) 88 assert.Equal(t, "hello world", w.Body.String()) 89 assert.Equal(t, "text/plain", w.Header().Get("Content-Type")) 90 } 91 92 func TestStreamWriteYAMLResponse(t *testing.T) { 93 type testStruct struct { 94 Name string `yaml:"name"` 95 Value int `yaml:"value"` 96 } 97 tt := struct { 98 name string 99 headers map[string]string 100 expectedOutput string 101 expectedContentType string 102 value map[string]*testStruct 103 }{ 104 name: "Test Stream Render YAML", 105 headers: map[string]string{ 106 "Content-Type": "application/yaml", 107 }, 108 expectedContentType: "application/yaml", 109 value: make(map[string]*testStruct), 110 } 111 112 // Generate some data to serialize. 113 for i := 0; i < rand.Intn(100)+1; i++ { 114 ts := testStruct{ 115 Name: "testName" + strconv.Itoa(i), 116 Value: i, 117 } 118 tt.value[ts.Name] = &ts 119 } 120 d, err := yaml.Marshal(tt.value) 121 require.NoError(t, err) 122 tt.expectedOutput = string(d) 123 w := httptest.NewRecorder() 124 125 done := make(chan struct{}) 126 iter := make(chan interface{}) 127 go func() { 128 util.StreamWriteYAMLResponse(w, iter, util_log.Logger) 129 close(done) 130 }() 131 for k, v := range tt.value { 132 iter <- map[string]*testStruct{k: v} 133 } 134 close(iter) 135 <-done 136 assert.Equal(t, tt.expectedContentType, w.Header().Get("Content-Type")) 137 assert.Equal(t, 200, w.Code) 138 assert.YAMLEq(t, tt.expectedOutput, w.Body.String()) 139 } 140 141 func TestParseProtoReader(t *testing.T) { 142 // 47 bytes compressed and 53 uncompressed 143 req := &logproto.PreallocWriteRequest{ 144 WriteRequest: logproto.WriteRequest{ 145 Timeseries: []logproto.PreallocTimeseries{ 146 { 147 TimeSeries: &logproto.TimeSeries{ 148 Labels: []logproto.LabelAdapter{ 149 {Name: "foo", Value: "bar"}, 150 }, 151 Samples: []logproto.LegacySample{ 152 {Value: 10, TimestampMs: 1}, 153 {Value: 20, TimestampMs: 2}, 154 {Value: 30, TimestampMs: 3}, 155 }, 156 }, 157 }, 158 }, 159 }, 160 } 161 162 for _, tt := range []struct { 163 name string 164 compression util.CompressionType 165 maxSize int 166 expectErr bool 167 useBytesBuffer bool 168 }{ 169 {"rawSnappy", util.RawSnappy, 53, false, false}, 170 {"noCompression", util.NoCompression, 53, false, false}, 171 {"too big rawSnappy", util.RawSnappy, 10, true, false}, 172 {"too big decoded rawSnappy", util.RawSnappy, 50, true, false}, 173 {"too big noCompression", util.NoCompression, 10, true, false}, 174 175 {"bytesbuffer rawSnappy", util.RawSnappy, 53, false, true}, 176 {"bytesbuffer noCompression", util.NoCompression, 53, false, true}, 177 {"bytesbuffer too big rawSnappy", util.RawSnappy, 10, true, true}, 178 {"bytesbuffer too big decoded rawSnappy", util.RawSnappy, 50, true, true}, 179 {"bytesbuffer too big noCompression", util.NoCompression, 10, true, true}, 180 } { 181 t.Run(tt.name, func(t *testing.T) { 182 w := httptest.NewRecorder() 183 assert.Nil(t, util.SerializeProtoResponse(w, req, tt.compression)) 184 var fromWire logproto.PreallocWriteRequest 185 186 reader := w.Result().Body 187 if tt.useBytesBuffer { 188 buf := bytes.Buffer{} 189 _, err := buf.ReadFrom(reader) 190 assert.Nil(t, err) 191 reader = bytesBuffered{Buffer: &buf} 192 } 193 194 err := util.ParseProtoReader(context.Background(), reader, 0, tt.maxSize, &fromWire, tt.compression) 195 if tt.expectErr { 196 assert.NotNil(t, err) 197 return 198 } 199 assert.Nil(t, err) 200 assert.Equal(t, req, &fromWire) 201 }) 202 } 203 } 204 205 type bytesBuffered struct { 206 *bytes.Buffer 207 } 208 209 func (b bytesBuffered) Close() error { 210 return nil 211 } 212 213 func (b bytesBuffered) BytesBuffer() *bytes.Buffer { 214 return b.Buffer 215 } 216 217 func TestIsRequestBodyTooLargeRegression(t *testing.T) { 218 _, err := ioutil.ReadAll(http.MaxBytesReader(httptest.NewRecorder(), ioutil.NopCloser(bytes.NewReader([]byte{1, 2, 3, 4})), 1)) 219 assert.True(t, util.IsRequestBodyTooLarge(err)) 220 }