Add Context cancellation to clock.

common/clock/testclock/testclock_test.go

Index: common/clock/testclock/testclock_test.go
diff --git a/common/clock/testclock/testclock_test.go b/common/clock/testclock/testclock_test.go
index f063c02f3d3e33b9f2e245fb4228f0f4a5037443..f1d19c2452c11d84d93ea0adab77ff2cf69e9ffc 100644
--- a/common/clock/testclock/testclock_test.go
+++ b/common/clock/testclock/testclock_test.go
@@ -10,6 +10,7 @@ import (
 
 	"github.com/luci/luci-go/common/clock"
 	. "github.com/smartystreets/goconvey/convey"
+	"golang.org/x/net/context"
 )
 
 func TestTestClock(t *testing.T) {
@@ -17,41 +18,55 @@ func TestTestClock(t *testing.T) {
 
 	Convey(`A testing clock instance`, t, func() {
 		now := time.Date(2015, 01, 01, 00, 00, 00, 00, time.UTC)
-		c := New(now)
+		ctx, clk := UseTime(context.Background(), now)
 
 		Convey(`Returns the current time.`, func() {
-			So(c.Now(), ShouldResemble, now)
+			So(clk.Now(), ShouldResemble, now)
 		})
 
 		Convey(`When sleeping with a time of zero, immediately awakens.`, func() {
-			c.Sleep(0)
-			So(c.Now(), ShouldResemble, now)
+			clk.Sleep(ctx, 0)
+			So(clk.Now(), ShouldResemble, now)
+		})
+
+		Convey(`Will panic if going backwards in time.`, func() {
+			So(func() { clk.Add(-1 * time.Second) }, ShouldPanic)
 		})
 
 		Convey(`When sleeping for a period of time, awakens when signalled.`, func() {
 			sleepingC := make(chan struct{})
-			c.SetTimerCallback(func(_ time.Duration, _ clock.Timer) {
+			clk.SetTimerCallback(func(_ time.Duration, _ clock.Timer) {
 				close(sleepingC)
 			})
 
 			awakeC := make(chan time.Time)
 			go func() {
-				c.Sleep(2 * time.Second)
-				awakeC <- c.Now()
+				clk.Sleep(ctx, 2*time.Second)
+				awakeC <- clk.Now()
 			}()
 
 			<-sleepingC
-			c.Set(now.Add(1 * time.Second))
-			c.Set(now.Add(2 * time.Second))
+			clk.Set(now.Add(1 * time.Second))
+			clk.Set(now.Add(2 * time.Second))
 			So(<-awakeC, ShouldResemble, now.Add(2*time.Second))
 		})
 
 		Convey(`Awakens after a period of time.`, func() {
-			afterC := c.After(2 * time.Second)
+			afterC := clk.After(ctx, 2*time.Second)
+
+			clk.Set(now.Add(1 * time.Second))
+			clk.Set(now.Add(2 * time.Second))
+			So(<-afterC, ShouldResemble, clock.TimerResult{now.Add(2 * time.Second), nil})
+		})
+
+		Convey(`When sleeping, awakens if canceled.`, func() {
+			ctx, cancelFunc := context.WithCancel(ctx)
+
+			clk.SetTimerCallback(func(_ time.Duration, _ clock.Timer) {
+				cancelFunc()
+			})
 
-			c.Set(now.Add(1 * time.Second))
-			c.Set(now.Add(2 * time.Second))
-			So(<-afterC, ShouldResemble, now.Add(2*time.Second))
+			So(clk.Sleep(ctx, time.Second), ShouldEqual, context.Canceled)
 		})
 	})
 }