From e41db4e1d9b44493ef111512e405277aced80475 Mon Sep 17 00:00:00 2001
From: Michael Li <michael.li@hashicorp.com>
Date: Thu, 18 Apr 2024 18:10:04 +0000
Subject: [PATCH 1/3] backport of commit
 dd3a0738771f25b4ac4a1f6e9943bb150240de2e

---
 internal/daemon/worker/controller_connection_test.go | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/internal/daemon/worker/controller_connection_test.go b/internal/daemon/worker/controller_connection_test.go
index d90e35b180..da11d848e2 100644
--- a/internal/daemon/worker/controller_connection_test.go
+++ b/internal/daemon/worker/controller_connection_test.go
@@ -12,7 +12,6 @@ import (
 	"testing"
 	"time"
 
-	opsservices "github.com/hashicorp/boundary/internal/gen/ops/services"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/net/nettest"
@@ -60,10 +59,9 @@ func TestMonitorUpstreamConnectionState(t *testing.T) {
 	})
 
 	tests := []struct {
-		name             string
-		expectedResponse *opsservices.GetHealthResponse
-		addresses        []string
-		expectedState    connectivity.State
+		name          string
+		addresses     []string
+		expectedState connectivity.State
 	}{
 		{
 			name:          "connection with 1 good address",
@@ -100,9 +98,6 @@ func TestMonitorUpstreamConnectionState(t *testing.T) {
 			case <-time.After(2 * time.Second):
 				t.Error("Time out waiting for condition")
 			}
-
-			got := upstreamConnectionState.Load()
-			assert.Equal(t, tt.expectedState, got)
 		})
 	}
 }

From 6e08be789be9f69936b314d0a75104fe89bcb6e5 Mon Sep 17 00:00:00 2001
From: Elim Tsiagbey <elim.tsiagbey@hashicorp.com>
Date: Thu, 18 Apr 2024 20:14:09 +0000
Subject: [PATCH 2/3] backport of commit
 14198b6e49a31ed5bcf09fc8d1258339e43d645b

---
 .../worker/controller_connection_test.go      | 45 +++++++++----------
 1 file changed, 20 insertions(+), 25 deletions(-)

diff --git a/internal/daemon/worker/controller_connection_test.go b/internal/daemon/worker/controller_connection_test.go
index da11d848e2..ba07127629 100644
--- a/internal/daemon/worker/controller_connection_test.go
+++ b/internal/daemon/worker/controller_connection_test.go
@@ -21,38 +21,13 @@ import (
 )
 
 func TestMonitorUpstreamConnectionState(t *testing.T) {
-	ctx := context.Background()
-	stateCtx, cancelStateCtx := context.WithCancel(ctx)
-
-	upstreamConnectionState := new(atomic.Value)
-
 	servers, err := createTestServers(t)
 	require.NoError(t, err)
 
 	scheme := strconv.FormatInt(time.Now().UnixNano(), 36)
 	res := manual.NewBuilderWithScheme(scheme)
-	grpcResolver := &grpcResolverReceiver{res}
-
-	dialOpts := createDefaultGRPCDialOptions(res, nil)
-
-	cc, err := grpc.Dial(
-		fmt.Sprintf("%s:///%s", res.Scheme(), servers[0].address),
-		dialOpts...,
-	)
-
-	require.NoError(t, err)
-
-	// track GRPC state changes
-	go monitorUpstreamConnectionState(stateCtx, cc, upstreamConnectionState)
-
-	grpcResolver.InitialAddresses([]string{servers[0].address})
 
 	t.Cleanup(func() {
-		cc.Close()
-		cancelStateCtx()
-
-		assert.Equal(t, connectivity.Shutdown, cc.GetState())
-
 		for _, s := range servers {
 			s.srv.GracefulStop()
 		}
@@ -87,7 +62,23 @@ func TestMonitorUpstreamConnectionState(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			stateCtx, cancelStateCtx := context.WithCancel(context.Background())
+			upstreamConnectionState := new(atomic.Value)
 			doneWait := make(chan struct{})
+
+			grpcResolver := &grpcResolverReceiver{res}
+			grpcResolver.InitialAddresses([]string{servers[0].address})
+
+			dialOpts := createDefaultGRPCDialOptions(res, nil)
+			cc, err := grpc.Dial(
+				fmt.Sprintf("%s:///%s", res.Scheme(), servers[0].address),
+				dialOpts...,
+			)
+			require.NoError(t, err)
+
+			// track GRPC state changes
+			go monitorUpstreamConnectionState(stateCtx, cc, upstreamConnectionState)
+
 			grpcResolver.SetAddresses(tt.addresses)
 
 			go waitForConnectionStateCondition(upstreamConnectionState, tt.expectedState, doneWait)
@@ -98,6 +89,10 @@ func TestMonitorUpstreamConnectionState(t *testing.T) {
 			case <-time.After(2 * time.Second):
 				t.Error("Time out waiting for condition")
 			}
+
+			require.NoError(t, cc.Close())
+			cancelStateCtx()
+			assert.Equal(t, connectivity.Shutdown, cc.GetState())
 		})
 	}
 }

From 69f602c2270f92025a503904c999b0ae14721113 Mon Sep 17 00:00:00 2001
From: Elim Tsiagbey <elim.tsiagbey@hashicorp.com>
Date: Thu, 18 Apr 2024 20:26:15 +0000
Subject: [PATCH 3/3] backport of commit
 c1e334169d94bf96f068afc8f81e3efb3ace5b1b

---
 internal/daemon/worker/controller_connection_test.go | 2 --
 1 file changed, 2 deletions(-)

diff --git a/internal/daemon/worker/controller_connection_test.go b/internal/daemon/worker/controller_connection_test.go
index ba07127629..f03557f366 100644
--- a/internal/daemon/worker/controller_connection_test.go
+++ b/internal/daemon/worker/controller_connection_test.go
@@ -12,7 +12,6 @@ import (
 	"testing"
 	"time"
 
-	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/net/nettest"
 	"google.golang.org/grpc"
@@ -92,7 +91,6 @@ func TestMonitorUpstreamConnectionState(t *testing.T) {
 
 			require.NoError(t, cc.Close())
 			cancelStateCtx()
-			assert.Equal(t, connectivity.Shutdown, cc.GetState())
 		})
 	}
 }