tests: stop using UpdateSubConnState (#6527)

grpc · Aug 9, 2023 · 8f51ca8 · 8f51ca8
1 parent cea77bb
commit 8f51ca8
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 50 deletions.
diff --git a/clientconn_test.go b/clientconn_test.go
@@ -1146,7 +1146,9 @@ type stateRecordingBalancer struct {
 	balancer.Balancer
 }
 
-func (b *stateRecordingBalancer) UpdateSubConnState(sc balancer.SubConn, s balancer.SubConnState) {}
+func (b *stateRecordingBalancer) UpdateSubConnState(sc balancer.SubConn, s balancer.SubConnState) {
+	panic(fmt.Sprintf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, s))
+}
 
 func (b *stateRecordingBalancer) Close() {
 	b.Balancer.Close()

diff --git a/test/balancer_switching_test.go b/test/balancer_switching_test.go
@@ -483,10 +483,6 @@ func (s) TestBalancerSwitch_Graceful(t *testing.T) {
 			}()
 			return nil
 		},
-		UpdateSubConnState: func(bd *stub.BalancerData, sc balancer.SubConn, state balancer.SubConnState) {
-			bal := bd.Data.(balancer.Balancer)
-			bal.UpdateSubConnState(sc, state)
-		},
 	})
 
 	// Push a resolver update with the service config specifying our stub

