Fix tests

arjan-bal · Aug 9, 2024 · f1daad1 · f1daad1
1 parent 17c63d2
commit f1daad1
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/balancer/rls/balancer_test.go b/balancer/rls/balancer_test.go
@@ -923,9 +923,9 @@ func (s) TestUpdateStatePauses(t *testing.T) {
 		Init: func(bd *stub.BalancerData) {
 			bd.Data = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
 		},
-        Close: func(bd *stub.BalancerData) {
-            bd.Data.(balancer.Balancer).Close()
-        },
+		Close: func(bd *stub.BalancerData) {
+			bd.Data.(balancer.Balancer).Close()
+		},
 		ParseConfig: func(sc json.RawMessage) (serviceconfig.LoadBalancingConfig, error) {
 			cfg := &childPolicyConfig{}
 			if err := json.Unmarshal(sc, cfg); err != nil {

diff --git a/xds/internal/balancer/clustermanager/clustermanager_test.go b/xds/internal/balancer/clustermanager/clustermanager_test.go
@@ -607,6 +607,7 @@ func TestClusterGracefulSwitch(t *testing.T) {
 	builder := balancer.Get(balancerName)
 	parser := builder.(balancer.ConfigParser)
 	bal := builder.Build(cc, balancer.BuildOptions{})
+	defer bal.Close()
 
 	configJSON1 := `{
 "children": {
@@ -644,6 +645,9 @@ func TestClusterGracefulSwitch(t *testing.T) {
 		Init: func(bd *stub.BalancerData) {
 			bd.Data = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
 		},
+		Close: func(bd *stub.BalancerData) {
+			bd.Data.(balancer.Balancer).Close()
+		},
 		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
 			bal := bd.Data.(balancer.Balancer)
 			return bal.UpdateClientConnState(ccs)
@@ -730,6 +734,7 @@ func (s) TestUpdateStatePauses(t *testing.T) {
 	builder := balancer.Get(balancerName)
 	parser := builder.(balancer.ConfigParser)
 	bal := builder.Build(cc, balancer.BuildOptions{})
+	defer bal.Close()
 
 	configJSON1 := `{
 "children": {