Add Context cancellation to clock.

common/clock/testclock/testclock_test.go

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 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) {
                                 close(sleepingC)
                         })
 
                         awakeC := make(chan time.Time)
                         go func() {
-                                c.Sleep(2 * time.Second)
+                                c.Sleep(ctx, 2*time.Second)
                                 awakeC <- c.Now()
                         }()
 
                         <-sleepingC
                         c.Set(now.Add(1 * time.Second))
                         c.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 := 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)
                 })
         })
 }