diff --git a/test/balancer_test.go b/test/balancer_test.go
@@ -87,6 +87,7 @@ func (b *testBalancer) UpdateClientConnState(state balancer.ClientConnState) err
 	// Only create a subconn at the first time.
 	if b.sc == nil {
 		var err error
+		b.newSubConnOptions.StateListener = b.updateSubConnState
 		b.sc, err = b.cc.NewSubConn(state.ResolverState.Addresses, b.newSubConnOptions)
 		if err != nil {
 			logger.Errorf("testBalancer: failed to NewSubConn: %v", err)
@@ -99,21 +100,17 @@ func (b *testBalancer) UpdateClientConnState(state balancer.ClientConnState) err
 }
 
 func (b *testBalancer) UpdateSubConnState(sc balancer.SubConn, s balancer.SubConnState) {
-	logger.Infof("testBalancer: UpdateSubConnState: %p, %v", sc, s)
-	if b.sc != sc {
-		logger.Infof("testBalancer: ignored state change because sc is not recognized")
-		return
-	}
-	if s.ConnectivityState == connectivity.Shutdown {
-		b.sc = nil
-		return
-	}
+	panic(fmt.Sprintf("UpdateSubConnState(%v, %+v) called unexpectedly", sc, s))
+}
+
+func (b *testBalancer) updateSubConnState(s balancer.SubConnState) {
+	logger.Infof("testBalancer: updateSubConnState: %v", s)
 
 	switch s.ConnectivityState {
 	case connectivity.Ready:
-		b.cc.UpdateState(balancer.State{ConnectivityState: s.ConnectivityState, Picker: &picker{sc: sc, bal: b}})
+		b.cc.UpdateState(balancer.State{ConnectivityState: s.ConnectivityState, Picker: &picker{bal: b}})
 	case connectivity.Idle:
-		b.cc.UpdateState(balancer.State{ConnectivityState: s.ConnectivityState, Picker: &picker{sc: sc, bal: b, idle: true}})
+		b.cc.UpdateState(balancer.State{ConnectivityState: s.ConnectivityState, Picker: &picker{bal: b, idle: true}})
 	case connectivity.Connecting:
 		b.cc.UpdateState(balancer.State{ConnectivityState: s.ConnectivityState, Picker: &picker{err: balancer.ErrNoSubConnAvailable, bal: b}})
 	case connectivity.TransientFailure:
@@ -127,7 +124,6 @@ func (b *testBalancer) ExitIdle() {}
 
 type picker struct {
 	err  error
-	sc   balancer.SubConn
 	bal  *testBalancer
 	idle bool
 }
@@ -137,14 +133,14 @@ func (p *picker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
 		return balancer.PickResult{}, p.err
 	}
 	if p.idle {
-		p.sc.Connect()
+		p.bal.sc.Connect()
 		return balancer.PickResult{}, balancer.ErrNoSubConnAvailable
 	}
 	extraMD, _ := grpcutil.ExtraMetadata(info.Ctx)
 	info.Ctx = nil // Do not validate context.
 	p.bal.pickInfos = append(p.bal.pickInfos, info)
 	p.bal.pickExtraMDs = append(p.bal.pickExtraMDs, extraMD)
-	return balancer.PickResult{SubConn: p.sc, Done: func(d balancer.DoneInfo) { p.bal.doneInfo = append(p.bal.doneInfo, d) }}, nil
+	return balancer.PickResult{SubConn: p.bal.sc, Done: func(d balancer.DoneInfo) { p.bal.doneInfo = append(p.bal.doneInfo, d) }}, nil
 }
 
 func (s) TestCredsBundleFromBalancer(t *testing.T) {
@@ -397,16 +393,18 @@ func (s) TestAddressAttributesInNewSubConn(t *testing.T) {
 			// Only use the first address.
 			attr := attributes.New(testAttrKey, testAttrVal)
 			addrs[0].Attributes = attr
-			sc, err := bd.ClientConn.NewSubConn([]resolver.Address{addrs[0]}, balancer.NewSubConnOptions{})
+			var sc balancer.SubConn
+			sc, err := bd.ClientConn.NewSubConn([]resolver.Address{addrs[0]}, balancer.NewSubConnOptions{
+				StateListener: func(state balancer.SubConnState) {
+					bd.ClientConn.UpdateState(balancer.State{ConnectivityState: state.ConnectivityState, Picker: &aiPicker{result: balancer.PickResult{SubConn: sc}, err: state.ConnectionError}})
+				},
+			})
 			if err != nil {
 				return err
 			}
 			sc.Connect()
 			return nil
 		},
-		UpdateSubConnState: func(bd *stub.BalancerData, sc balancer.SubConn, state balancer.SubConnState) {
-			bd.ClientConn.UpdateState(balancer.State{ConnectivityState: state.ConnectivityState, Picker: &aiPicker{result: balancer.PickResult{SubConn: sc}, err: state.ConnectionError}})
-		},
 	}
 	stub.Register(attrBalancerName, bf)
 	t.Logf("Registered balancer %s...", attrBalancerName)
@@ -483,18 +481,20 @@ func (s) TestMetadataInAddressAttributes(t *testing.T) {
 				return nil
 			}
 			// Only use the first address.
+			var sc balancer.SubConn
 			sc, err := bd.ClientConn.NewSubConn([]resolver.Address{
 				imetadata.Set(addrs[0], metadata.Pairs(testMDKey, testMDValue)),
-			}, balancer.NewSubConnOptions{})
+			}, balancer.NewSubConnOptions{
+				StateListener: func(state balancer.SubConnState) {
+					bd.ClientConn.UpdateState(balancer.State{ConnectivityState: state.ConnectivityState, Picker: &aiPicker{result: balancer.PickResult{SubConn: sc}, err: state.ConnectionError}})
+				},
+			})
 			if err != nil {
 				return err
 			}
 			sc.Connect()
 			return nil
 		},
-		UpdateSubConnState: func(bd *stub.BalancerData, sc balancer.SubConn, state balancer.SubConnState) {
-			bd.ClientConn.UpdateState(balancer.State{ConnectivityState: state.ConnectivityState, Picker: &aiPicker{result: balancer.PickResult{SubConn: sc}, err: state.ConnectionError}})
-		},
 	}
 	stub.Register(mdBalancerName, bf)
 	t.Logf("Registered balancer %s...", mdBalancerName)
@@ -717,16 +717,18 @@ func (s) TestAuthorityInBuildOptions(t *testing.T) {
 					}
 
 					// Only use the first address.
