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..6bfc090d2b8ac94b67d97c9f23f8480944e71e7a 100644
--- a/common/clock/testclock/testclock_test.go
+++ b/common/clock/testclock/testclock_test.go
@@ -5,29 +5,44 @@
 package testclock
 
 import (
+	"fmt"
+	"strings"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/luci/luci-go/common/clock"
 	. "github.com/smartystreets/goconvey/convey"
+	"golang.org/x/net/context"
 )
 
+// trashTimer is a useless implementation of clock.Timer specifically designed
+// to exist and not be a test timer type.
+type trashTimer struct {
+	clock.Timer
+}
+
 func TestTestClock(t *testing.T) {
 	t.Parallel()
 
 	Convey(`A testing clock instance`, t, func() {
 		now := time.Date(2015, 01, 01, 00, 00, 00, 00, time.UTC)
 		c := New(now)
+		ctx := context.Background()
 
 		Convey(`Returns the current time.`, func() {
 			So(c.Now(), ShouldResemble, now)
 		})
 
 		Convey(`When sleeping with a time of zero, immediately awakens.`, func() {
-			c.Sleep(0)
+			c.Sleep(ctx, 0)
 			So(c.Now(), ShouldResemble, now)
 		})
 
+		Convey(`Will panic if going backwards in time.`, func() {
+			So(func() { c.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) {
@@ -36,7 +51,7 @@ func TestTestClock(t *testing.T) {
 
 			awakeC := make(chan time.Time)
 			go func() {
-				c.Sleep(2 * time.Second)
+				c.Sleep(ctx, 2*time.Second)
 				awakeC <- c.Now()
 			}()
 
@@ -47,11 +62,62 @@ func TestTestClock(t *testing.T) {
 		})
 
 		Convey(`Awakens after a period of time.`, func() {
-			afterC := c.After(2 * time.Second)
+			afterC := c.After(ctx, 2*time.Second)
 
 			c.Set(now.Add(1 * time.Second))
 			c.Set(now.Add(2 * time.Second))
-			So(<-afterC, ShouldResemble, 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)
+
+			c.SetTimerCallback(func(_ time.Duration, _ clock.Timer) {
+				cancelFunc()
+			})
+
+			So(c.Sleep(ctx, time.Second), ShouldEqual, context.Canceled)
+		})
+
+		Convey(`Can set and retrieve timer tags.`, func() {
+			var tagMu sync.Mutex
+			tagMap := map[string]struct{}{}
+
+			// On the last timer callback, advance time past the timer threshold.
+			timers := make([]clock.Timer, 10)
+			count := 0
+			c.SetTimerCallback(func(_ time.Duration, t clock.Timer) {
+				tagMu.Lock()
+				defer tagMu.Unlock()
+
+				tagMap[strings.Join(GetTags(t), "")] = struct{}{}
+				count++
+				if count == len(timers) {
+					c.Add(time.Second)
+				}
+			})
+
+			for i := range timers {
+				t := c.NewTimer(clock.Tag(ctx, fmt.Sprintf("%d", i)))
+				t.Reset(time.Second)
+				timers[i] = t
+			}
+
+			// Wait until all timers have expired.
+			for _, t := range timers {
+				<-t.GetC()
+			}
+
+			// Did we see all of their tags?
+			for i := range timers {
+				_, ok := tagMap[fmt.Sprintf("%d", i)]
+				So(ok, ShouldBeTrue)
+			}
+		})
+
+		Convey(`A non-test timer has a nil tag.`, func() {
+			t := trashTimer{}
+			So(GetTags(t), ShouldBeNil)
 		})
 	})
 }