Adding custom context to Request (#126) (#127)

closes #126

README.md

@@ -237,6 +237,19 @@ func TestApi(t *testing.T) {
 }
 ```
 
+#### Pass a custom context to the request
+
+```go
+func TestApi(t *testing.T) {
+	apitest.Handler(handler).
+		Get("/hello").
+		WithContext(context.TODO()).
+		Expect(t).
+		Status(http.StatusOK).
+		End()
+}
+```
+
 #### Provide cookies in the request
 
 ```go

apitest.go

@@ -2,6 +2,7 @@ package apitest
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"hash/fnv"
@@ -229,6 +230,7 @@ type Request struct {
 	multipart       *multipart.Writer
 	cookies         []*Cookie
 	basicAuth       string
+	context         context.Context
 	apiTest         *APITest
 }
 
@@ -484,6 +486,12 @@ func (r *Request) BasicAuth(username, password string) *Request {
 	return r
 }
 
+// WithContext is a builder method to set a context on the request
+func (r *Request) WithContext(ctx context.Context) *Request {
+	r.context = ctx
+	return r
+}
+
 // FormData is a builder method to set the body form data
 // Also sets the content type of the request to application/x-www-form-urlencoded
 func (r *Request) FormData(name string, values ...string) *Request {
@@ -976,6 +984,10 @@ func (a *APITest) buildRequest() *http.Request {
 	}
 
 	req, _ := http.NewRequest(a.request.method, a.request.url, bytes.NewBufferString(a.request.body))
+	if a.request.context != nil {
+		req = req.WithContext(a.request.context)
+	}
+
 	req.URL.RawQuery = formatQuery(a.request)
 	req.Host = SystemUnderTestDefaultName
 	if a.networkingEnabled {

apitest_test.go

@@ -428,6 +428,51 @@ func TestApiTest_AddsBasicAuthToRequest(t *testing.T) {
 		End()
 }
 
+func TestApiTest_AddsTimedOutContextToRequest(t *testing.T) {
+
+	handler := http.NewServeMux()
+	handler.HandleFunc("/hello", func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(time.Nanosecond * 10)
+		if r.Context().Err() == context.DeadlineExceeded {
+			w.WriteHeader(http.StatusRequestTimeout)
+		}
+		w.WriteHeader(http.StatusOK)
+	})
+
+	timeoutCtx, cancel := context.WithTimeout(context.Background(), time.Nanosecond*1)
+	defer cancel()
+
+	apitest.New("test with timed out context").
+		Handler(handler).
+		Get("/hello").
+		WithContext(timeoutCtx).
+		Expect(t).
+		Status(http.StatusRequestTimeout).
+		End()
+}
+
+func TestApiTest_AddsCancelledContextToRequest(t *testing.T) {
+
+	handler := http.NewServeMux()
+	handler.HandleFunc("/hello", func(w http.ResponseWriter, r *http.Request) {
+		if r.Context().Err() == context.Canceled {
+			w.WriteHeader(http.StatusNoContent)
+		}
+		w.WriteHeader(http.StatusOK)
+	})
+
+	cancelCtx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	apitest.New("test with canceled context").
+		Handler(handler).
+		Get("/hello").
+		WithContext(cancelCtx).
+		Expect(t).
+		Status(http.StatusNoContent).
+		End()
+}
+
 func TestApiTest_GraphQLQuery(t *testing.T) {
 	apitest.New().
 		HandlerFunc(func(w http.ResponseWriter, r *http.Request) {