Files
grafana/pkg/infra/httpclient/httpclientprovider/response_limit_middleware_test.go
Dimitris Sotirakis ba9d5540b8 Datasources: Introduce response_limit for datasource responses (#38962)
* Introduce response_limit for datasource responses

* Fix lint

* Fix tests

* Add case where limit <= 0 - added parametrized tests

* Add max_bytes_reader.go

* Use new httpclient.MaxBytesReader instead of net/http one

* Fixes according to reviewer's comments

* Add tests for max_bytes_reader

* Add small piece in configuration.md

* Further fixes according to reviewer's comments

* Fix linting - fix test
2021-09-10 16:51:06 +03:00

61 lines
1.8 KiB
Go

package httpclientprovider
import (
"context"
"errors"
"fmt"
"io/ioutil"
"net/http"
"strings"
"testing"
"github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
"github.com/stretchr/testify/require"
)
func TestResponseLimitMiddleware(t *testing.T) {
tcs := []struct {
limit int64
bodyLength int
body string
err error
}{
{limit: 1, bodyLength: 1, body: "d", err: errors.New("error: http: response body too large, response limit is set to: 1")},
{limit: 1000000, bodyLength: 5, body: "dummy", err: nil},
{limit: 0, bodyLength: 5, body: "dummy", err: nil},
}
for _, tc := range tcs {
t.Run(fmt.Sprintf("Test ResponseLimitMiddleware with limit: %d", tc.limit), func(t *testing.T) {
finalRoundTripper := httpclient.RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
return &http.Response{StatusCode: http.StatusOK, Request: req, Body: ioutil.NopCloser(strings.NewReader("dummy"))}, nil
})
mw := ResponseLimitMiddleware(tc.limit)
rt := mw.CreateMiddleware(httpclient.Options{}, finalRoundTripper)
require.NotNil(t, rt)
middlewareName, ok := mw.(httpclient.MiddlewareName)
require.True(t, ok)
require.Equal(t, ResponseLimitMiddlewareName, middlewareName.MiddlewareName())
ctx := context.Background()
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://test.com/query", nil)
require.NoError(t, err)
res, err := rt.RoundTrip(req)
require.NoError(t, err)
require.NotNil(t, res)
require.NotNil(t, res.Body)
require.NoError(t, res.Body.Close())
bodyBytes, err := ioutil.ReadAll(res.Body)
if err != nil {
require.EqualError(t, tc.err, err.Error())
} else {
require.NoError(t, tc.err)
}
require.Len(t, bodyBytes, tc.bodyLength)
require.Equal(t, string(bodyBytes), tc.body)
})
}
}