Add Context cancellation to clock.

common/clock/cancelSleep_test.go

+// 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)
+                })
+        })
+}