WIP: Add withLimit functionality

body_wrapper.go

@@ -3,6 +3,7 @@ package httpexpect
 import (
 	"bytes"
 	"context"
+	"errors"
 	"io"
 	"io/ioutil"
 	"runtime"
@@ -41,6 +42,25 @@ func newBodyWrapper(reader io.ReadCloser, cancelFunc context.CancelFunc) *bodyWr
 	return bw
 }
 
+func newBodyWrapperLimit(reader io.ReadCloser, cancelFunc context.CancelFunc,
+	limit int64) (*bodyWrapper, error) {
+
+	b := make([]byte, limit)
+	_, err := io.ReadFull(reader, b)
+	if err != nil {
+		bw := &bodyWrapper{
+			origReader: reader,
+			cancelFunc: cancelFunc,
+		}
+
+		// Finalizer will close body if closeAndCancel was never called.
+		runtime.SetFinalizer(bw, (*bodyWrapper).Close)
+
+		return bw, nil
+	}
+	return nil, errors.New("unexpected: response body size larger than limit")
+}
+
 // Read body contents
 func (bw *bodyWrapper) Read(p []byte) (n int, err error) {
 	bw.mu.Lock()

request.go

@@ -53,6 +53,7 @@ type Request struct {
 	formbuf   *bytes.Buffer
 	multipart *multipart.Writer
 
+	limit        int64
 	bodySetter   string
 	typeSetter   string
 	forceType    bool
@@ -235,6 +236,21 @@ func (r *Request) WithName(name string) *Request {
 	return r
 }
 
+func (r *Request) WithLimit(limit int64) *Request {
+	opChain := r.chain.enter("WithLimit()")
+	defer opChain.leave()
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	if !r.checkOrder(opChain, "WithLimit()") {
+		return r
+	}
+
+	r.limit = limit
+	return r
+}
+
 // WithMatcher attaches a matcher to the request.
 // All attached matchers are invoked in the Expect method for a newly
 // created Response.
@@ -2203,7 +2219,15 @@ func (r *Request) retryRequest(reqFunc func() (*http.Response, error)) (
 		elapsed := time.Since(start)
 
 		if resp != nil && resp.Body != nil {
-			resp.Body = newBodyWrapper(resp.Body, cancelFn)
+			if r.limit != 0 {
+				resp.Body, err = newBodyWrapperLimit(resp.Body, cancelFn, r.limit)
+				if err != nil {
+					return nil, 0, err
+				}
+			} else {
+				resp.Body = newBodyWrapper(resp.Body, cancelFn)
+			}
+
 		} else if cancelFn != nil {
 			cancelFn()
 		}

request_test.go

@@ -265,6 +265,46 @@ func TestRequest_Basic(t *testing.T) {
 
 		assert.Nil(t, resp.Raw())
 	})
+
+}
+
+func TestResponse_Limit(t *testing.T) {
+	t.Run("response size exceeds limit", func(t *testing.T) {
+		client := &mockClient{}
+
+		config := Config{
+			Client:   client,
+			Reporter: newMockReporter(t),
+		}
+
+		req := NewRequestC(config, "METHOD", "/")
+		req.WithText("test string")
+		req.WithLimit(3)
+
+		resp := req.Expect()
+
+		req.chain.assertFailed(t)
+		resp.chain.assertFailed(t)
+	})
+
+	t.Run("response size doesn't exceed limit", func(t *testing.T) {
+		client := &mockClient{}
+
+		config := Config{
+			Client:   client,
+			Reporter: newMockReporter(t),
+		}
+
+		req := NewRequestC(config, "METHOD", "/")
+		req.WithText("1234")
+		req.WithLimit(5)
+
+		resp := req.Expect()
+
+		req.chain.assertNotFailed(t)
+		resp.chain.assertNotFailed(t)
+	})
+
 }
 
 func TestRequest_Matchers(t *testing.T) {