-					sc, err := bd.ClientConn.NewSubConn([]resolver.Address{addrs[0]}, balancer.NewSubConnOptions{})
+					var sc balancer.SubConn
+					sc, err := bd.ClientConn.NewSubConn([]resolver.Address{addrs[0]}, balancer.NewSubConnOptions{
+						StateListener: func(state balancer.SubConnState) {
+							bd.ClientConn.UpdateState(balancer.State{ConnectivityState: state.ConnectivityState, Picker: &aiPicker{result: balancer.PickResult{SubConn: sc}, err: state.ConnectionError}})
+						},
+					})
 					if err != nil {
 						return err
 					}
 					sc.Connect()
 					return nil
 				},
-				UpdateSubConnState: func(bd *stub.BalancerData, sc balancer.SubConn, state balancer.SubConnState) {
-					bd.ClientConn.UpdateState(balancer.State{ConnectivityState: state.ConnectivityState, Picker: &aiPicker{result: balancer.PickResult{SubConn: sc}, err: state.ConnectionError}})
-				},
 			}
 			balancerName := "stub-balancer-" + test.name
 			stub.Register(balancerName, bf)

diff --git a/test/resolver_update_test.go b/test/resolver_update_test.go
@@ -177,10 +177,6 @@ func (s) TestResolverUpdate_InvalidServiceConfigAfterGoodUpdate(t *testing.T) {
 			ccs.BalancerConfig = nil
 			return bal.UpdateClientConnState(ccs)
 		},
-		UpdateSubConnState: func(bd *stub.BalancerData, sc balancer.SubConn, state balancer.SubConnState) {
-			bal := bd.Data.(balancer.Balancer)
-			bal.UpdateSubConnState(sc, state)
-		},
 	})
 
 	// Start a backend exposing the test service.

diff --git a/test/subconn_test.go b/test/subconn_test.go
@@ -59,7 +59,22 @@ func (s) TestSubConnEmpty(t *testing.T) {
 		UpdateClientConnState: func(d *stub.BalancerData, ccs balancer.ClientConnState) error {
 			if sc == nil {
 				var err error
-				sc, err = d.ClientConn.NewSubConn(ccs.ResolverState.Addresses, balancer.NewSubConnOptions{})
+				sc, err = d.ClientConn.NewSubConn(ccs.ResolverState.Addresses, balancer.NewSubConnOptions{
+					StateListener: func(state balancer.SubConnState) {
+						switch state.ConnectivityState {
+						case connectivity.Ready:
+							d.ClientConn.UpdateState(balancer.State{
+								ConnectivityState: connectivity.Ready,
+								Picker:            &tsccPicker{sc: sc},
+							})
+						case connectivity.TransientFailure:
+							d.ClientConn.UpdateState(balancer.State{
+								ConnectivityState: connectivity.TransientFailure,
+								Picker:            base.NewErrPicker(fmt.Errorf("error connecting: %v", state.ConnectionError)),
+							})
+						}
+					},
+				})
 				if err != nil {
 					t.Errorf("error creating initial subconn: %v", err)
 				}
@@ -81,20 +96,6 @@ func (s) TestSubConnEmpty(t *testing.T) {
 			}
 			return nil
 		},
-		UpdateSubConnState: func(d *stub.BalancerData, sc balancer.SubConn, scs balancer.SubConnState) {
-			switch scs.ConnectivityState {
-			case connectivity.Ready:
-				d.ClientConn.UpdateState(balancer.State{
-					ConnectivityState: connectivity.Ready,
-					Picker:            &tsccPicker{sc: sc},
-				})
-			case connectivity.TransientFailure:
-				d.ClientConn.UpdateState(balancer.State{
-					ConnectivityState: connectivity.TransientFailure,
-					Picker:            base.NewErrPicker(fmt.Errorf("error connecting: %v", scs.ConnectionError)),
-				})
-			}
-		},
 	}
 	stub.Register("tscc", bal)