From 3abe4e43d373fb4dfa784f30abe73c81ce57ae0d Mon Sep 17 00:00:00 2001
From: Daniel Nephin <dnephin@hashicorp.com>
Date: Thu, 13 Aug 2020 17:09:37 -0400
Subject: [PATCH] testing: use t.Cleanup in TestAgent for returnPorts

---
 agent/testagent.go | 24 ++++--------------------
 1 file changed, 4 insertions(+), 20 deletions(-)

diff --git a/agent/testagent.go b/agent/testagent.go
index 387c3efbe..74570a368 100644
--- a/agent/testagent.go
+++ b/agent/testagent.go
@@ -58,12 +58,9 @@ type TestAgent struct {
 	// the callers responsibility to clean up the data directory.
 	// Otherwise, a temporary data directory is created and removed
 	// when Shutdown() is called.
+	// TODO:
 	Config *config.RuntimeConfig
 
-	// returnPortsFn will put the ports claimed for the test back into the
-	// general freeport pool
-	returnPortsFn func()
-
 	// LogOutput is the sink for the logs. If nil, logs are written
 	// to os.Stderr.
 	LogOutput io.Writer
@@ -71,9 +68,11 @@ type TestAgent struct {
 	// DataDir is the data directory which is used when Config.DataDir
 	// is not set. It is created automatically and removed when
 	// Shutdown() is called.
+	// TODO:
 	DataDir string
 
 	// Key is the optional encryption key for the LAN and WAN keyring.
+	// TODO:
 	Key string
 
 	// UseTLS, if true, will disable the HTTP port and enable the HTTPS
@@ -198,7 +197,7 @@ func (a *TestAgent) Start(t *testing.T) (err error) {
 	})
 
 	portsConfig, returnPortsFn := randomPortsSource(a.UseTLS)
-	a.returnPortsFn = returnPortsFn
+	t.Cleanup(returnPortsFn)
 
 	nodeID := NodeID()
 
@@ -221,13 +220,6 @@ func (a *TestAgent) Start(t *testing.T) (err error) {
 		),
 	}
 
-	defer func() {
-		if err != nil && a.returnPortsFn != nil {
-			a.returnPortsFn()
-			a.returnPortsFn = nil
-		}
-	}()
-
 	// write the keyring
 	if a.Key != "" {
 		writeKey := func(key, filename string) error {
@@ -357,14 +349,6 @@ func (a *TestAgent) Shutdown() error {
 		return nil
 	}
 
-	// Return ports last of all
-	defer func() {
-		if a.returnPortsFn != nil {
-			a.returnPortsFn()
-			a.returnPortsFn = nil
-		}
-	}()
-
 	// shutdown agent before endpoints
 	defer a.Agent.ShutdownEndpoints()
 	if err := a.Agent.ShutdownAgent(); err != nil {