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,29 +1,66 @@
 package context
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 )
 
 type SpyStore struct {
-	response string
+	response  string
+	cancelled bool
 }
 
 func (s *SpyStore) Fetch() string {
+	time.Sleep(100 * time.Millisecond)
 	return s.response
 }
 
-func TestServer(t *testing.T) {
-	data := "hello, world"
-	srv := Server(&SpyStore{data})
-
-	request := httptest.NewRequest(http.MethodGet, "/", nil)
-	response := httptest.NewRecorder()
-
-	srv.ServeHTTP(response, request)
-
-	if response.Body.String() != data {
-		t.Errorf(`got "%s", want "%s"`, response.Body.String(), data)
-	}
+func (s *SpyStore) Cancel() {
+	s.cancelled = true
+}
+
+func TestServer(t *testing.T) {
+	t.Run("basic get store", func(t *testing.T) {
+		data := "hello, world"
+		store := &SpyStore{response: data}
+		srv := Server(store)
+
+		request := httptest.NewRequest(http.MethodGet, "/", nil)
+		response := httptest.NewRecorder()
+
+		srv.ServeHTTP(response, request)
+
+		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")
+		}
+	})
 }