Move time.Sleep stub into client-specific option

This ensures that:
 - tests cannot impact either each other by stubbing the same sleep
   function
 - external packages can stub the time.Sleep call

hyperbahn/advertise.go

@@ -41,9 +41,6 @@ const (
 	advertiseRetryInterval = 1 * time.Second
 )
 
-// timeSleep is a variable for stubbing in unit tests.
-var timeSleep = time.Sleep
-
 // ErrAdvertiseFailed is triggered when advertise fails.
 type ErrAdvertiseFailed struct {
 	// WillRetry is set to true if advertise will be retried.
@@ -84,7 +81,7 @@ func (c *Client) advertiseLoop() {
 	consecutiveFailures := uint(0)
 
 	for {
-		timeSleep(sleepFor)
+		c.sleep(sleepFor)
 		if c.IsClosed() {
 			c.tchan.Logger().Infof("Hyperbahn client closed")
 			return
@@ -129,7 +126,11 @@ func (c *Client) initialAdvertise() error {
 
 		// Back off for a while.
 		sleepFor := fuzzInterval(advertiseRetryInterval * time.Duration(1<<attempt))
-		timeSleep(sleepFor)
+		c.sleep(sleepFor)
 	}
 	return err
 }
+
+func (c *Client) sleep(d time.Duration) {
+	c.opts.TimeSleep(d)
+}

hyperbahn/advertise_test.go

@@ -39,8 +39,8 @@ import (
 func TestInitialAdvertiseFailedRetryBackoff(t *testing.T) {
 	defer testutils.SetTimeout(t, time.Second)()
 
-	sleepArgs, sleepBlock, sleepClose := testutils.SleepStub(&timeSleep)
-	defer testutils.ResetSleepStub(&timeSleep)
+	clientOpts := stubbedSleep()
+	sleepArgs, sleepBlock, sleepClose := testutils.SleepStub(&clientOpts.TimeSleep)
 
 	// We expect to retry 5 times,
 	go func() {
@@ -58,17 +58,14 @@ func TestInitialAdvertiseFailedRetryBackoff(t *testing.T) {
 		serverCh := testutils.NewServer(t, nil)
 		defer serverCh.Close()
 
-		client, err := NewClient(serverCh, configFor(hostPort), nil)
+		client, err := NewClient(serverCh, configFor(hostPort), clientOpts)
 		require.NoError(t, err, "NewClient")
 		defer client.Close()
 		assert.Error(t, client.Advertise(), "Advertise without handler should fail")
 	})
 }
 
 func TestInitialAdvertiseFailedRetryTimeout(t *testing.T) {
-	timeSleep = func(d time.Duration) {}
-	defer testutils.ResetSleepStub(&timeSleep)
-
 	withSetup(t, func(hypCh *tchannel.Channel, hyperbahnHostPort string) {
 		started := time.Now()
 		count := 0
@@ -86,7 +83,7 @@ func TestInitialAdvertiseFailedRetryTimeout(t *testing.T) {
 		json.Register(hypCh, json.Handlers{"ad": adHandler}, nil)
 
 		ch := testutils.NewServer(t, nil)
-		client, err := NewClient(ch, configFor(hyperbahnHostPort), nil)
+		client, err := NewClient(ch, configFor(hyperbahnHostPort), stubbedSleep())
 		assert.NoError(t, err, "hyperbahn NewClient failed")
 		defer client.Close()
 
@@ -104,7 +101,7 @@ func TestNotListeningChannel(t *testing.T) {
 		json.Register(hypCh, json.Handlers{"ad": adHandler}, nil)
 
 		ch := testutils.NewClient(t, nil)
-		client, err := NewClient(ch, configFor(hyperbahnHostPort), nil)
+		client, err := NewClient(ch, configFor(hyperbahnHostPort), stubbedSleep())
 		assert.NoError(t, err, "hyperbahn NewClient failed")
 		defer client.Close()
 
@@ -122,6 +119,7 @@ type retryTest struct {
 	sleepArgs  <-chan time.Duration
 	sleepBlock chan<- struct{}
 	sleepClose func()
+	timeSleep  func(time.Duration)
 
 	ch     *tchannel.Channel
 	client *Client
@@ -147,7 +145,7 @@ func (r *retryTest) adHandler(ctx json.Context, req *AdRequest) (*AdResponse, er
 func (r *retryTest) setup() {
 	r.respCh = make(chan int, 1)
 	r.reqCh = make(chan *AdRequest, 1)
-	r.sleepArgs, r.sleepBlock, r.sleepClose = testutils.SleepStub(&timeSleep)
+	r.sleepArgs, r.sleepBlock, r.sleepClose = testutils.SleepStub(&r.timeSleep)
 }
 
 func (r *retryTest) setAdvertiseSuccess() {
@@ -162,7 +160,6 @@ func runRetryTest(t *testing.T, f func(r *retryTest)) {
 	r := &retryTest{}
 	defer testutils.SetTimeout(t, time.Second)()
 	r.setup()
-	defer testutils.ResetSleepStub(&timeSleep)
 
 	withSetup(t, func(hypCh *tchannel.Channel, hostPort string) {
 		json.Register(hypCh, json.Handlers{"ad": r.adHandler}, nil)
@@ -179,6 +176,7 @@ func runRetryTest(t *testing.T, f func(r *retryTest)) {
 		r.client, err = NewClient(serverCh, configFor(hostPort), &ClientOptions{
 			Handler:      r,
 			FailStrategy: FailStrategyIgnore,
+			TimeSleep:    r.timeSleep,
 		})
 		require.NoError(t, err, "NewClient")
 		defer r.client.Close()
@@ -402,6 +400,12 @@ func configFor(node string) Configuration {
 	}
 }
 
+func stubbedSleep() *ClientOptions {
+	return &ClientOptions{
+		TimeSleep: func(_ time.Duration) {},
+	}
+}
+
 func withSetup(t *testing.T, f func(ch *tchannel.Channel, hostPort string)) {
 	serverCh, err := tchannel.NewChannel(hyperbahnServiceName, nil)
 	require.NoError(t, err)

hyperbahn/client.go

@@ -86,6 +86,10 @@ type ClientOptions struct {
 	TimeoutPerAttempt time.Duration
 	Handler           Handler
 	FailStrategy      FailStrategy
+
+	// The following are variables for stubbing in unit tests.
+	// They are not part of the stable API and may change.
+	TimeSleep func(d time.Duration)
 }
 
 // NewClient creates a new Hyperbahn client using the given channel.
@@ -105,6 +109,9 @@ func NewClient(ch *tchannel.Channel, config Configuration, opts *ClientOptions)
 	if client.opts.Handler == nil {
 		client.opts.Handler = nullHandler{}
 	}
+	if client.opts.TimeSleep == nil {
+		client.opts.TimeSleep = time.Sleep
+	}
 
 	if err := parseConfig(&config); err != nil {
 		return nil, err