Add Context cancellation to clock.

common/clock/clockcontext_test.go

Index: common/clock/clockcontext_test.go
diff --git a/common/clock/clockcontext_test.go b/common/clock/clockcontext_test.go
index c550a45e94bf29b40da99d6bc2d00b45992dff1d..c360d04c49d7918903794d8c710040794eee01fe 100644
--- a/common/clock/clockcontext_test.go
+++ b/common/clock/clockcontext_test.go
@@ -12,99 +12,117 @@ import (
 	"golang.org/x/net/context"
 )
 
-// A test blocking operation.
-func blockingOperation(ctx context.Context, releaseC chan struct{}) (bool, error) {
-	select {
-	case <-ctx.Done():
-		return false, ctx.Err()
-
-	case <-releaseC:
-		return true, nil
-	}
+// manualClock is a partial Clock implementation that allows us to release
+// blocking calls.
+type manualClock struct {
+	Clock
+
+	now             time.Time
+	timeoutCallback func(time.Duration) bool
+	testFinishedC   chan struct{}
 }
 
-func TestClockContext(t *testing.T) {
-	t.Parallel()
+func (mc *manualClock) Now() time.Time {
+	return mc.now
+}
 
-	Convey(`A context with a deadline wrapping a cancellable parent`, t, func() {
-		now := time.Now()
-		cctx, pcf := context.WithCancel(context.Background())
-		ctx, cf := WithTimeout(cctx, 10*time.Millisecond)
-		releaseC := make(chan struct{})
+func (mc *manualClock) After(ctx context.Context, d time.Duration) <-chan TimerResult {
+	resultC := make(chan TimerResult)
+	go func() {
+		ar := TimerResult{}
+		defer func() {
+			resultC <- ar
+		}()
+
+		// If we are instructed to immediately timeout, do so.
+		if cb := mc.timeoutCallback; cb != nil && cb(d) {
+			return
+		}
+
+		select {
+		case <-ctx.Done():
+			ar.Err = ctx.Err()
+		case <-mc.testFinishedC:
+			break
+		}
+	}()
+
+	return resultC
+}
 
-		Convey(`Successfully reports its deadline.`, func() {
-			deadline, ok := ctx.Deadline()
-			So(ok, ShouldBeTrue)
-			So(deadline.After(now), ShouldBeTrue)
-		})
+func wait(c context.Context) error {
+	<-c.Done()
+	return c.Err()
+}
 
-		Convey(`Will successfully time out.`, func() {
-			released, err := blockingOperation(ctx, releaseC)
-			So(released, ShouldBeFalse)
-			So(err, ShouldEqual, context.DeadlineExceeded)
-		})
+func TestClockContext(t *testing.T) {
+	t.Parallel()
 
-		Convey(`Will successfully cancel with its cancel func.`, func() {
-			go func() {
-				cf()
-			}()
-			released, err := blockingOperation(ctx, releaseC)
-			So(released, ShouldBeFalse)
-			So(err, ShouldEqual, context.Canceled)
+	Convey(`A manual testing clock`, t, func() {
+		mc := manualClock{
+			now:           time.Date(2016, 1, 1, 0, 0, 0, 0, time.Local),
+			testFinishedC: make(chan struct{}),
+		}
+		defer close(mc.testFinishedC)
+
+		Convey(`A context with a deadline wrapping a cancellable parent`, func() {
+			cctx, pcf := context.WithCancel(Set(context.Background(), &mc))
+			ctx, cf := WithTimeout(cctx, 10*time.Millisecond)
+
+			Convey(`Successfully reports its deadline.`, func() {
+				deadline, ok := ctx.Deadline()
+				So(ok, ShouldBeTrue)
+				So(deadline.After(mc.now), ShouldBeTrue)
+			})
+
+			Convey(`Will successfully time out.`, func() {
+				mc.timeoutCallback = func(time.Duration) bool {
+					return true
+				}
+				So(wait(ctx), ShouldEqual, context.DeadlineExceeded)
+			})
+
+			Convey(`Will successfully cancel with its cancel func.`, func() {
+				go func() {
+					cf()
+				}()
+				So(wait(ctx), ShouldEqual, context.Canceled)
+			})
+
+			Convey(`Will successfully cancel if the parent is canceled.`, func() {
+				go func() {
+					pcf()
+				}()
+				So(wait(ctx), ShouldEqual, context.Canceled)
+			})
 		})
 
-		Convey(`Will successfully cancel if the parent is cancelled.`, func() {
-			go func() {
-				pcf()
-			}()
-			released, err := blockingOperation(ctx, releaseC)
-			So(released, ShouldBeFalse)
-			So(err, ShouldEqual, context.Canceled)
-		})
+		Convey(`A context with a deadline wrapping a parent with a shorter deadline`, func() {
+			cctx, _ := context.WithTimeout(context.Background(), 10*time.Millisecond)
+			ctx, cf := WithTimeout(cctx, 1*time.Hour)
 
-		Convey(`Will release before the deadline.`, func() {
-			go func() {
-				close(releaseC)
-			}()
-			released, err := blockingOperation(ctx, releaseC)
-			So(released, ShouldBeTrue)
-			So(err, ShouldBeNil)
-		})
-	})
+			Convey(`Will successfully time out.`, func() {
+				mc.timeoutCallback = func(d time.Duration) bool {
+					return d == 10*time.Millisecond
+				}
 
-	Convey(`A context with a deadline wrapping a parent with a shorter deadline`, t, func() {
-		cctx, _ := context.WithTimeout(context.Background(), 10*time.Millisecond)
-		ctx, cf := WithTimeout(cctx, 1*time.Hour)
-		releaseC := make(chan struct{})
+				So(wait(ctx), ShouldEqual, context.DeadlineExceeded)
+			})
 
-		Convey(`Will successfully time out.`, func() {
-			released, err := blockingOperation(ctx, releaseC)
-			So(released, ShouldBeFalse)
-			So(err, ShouldEqual, context.DeadlineExceeded)
+			Convey(`Will successfully cancel with its cancel func.`, func() {
+				go func() {
+					cf()
+				}()
+				So(wait(ctx), ShouldEqual, context.Canceled)
+			})
 		})
 
-		Convey(`Will successfully cancel with its cancel func.`, func() {
-			go func() {
-				cf()
-			}()
-			released, err := blockingOperation(ctx, releaseC)
-			So(released, ShouldBeFalse)
-			So(err, ShouldEqual, context.Canceled)
-		})
-	})
+		Convey(`A context with a deadline in the past`, func() {
+			ctx, _ := WithDeadline(context.Background(), mc.now.Add(-time.Second))
 
-	Convey(`A context with a deadline in the past`, t, func() {
-		ctx, _ := WithDeadline(context.Background(), time.Unix(0, 0))
-		releaseC := make(chan struct{})
-
-		Convey(`Will time out immediately.`, func() {
-			go func() {
-				defer close(releaseC)
-				time.Sleep(50 * time.Millisecond)
-			}()
-			released, err := blockingOperation(ctx, releaseC)
-			So(released, ShouldBeFalse)
-			So(err, ShouldEqual, context.DeadlineExceeded)
+			Convey(`Will time out immediately.`, func() {
+				So(wait(ctx), ShouldEqual, context.DeadlineExceeded)
+			})
 		})
 	})
 }