package util_test import ( "bytes" "context" "html/template" "io/ioutil" "math/rand" "net/http" "net/http/httptest" "strconv" "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gopkg.in/yaml.v2" "github.com/grafana/loki/pkg/logproto" "github.com/grafana/loki/pkg/util" util_log "github.com/grafana/loki/pkg/util/log" ) func TestRenderHTTPResponse(t *testing.T) { type testStruct struct { Name string `json:"name"` Value int `json:"value"` } tests := []struct { name string headers map[string]string tmpl string expectedOutput string expectedContentType string value testStruct }{ { name: "Test Renders json", headers: map[string]string{ "Accept": "application/json", }, tmpl: "", expectedOutput: `{"name":"testName","value":42}`, expectedContentType: "application/json", value: testStruct{ Name: "testName", Value: 42, }, }, { name: "Test Renders html", headers: map[string]string{}, tmpl: "{{ .Name }}", expectedOutput: "testName", expectedContentType: "text/html; charset=utf-8", value: testStruct{ Name: "testName", Value: 42, }, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { tmpl := template.Must(template.New("webpage").Parse(tt.tmpl)) writer := httptest.NewRecorder() request := httptest.NewRequest("GET", "/", nil) for k, v := range tt.headers { request.Header.Add(k, v) } util.RenderHTTPResponse(writer, tt.value, tmpl, request) assert.Equal(t, tt.expectedContentType, writer.Header().Get("Content-Type")) assert.Equal(t, 200, writer.Code) assert.Equal(t, tt.expectedOutput, writer.Body.String()) }) } } func TestWriteTextResponse(t *testing.T) { w := httptest.NewRecorder() util.WriteTextResponse(w, "hello world") assert.Equal(t, 200, w.Code) assert.Equal(t, "hello world", w.Body.String()) assert.Equal(t, "text/plain", w.Header().Get("Content-Type")) } func TestStreamWriteYAMLResponse(t *testing.T) { type testStruct struct { Name string `yaml:"name"` Value int `yaml:"value"` } tt := struct { name string headers map[string]string expectedOutput string expectedContentType string value map[string]*testStruct }{ name: "Test Stream Render YAML", headers: map[string]string{ "Content-Type": "application/yaml", }, expectedContentType: "application/yaml", value: make(map[string]*testStruct), } // Generate some data to serialize. for i := 0; i < rand.Intn(100)+1; i++ { ts := testStruct{ Name: "testName" + strconv.Itoa(i), Value: i, } tt.value[ts.Name] = &ts } d, err := yaml.Marshal(tt.value) require.NoError(t, err) tt.expectedOutput = string(d) w := httptest.NewRecorder() done := make(chan struct{}) iter := make(chan interface{}) go func() { util.StreamWriteYAMLResponse(w, iter, util_log.Logger) close(done) }() for k, v := range tt.value { iter <- map[string]*testStruct{k: v} } close(iter) <-done assert.Equal(t, tt.expectedContentType, w.Header().Get("Content-Type")) assert.Equal(t, 200, w.Code) assert.YAMLEq(t, tt.expectedOutput, w.Body.String()) } func TestParseProtoReader(t *testing.T) { // 47 bytes compressed and 53 uncompressed req := &logproto.PreallocWriteRequest{ WriteRequest: logproto.WriteRequest{ Timeseries: []logproto.PreallocTimeseries{ { TimeSeries: &logproto.TimeSeries{ Labels: []logproto.LabelAdapter{ {Name: "foo", Value: "bar"}, }, Samples: []logproto.LegacySample{ {Value: 10, TimestampMs: 1}, {Value: 20, TimestampMs: 2}, {Value: 30, TimestampMs: 3}, }, }, }, }, }, } for _, tt := range []struct { name string compression util.CompressionType maxSize int expectErr bool useBytesBuffer bool }{ {"rawSnappy", util.RawSnappy, 53, false, false}, {"noCompression", util.NoCompression, 53, false, false}, {"too big rawSnappy", util.RawSnappy, 10, true, false}, {"too big decoded rawSnappy", util.RawSnappy, 50, true, false}, {"too big noCompression", util.NoCompression, 10, true, false}, {"bytesbuffer rawSnappy", util.RawSnappy, 53, false, true}, {"bytesbuffer noCompression", util.NoCompression, 53, false, true}, {"bytesbuffer too big rawSnappy", util.RawSnappy, 10, true, true}, {"bytesbuffer too big decoded rawSnappy", util.RawSnappy, 50, true, true}, {"bytesbuffer too big noCompression", util.NoCompression, 10, true, true}, } { t.Run(tt.name, func(t *testing.T) { w := httptest.NewRecorder() assert.Nil(t, util.SerializeProtoResponse(w, req, tt.compression)) var fromWire logproto.PreallocWriteRequest reader := w.Result().Body if tt.useBytesBuffer { buf := bytes.Buffer{} _, err := buf.ReadFrom(reader) assert.Nil(t, err) reader = bytesBuffered{Buffer: &buf} } err := util.ParseProtoReader(context.Background(), reader, 0, tt.maxSize, &fromWire, tt.compression) if tt.expectErr { assert.NotNil(t, err) return } assert.Nil(t, err) assert.Equal(t, req, &fromWire) }) } } type bytesBuffered struct { *bytes.Buffer } func (b bytesBuffered) Close() error { return nil } func (b bytesBuffered) BytesBuffer() *bytes.Buffer { return b.Buffer } func TestIsRequestBodyTooLargeRegression(t *testing.T) { _, err := ioutil.ReadAll(http.MaxBytesReader(httptest.NewRecorder(), ioutil.NopCloser(bytes.NewReader([]byte{1, 2, 3, 4})), 1)) assert.True(t, util.IsRequestBodyTooLarge(err)) }