From bef3a4c6afb42ebdc0dacecda42a33df35514fdd Mon Sep 17 00:00:00 2001
From: Dave Collins <davec@conformal.com>
Date: Mon, 4 May 2020 18:27:00 -0500
Subject: [PATCH] connmgr: Tighten TestNetworkFailure.

This reworks the TestNetworkFailure test to more accurately test the
intended condition by signaling once the max number of failed attempts
has been reached prior to shutting down the connection manager and
checking the expected number of total dials attempted.

It also tightens the number of expected dials to a more accurate value
to better ensure an accurate result.
---
 connmgr/connmanager_test.go | 37 +++++++++++++++++++++++++++++--------
 1 file changed, 29 insertions(+), 8 deletions(-)

diff --git a/connmgr/connmanager_test.go b/connmgr/connmanager_test.go
index cf87bf82..d2756993 100644
--- a/connmgr/connmanager_test.go
+++ b/connmgr/connmanager_test.go
@@ -444,14 +444,23 @@ func TestMaxRetryDuration(t *testing.T) {
 // TestNetworkFailure tests that the connection manager handles a network
 // failure gracefully.
 func TestNetworkFailure(t *testing.T) {
+	var closeOnce sync.Once
+	const targetOutbound = 5
+	const retryTimeout = time.Millisecond * 5
 	var dials uint32
+	reachedMaxFailedAttempts := make(chan struct{})
+	connMgrDone := make(chan struct{})
 	errDialer := func(ctx context.Context, network, addr string) (net.Conn, error) {
-		atomic.AddUint32(&dials, 1)
+		totalDials := atomic.AddUint32(&dials, 1)
+		if totalDials >= maxFailedAttempts {
+			closeOnce.Do(func() { close(reachedMaxFailedAttempts) })
+			<-connMgrDone
+		}
 		return nil, errors.New("network down")
 	}
 	cmgr, err := New(&Config{
-		TargetOutbound: 5,
-		RetryDuration:  5 * time.Millisecond,
+		TargetOutbound: targetOutbound,
+		RetryDuration:  retryTimeout,
 		Dial:           errDialer,
 		GetNewAddress: func() (net.Addr, error) {
 			return &net.TCPAddr{
@@ -467,12 +476,24 @@ func TestNetworkFailure(t *testing.T) {
 		t.Fatalf("New error: %v", err)
 	}
 	cmgr.Start()
-	time.AfterFunc(10*time.Millisecond, cmgr.Stop)
+
+	// Shutdown the connection manager after the max failed attempts is reached
+	// and an additional retry duration has passed and then wait for the
+	// shutdown to complete.
+	<-reachedMaxFailedAttempts
+	time.Sleep(retryTimeout)
+	cmgr.Stop()
+	close(connMgrDone)
 	cmgr.Wait()
-	wantMaxDials := uint32(75)
-	if atomic.LoadUint32(&dials) > wantMaxDials {
-		t.Fatalf("network failure: unexpected number of dials - got %v, want < %v",
-			atomic.LoadUint32(&dials), wantMaxDials)
+
+	// Ensure the number of dial attempts does not exceed the max number of
+	// failed attempts plus the number of potential retries during the
+	// additional waiting period.
+	gotDials := atomic.LoadUint32(&dials)
+	wantMaxDials := uint32(maxFailedAttempts + targetOutbound)
+	if gotDials > wantMaxDials {
+		t.Fatalf("unexpected number of dials - got %v, want <= %v", gotDials,
+			wantMaxDials)
 	}
 }