Add Context cancellation to clock.

common/clock/cancelSleep_test.go

Index: common/clock/cancelSleep_test.go
diff --git a/common/clock/cancelSleep_test.go b/common/clock/cancelSleep_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..c6e57ba09713e4df2e912936ce288e8a56618ccd
--- /dev/null
+++ b/common/clock/cancelSleep_test.go
@@ -0,0 +1,76 @@
+// Copyright 2016 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 clock
+
+import (
+	"testing"
+	"time"
+
+	"golang.org/x/net/context"
+
+	. "github.com/smartystreets/goconvey/convey"
+)
+
+// instrumentedAfterClock is a Clock implementation that implements After. It
+// is instrumented to allow the test code to enforce operational order.
+type instrumentedAfterClock struct {
+	Clock
+
+	sleepingC chan struct{}
+	wakeC     chan struct{}
+}
+
+func (tc *instrumentedAfterClock) After(time.Duration) <-chan time.Time {
+	awakeC := make(chan time.Time, 1)
+	go func() {
+		if tc.sleepingC != nil {
+			tc.sleepingC <- struct{}{}
+		}
+
+		// Wait for signal to wake.
+		<-tc.wakeC
+		awakeC <- time.Time{}
+	}()
+
+	return awakeC
+}
+
+// TestCancelSleep tests the CancelSleep function.
+func TestCancelSleep(t *testing.T) {
+	t.Parallel()
+
+	Convey(`A cancelable Context`, t, func() {
+		tc := instrumentedAfterClock{
+			wakeC: make(chan struct{}, 1),
+		}
+		defer close(tc.wakeC)
+
+		c := Set(context.Background(), &tc)
+		c, cancelFunc := context.WithCancel(c)
+
+		Convey(`Will perform a full sleep if the Context isn't canceled.`, func() {
+			tc.wakeC <- struct{}{}
+
+			So(CancelSleep(c, time.Minute), ShouldBeNil)
+		})
+
+		Convey(`Will return Context error if canceled.`, func() {
+			tc.wakeC <- struct{}{}
+			cancelFunc()
+
+			So(CancelSleep(c, time.Minute), ShouldEqual, context.Canceled)
+		})
+
+		Convey(`Will terminate the Sleep prematurely if the Context is canceled.`, func() {
+			tc.sleepingC = make(chan struct{})
+			go func() {
+				<-tc.sleepingC
+				cancelFunc()
+			}()
+
+			So(CancelSleep(c, time.Minute), ShouldEqual, context.Canceled)
+		})
+	})
+}