Repository: gavv/httpexpect
Pull request #97 (2021-03-14 14:44:30)

Add WithContext and WithTimeout #92
Overall I've followed the in-issue guide. Still need to look into a suspicious defer inside retryRequest but this will suffice for review. 

gtourkas: #97 
gtourkas: @gavv Pushed all the changes you requested. 
gavv: Thanks for update! A few final minor comments.
gavv: > Still need to look into a suspicious defer inside retryRequest

Looks legit
gavv: It seems there is a possible deadlock in the new tests.

I've run this:

```
go test -race -count=100 -v |& pp
```

(`pp` is https://github.com/maruel/panicparse)

And eventually `TestPerRequestWithTimeoutAndWithRetries` stucked and then terminated by timeout: https://gist.github.com/gavv/e9aa492ed0ca50e8073251099e9f9f7c
gavv: (Also it seems that goroutines are leaking in TestE2EWebsocketTimeouts, but I think this is unrelated; those goroutines are excluded from trace)
gtourkas: > It seems there is a possible deadlock in the new tests.
> 
> I've run this:
> 
> ```
> go test -race -count=100 -v |& pp
> ```
> 
> (`pp` is https://github.com/maruel/panicparse)
> 
> And eventually `TestPerRequestWithTimeoutAndWithRetries` stucked and then terminated by timeout: https://gist.github.com/gavv/e9aa492ed0ca50e8073251099e9f9f7c

I'll look into this shortly.
gtourkas: > > It seems there is a possible deadlock in the new tests.
> > I've run this:
> > ```
> > go test -race -count=100 -v |& pp
> > ```
> > 
> > 
> > (`pp` is https://github.com/maruel/panicparse)
> > And eventually `TestPerRequestWithTimeoutAndWithRetries` stucked and then terminated by timeout: https://gist.github.com/gavv/e9aa492ed0ca50e8073251099e9f9f7c
> 
> I'll look into this shortly.

@gavv r u sure you're testing with the latest commit ? cannot reproduce so far with

`go test -race -count=100 -v -run TestPerRequestWithTimeoutAndWithRetries`


gtourkas: > TestPerRequestWithTimeoutAndWithRetries



> > It seems there is a possible deadlock in the new tests.
> > I've run this:
> > ```
> > go test -race -count=100 -v |& pp
> > ```
> > 
> > 
> > (`pp` is https://github.com/maruel/panicparse)
> > And eventually `TestPerRequestWithTimeoutAndWithRetries` stucked and then terminated by timeout: https://gist.github.com/gavv/e9aa492ed0ca50e8073251099e9f9f7c
> 
> I'll look into this shortly.

@gavv can you reproduce with 

`go test -race -count=100 -v -run TestPerRequestWithTimeoutAndWithRetries`

?

gtourkas: @gavv I can reproduce as well. The problem is with the WithTimeout tests. There is no problem when omitted. I'll try to debug as well. 
gtourkas: @gavv It seems that the problem is the default timeout for go test, which is 10m(inutes). A single test lasts in my environment 6.317s(econds). For -count=100, this will be at least 631.7s>10m. Simply increasing the timeout to 11m will result in passing the test:

`go test -count=100 -timeout=11m -v`


![image](https://user-images.githubusercontent.com/3426478/112905589-51dfb480-90f3-11eb-83d0-02064469dbf0.png)

In any case, I've set the various WithTimeouts to 500ms so that the tests complete under the default 10m for count=100. 

gavv: Never mind, I expected the timeout is per-test... Thanks for looking into it and for the PR!

Patch:
diff --git a/README.md b/README.md
index caea392..61d6fd2 100644
--- a/README.md
+++ b/README.md
@@ -585,6 +585,48 @@ e := httpexpect.WithConfig(httpexpect.Config{
 	},
 })
 ```
+##### Global time-out/cancellation
+
+```go
+handler := FruitsHandler()
+
+server := httptest.NewServer(handler)
+defer server.Close()
+
+e := WithConfig(Config{
+  BaseURL:  server.URL,
+  Reporter: httpexpect.NewAssertReporter(t),
+  Context: ctx,
+})
+
+ctx, cancel := context.WithCancel(context.Background())
+
+go func() {
+  time.Sleep(time.Duration(5)*time.Seconds)
+  cancel()
+}()
+
+e.GET("/fruits").
+  Expect().
+  Status(http.StatusOK)
+```
+
+
+##### Per-request time-out 
+
+```go
+e.GET("/fruits").
+  WithTimeout(time.Duration(5)*time.Seconds).
+  Expect().
+  Status(http.StatusOK)
+
+// combined with retries
+e.POST("/fruits").
+  WithMaxRetries(5).
+  WithTimeout(time.Duration(10)*time.Seconds).
+  Expect().
+  Status(http.StatusOK)
+```
 
 ## Similar packages
 
diff --git a/expect.go b/expect.go
index 4962e87..c43a6c6 100644
--- a/expect.go
+++ b/expect.go
@@ -66,6 +66,7 @@
 package httpexpect
 
 import (
+	"context"
 	"io"
 	"net/http"
 	"net/http/cookiejar"
@@ -129,6 +130,14 @@ type Config struct {
 	// you're happy with their format, but want to send logs somewhere
 	// else instead of testing.TB.
 	Printers []Printer
+
+	// Context is passed to all requests. It is typically used for request cancellation,
+	// either explicit or after a time-out.
+	// May be nil.
+	//
+	// You can use the Request.WithContext for per-request context and Request.WithTimeout
+	// for per-request timeout.
+	Context context.Context
 }
 
 // RequestFactory is used to create all http.Request objects.
diff --git a/request.go b/request.go
index 7ae59f5..eb50dca 100644
--- a/request.go
+++ b/request.go
@@ -2,6 +2,7 @@ package httpexpect
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -46,6 +47,7 @@ type Request struct {
 	wsUpgrade      bool
 	transforms     []func(*http.Request)
 	matchers       []func(*Response)
+	timeout        time.Duration
 }
 
 // NewRequest returns a new Request object.
@@ -448,6 +450,56 @@ func (r *Request) WithWebsocketDialer(dialer WebsocketDialer) *Request {
 	return r
 }
 
+// WithContext sets the context.
+//
+// Config.Context will be overwritten.
+//
+// Any retries will stop after one is cancelled.
+// If the intended behavior is to continue any further retries, use WithTimeout.
+//
+// Example:
+//  ctx, _ = context.WithTimeout(context.Background(), time.Duration(3)*time.Second)
+//  req := NewRequest(config, "GET", "/path")
+//  req.WithContext(ctx)
+//  req.Expect().Status(http.StatusOK)
+func (r *Request) WithContext(ctx context.Context) *Request {
+	if r.chain.failed() {
+		return r
+	}
+	if ctx == nil {
+		r.chain.fail("\nunexpected nil ctx in WithContext")
+		return r
+	}
+
+	r.config.Context = ctx
+
+	return r
+}
+
+// WithTimeout sets a timeout duration for the request.
+//
+// Will attach to the request a context.WithTimeout around the Config.Context
+// or any context set WithContext. If these are nil, the new context will be
+// created on top of a context.Background().
+//
+// Any retries will continue after one is cancelled.
+// If the intended behavior is to stop any further retries, use WithContext or
+// Config.Context.
+//
+// Example:
+//  req := NewRequest(config, "GET", "/path")
+//  req.WithTimeout(time.Duration(3)*time.Second)
+//  req.Expect().Status(http.StatusOK)
+func (r *Request) WithTimeout(timeout time.Duration) *Request {
+	if r.chain.failed() {
+		return r
+	}
+
+	r.timeout = timeout
+
+	return r
+}
+
 // WithPath substitutes named parameters in url path.
 //
 // value is converted to string using fmt.Sprint(). If there is no named
@@ -1183,6 +1235,10 @@ func (r *Request) encodeRequest() bool {
 		r.http.Body = http.NoBody
 	}
 
+	if r.config.Context != nil {
+		r.http = r.http.WithContext(r.config.Context)
+	}
+
 	r.setupRedirects()
 
 	return true
@@ -1270,11 +1326,24 @@ func (r *Request) retryRequest(reqFunc func() (resp *http.Response, err error))
 			printer.Request(r.http)
 		}
 
-		start := time.Now()
+		func() {
+			if r.timeout > 0 {
+				var ctx context.Context
+				var cancel context.CancelFunc
+				if r.config.Context != nil {
+					ctx, cancel = context.WithTimeout(r.config.Context, r.timeout)
+				} else {
+					ctx, cancel = context.WithTimeout(context.Background(), r.timeout)
+				}
 
-		resp, err = reqFunc()
+				defer cancel()
+				r.http = r.http.WithContext(ctx)
+			}
 
-		elapsed = time.Since(start)
+			start := time.Now()
+			resp, err = reqFunc()
+			elapsed = time.Since(start)
+		}()
 
 		if resp != nil {
 			for _, printer := range r.config.Printers {


Test patch:
diff --git a/e2e_context_test.go b/e2e_context_test.go
new file mode 100644
index 0000000..a3e4957
--- /dev/null
+++ b/e2e_context_test.go
@@ -0,0 +1,384 @@
+package httpexpect
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+const (
+	TimeOutDuration = 500 * time.Millisecond
+)
+
+type waitHandler struct {
+	mux           *http.ServeMux
+	callCount     int
+	retriesToFail int
+	retriesDone   chan struct{}
+	sync.RWMutex
+}
+
+func (h *waitHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	h.mux.ServeHTTP(w, r)
+}
+
+func (h *waitHandler) WaitForContextCancellation(w http.ResponseWriter, r *http.Request) {
+	callCount := h.IncrCallCount()
+	// if retries-to-fail are not set then simply wait for the cancellation
+	if h.retriesToFail == 0 {
+		<-r.Context().Done()
+	} else {
+		// if retries-to-fail are set then make sure they are exhausted before
+		// waiting for cancellation
+		if callCount < h.retriesToFail+1 {
+			w.WriteHeader(http.StatusInternalServerError)
+		} else {
+			h.retriesDone <- struct{}{}
+			<-r.Context().Done()
+		}
+	}
+}
+
+func (h *waitHandler) WaitForPerRequestTimeout(w http.ResponseWriter, r *http.Request) {
+	callCount := h.IncrCallCount()
+	// if retries-to-fail are not set or not exhausted yet, simply wait for
+	// the timeout
+	if h.retriesToFail == 0 || callCount < h.retriesToFail+1 {
+		<-r.Context().Done()
+	} else {
+		// otherwise succeed
+		w.WriteHeader(http.StatusOK)
+	}
+}
+
+func (h *waitHandler) IncrCallCount() int {
+	h.Lock()
+	defer h.Unlock()
+	h.callCount++
+	r := h.callCount
+	return r
+}
+
+func (h *waitHandler) GetCallCount() int {
+	h.RLock()
+	defer h.RUnlock()
+	return h.callCount
+}
+
+func newWaitHandler(retriesToFail int) *waitHandler {
+	mux := http.NewServeMux()
+
+	handler := &waitHandler{
+		mux:           mux,
+		retriesToFail: retriesToFail,
+		retriesDone:   make(chan struct{}),
+	}
+
+	mux.HandleFunc("/WaitForContextCancellation", handler.WaitForContextCancellation)
+	mux.HandleFunc("/WaitForPerRequestTimeout", handler.WaitForPerRequestTimeout)
+
+	return handler
+}
+
+func (h *waitHandler) waitForRetries() {
+	<-h.retriesDone
+}
+
+// the expErrorSuppressor is used as a Reporter to suppress an expected error
+type expErrorSuppressor struct {
+	backend          *assert.Assertions
+	isExpError       isExpErrorFunc
+	expErrorOccurred bool
+}
+
+type isExpErrorFunc func(message string, args ...interface{}) bool
+
+func newExpErrorSuppressor(t assert.TestingT,
+	isExpectedErr isExpErrorFunc) *expErrorSuppressor {
+	return &expErrorSuppressor{backend: assert.New(t),
+		isExpError: isExpectedErr}
+}
+
+func (r *expErrorSuppressor) Errorf(message string, args ...interface{}) {
+	if !r.isExpError(message, args) {
+		r.backend.Fail(fmt.Sprintf(message, args...))
+	} else {
+		r.expErrorOccurred = true
+	}
+}
+
+func TestGlobalContextCancel(t *testing.T) {
+	handler := newWaitHandler(0)
+
+	server := httptest.NewServer(handler)
+	defer server.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// config with context cancel suppression
+	reporter := newExpErrorSuppressor(t,
+		func(message string, args ...interface{}) bool {
+			return strings.HasSuffix(message, "context canceled")
+		})
+	e := WithConfig(Config{
+		BaseURL:  server.URL,
+		Reporter: reporter,
+		Context:  ctx,
+	})
+
+	done := make(chan struct{})
+
+	go func() {
+		e.GET("/WaitForContextCancellation").
+			Expect()
+		done <- struct{}{}
+	}()
+
+	cancel()
+
+	<-done
+
+	// expected error should occur
+	assert.True(t, reporter.expErrorOccurred)
+}
+
+func TestGlobalContextWithRetries(t *testing.T) {
+	maxRetries := 3
+	retriesToFail := 2
+	handler := newWaitHandler(retriesToFail)
+
+	server := httptest.NewServer(handler)
+	defer server.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// config with context cancel suppression
+	reporter := newExpErrorSuppressor(t,
+		func(message string, args ...interface{}) bool {
+			return strings.HasSuffix(message, "context canceled")
+		})
+	e := WithConfig(Config{
+		BaseURL:  server.URL,
+		Reporter: reporter,
+		Context:  ctx,
+	})
+
+	done := make(chan struct{})
+
+	go func() {
+		e.GET("/WaitForContextCancellation").
+			WithMaxRetries(maxRetries).
+			Expect()
+		done <- struct{}{}
+	}()
+
+	handler.waitForRetries() // wait for the retries-set-to-fail
+	cancel()                 // cancel the rest
+
+	<-done
+
+	// expected error should occur
+	assert.True(t, reporter.expErrorOccurred)
+	// first call + retries to fail should be the call count
+	assert.Equal(t, retriesToFail+1, handler.GetCallCount())
+}
+
+func TestPerRequestContext(t *testing.T) {
+	handler := newWaitHandler(0)
+
+	server := httptest.NewServer(handler)
+	defer server.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// config with context cancel suppression
+	reporter := newExpErrorSuppressor(t,
+		func(message string, args ...interface{}) bool {
+			return strings.HasSuffix(message, "context canceled")
+		})
+	e := WithConfig(Config{
+		BaseURL:  server.URL,
+		Reporter: reporter,
+	})
+
+	done := make(chan struct{})
+
+	go func() {
+		e.GET("/WaitForContextCancellation").
+			WithContext(ctx).
+			Expect()
+		done <- struct{}{}
+	}()
+
+	cancel()
+
+	<-done
+
+	// expected error should occur
+	assert.True(t, reporter.expErrorOccurred)
+}
+
+func TestPerRequestContextWithRetries(t *testing.T) {
+	maxRetries := 3
+	retriesToFail := 2
+	handler := newWaitHandler(retriesToFail)
+
+	server := httptest.NewServer(handler)
+	defer server.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// config with context cancel suppression
+	reporter := newExpErrorSuppressor(t,
+		func(message string, args ...interface{}) bool {
+			return strings.HasSuffix(message, "context canceled")
+		})
+	e := WithConfig(Config{
+		BaseURL:  server.URL,
+		Reporter: reporter,
+	})
+
+	done := make(chan struct{})
+
+	go func() {
+		e.GET("/WaitForContextCancellation").
+			WithMaxRetries(maxRetries).
+			WithContext(ctx).
+			Expect()
+		done <- struct{}{}
+	}()
+
+	handler.waitForRetries() // wait for the retries-set-to-fail
+	cancel()                 // cancel the rest
+
+	<-done
+
+	// expected error should occur
+	assert.True(t, reporter.expErrorOccurred)
+	// first call + retries to fail should be the call count
+	assert.Equal(t, retriesToFail+1, handler.GetCallCount())
+}
+
+func TestPerRequestWithTimeout(t *testing.T) {
+	handler := newWaitHandler(0)
+
+	server := httptest.NewServer(handler)
+	defer server.Close()
+
+	// config with context deadline expected error
+	reporter := newExpErrorSuppressor(t,
+		func(message string, args ...interface{}) bool {
+			return strings.HasSuffix(message, "context deadline exceeded")
+		})
+	e := WithConfig(Config{
+		BaseURL:  server.URL,
+		Reporter: reporter,
+	})
+
+	e.GET("/WaitForPerRequestTimeout").
+		WithTimeout(TimeOutDuration).
+		Expect()
+
+	// expected error should occur
+	assert.True(t, reporter.expErrorOccurred)
+}
+
+func TestPerRequestWithTimeoutAndWithRetries(t *testing.T) {
+	maxRetries := 3
+	retriesToFail := 2
+	handler := newWaitHandler(retriesToFail)
+
+	server := httptest.NewServer(handler)
+	defer server.Close()
+
+	// this call will terminate with success
+	e := WithConfig(Config{
+		BaseURL:  server.URL,
+		Reporter: NewAssertReporter(t),
+	})
+
+	e.GET("/WaitForPerRequestTimeout").
+		WithTimeout(TimeOutDuration).
+		WithMaxRetries(maxRetries).
+		Expect().
+		Status(http.StatusOK)
+
+	// first call + retries to fail should be the call count
+	assert.Equal(t, retriesToFail+1, handler.GetCallCount())
+}
+
+func TestPerRequestWithContextAndTimeout_CancelledByTimeout(t *testing.T) {
+	handler := newWaitHandler(0)
+
+	server := httptest.NewServer(handler)
+	defer server.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// config with context deadline expected error
+	reporter := newExpErrorSuppressor(t,
+		func(message string, args ...interface{}) bool {
+			return strings.HasSuffix(message, "context deadline exceeded")
+		})
+	e := WithConfig(Config{
+		BaseURL:  server.URL,
+		Reporter: reporter,
+	})
+
+	e.GET("/WaitForPerRequestTimeout").
+		WithContext(ctx).
+		WithTimeout(TimeOutDuration).
+		Expect()
+
+	// expected error should occur
+	assert.True(t, reporter.expErrorOccurred)
+}
+
+func TestPerRequestWithContextAndTimeout_CancelledByContext(t *testing.T) {
+	handler := newWaitHandler(0)
+
+	server := httptest.NewServer(handler)
+	defer server.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// config with context deadline expected error
+	reporter := newExpErrorSuppressor(t,
+		func(message string, args ...interface{}) bool {
+			return strings.HasSuffix(message, "context canceled")
+		})
+	e := WithConfig(Config{
+		BaseURL:  server.URL,
+		Reporter: reporter,
+	})
+
+	done := make(chan struct{})
+
+	go func() {
+		e.GET("/WaitForContextCancellation").
+			WithContext(ctx).
+			WithTimeout(TimeOutDuration).
+			Expect()
+		done <- struct{}{}
+	}()
+
+	cancel() // cancel the rest
+
+	<-done
+
+	// expected error should occur
+	assert.True(t, reporter.expErrorOccurred)
+}