context: setup timeout context

context/context.go

@@ -7,10 +7,24 @@ import (
 
 type Store interface {
 	Fetch() string
+	Cancel()
 }
 
 func Server(store Store) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
-		fmt.Fprint(w, store.Fetch())
+		ctx := r.Context()
+
+		data := make(chan string, 1)
+
+		go func() {
+			data <- store.Fetch()
+		}()
+
+		select {
+		case d := <-data:
+			fmt.Fprint(w, d)
+		case <-ctx.Done():
+			store.Cancel()
+		}
 	}
 }

context/context_test.go

@@ -1,22 +1,32 @@
 package context
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 )
 
 type SpyStore struct {
 	response  string
+	cancelled bool
 }
 
 func (s *SpyStore) Fetch() string {
+	time.Sleep(100 * time.Millisecond)
 	return s.response
 }
 
+func (s *SpyStore) Cancel() {
+	s.cancelled = true
+}
+
 func TestServer(t *testing.T) {
+	t.Run("basic get store", func(t *testing.T) {
 		data := "hello, world"
-	srv := Server(&SpyStore{data})
+		store := &SpyStore{response: data}
+		srv := Server(store)
 
 		request := httptest.NewRequest(http.MethodGet, "/", nil)
 		response := httptest.NewRecorder()
@@ -26,4 +36,31 @@ func TestServer(t *testing.T) {
 		if response.Body.String() != data {
 			t.Errorf(`got "%s", want "%s"`, response.Body.String(), data)
 		}
+
+		if store.cancelled {
+			t.Error("it should not have cancelled the store")
+		}
+	})
+
+	t.Run("tells store to cancel work if request is cancelled", func(t *testing.T) {
+		data := "hello, world"
+		store := &SpyStore{response: data}
+		srv := Server(store)
+
+		request := httptest.NewRequest(http.MethodGet, "/", nil)
+
+		cancellingCtx, cancel := context.WithCancel(request.Context())
+		// Wait 5ms to call cancel
+		time.AfterFunc(5*time.Millisecond, cancel)
+
+		request = request.WithContext(cancellingCtx)
+
+		response := httptest.NewRecorder()
+
+		srv.ServeHTTP(response, request)
+
+		if !store.cancelled {
+			t.Error("store was not told to cancel")
+		}
+	})
 }