Skip to content

Commit 9712ec5

Browse files
arjan-balarjan-bal
committed
refactor test
1 parent 67f7a1a commit 9712ec5

File tree

1 file changed

+81
-87
lines changed

1 file changed

+81
-87
lines changed

‎balancer/pickfirst/pickfirstleaf/pickfirstleaf_ext_test.go

+81-87
Original file line numberDiff line numberDiff line change
@@ -887,9 +887,9 @@ func (s) TestPickFirstLeaf_HappyEyeballs_TFAfterEndOfList(t *testing.T) {
887887
}()
888888

889889
servers := newHangingServerGroup(t, 3)
890-
defer servers.close()
891890
rb := manual.NewBuilderWithScheme("whatever")
892-
rb.InitialState(resolver.State{Addresses: servers.addrs})
891+
addrs := resolverAddrsFromHangingServers(servers)
892+
rb.InitialState(resolver.State{Addresses: addrs})
893893
cc, err := grpc.NewClient("whatever:///this-gets-overwritten",
894894
grpc.WithTransportCredentials(insecure.NewCredentials()),
895895
grpc.WithDefaultServiceConfig(fmt.Sprintf(`{"loadBalancingConfig": [{"%s":{}}]}`, pickfirstleaf.Name)),
@@ -904,37 +904,37 @@ func (s) TestPickFirstLeaf_HappyEyeballs_TFAfterEndOfList(t *testing.T) {
904904
testutils.AwaitState(ctx, t, cc, connectivity.Connecting)
905905

906906
// Verify that only the first server is contacted.
907-
if err := servers.awaitContacted(ctx, 0); err != nil {
908-
t.Fatalf("Server with address %q not contacted: %v", servers.addrs[0], err)
907+
if err := servers[0].awaitContacted(ctx); err != nil {
908+
t.Fatalf("Server with address %q not contacted: %v", addrs[0], err)
909909
}
910910

911911
// Ensure no other servers are contacted.
912-
if got, want := servers.isContacted(1), false; got != want {
913-
t.Fatalf("Servers.isContacted(%q) = %t, want %t", servers.addrs[1], got, want)
912+
if got, want := servers[1].isContacted(), false; got != want {
913+
t.Fatalf("Server[%q].isContacted() = %t, want %t", addrs[1], got, want)
914914
}
915-
if got, want := servers.isContacted(2), false; got != want {
916-
t.Fatalf("Servers.isContacted(%q) = %t, want %t", servers.addrs[2], got, want)
915+
if got, want := servers[2].isContacted(), false; got != want {
916+
t.Fatalf("Server[%q].isContacted() = %t, want %t", addrs[2], got, want)
917917
}
918918

919919
// Make the happy eyeballs timer fire twice so that pickfirst reaches the
920920
// last address in the list.
921921
timerCh <- struct{}{}
922922

923923
// Verify that the second server is contacted and 3rd isn't.
924-
if err := servers.awaitContacted(ctx, 1); err != nil {
925-
t.Fatalf("Server with address %q not contacted: %v", servers.addrs[1], err)
924+
if err := servers[1].awaitContacted(ctx); err != nil {
925+
t.Fatalf("Server with address %q not contacted: %v", addrs[1], err)
926926
}
927927

928-
if got, want := servers.isContacted(2), false; got != want {
929-
t.Fatalf("Servers.isContacted(%q) = %t, want %t", servers.addrs[2], got, want)
928+
if got, want := servers[2].isContacted(), false; got != want {
929+
t.Fatalf("Server[%q].isContacted() = %t, want %t", addrs[2], got, want)
930930
}
931931
timerCh <- struct{}{}
932-
if err := servers.awaitContacted(ctx, 2); err != nil {
933-
t.Fatalf("Server with address %q not contacted: %v", servers.addrs[2], err)
932+
if err := servers[2].awaitContacted(ctx); err != nil {
933+
t.Fatalf("Server with address %q not contacted: %v", addrs[2], err)
934934
}
935935

936936
// First SubConn Fails.
937-
servers.closeConn(0)
937+
servers[0].closeConn()
938938

939939
// No TF should be reported until the first pass is complete.
940940
shortCtx, shortCancel := context.WithTimeout(ctx, defaultTestShortTimeout)
@@ -947,12 +947,12 @@ func (s) TestPickFirstLeaf_HappyEyeballs_TFAfterEndOfList(t *testing.T) {
947947
timerCh <- struct{}{}
948948

949949
// Third SubConn fails.
950-
servers.closeConn(2)
950+
servers[2].closeConn()
951951

952952
testutils.AwaitNotState(shortCtx, t, cc, connectivity.TransientFailure)
953953

954954
// Last SubConn fails, this should result in a TF update.
955-
servers.closeConn(1)
955+
servers[1].closeConn()
956956
testutils.AwaitState(ctx, t, cc, connectivity.TransientFailure)
957957
}
958958

@@ -988,9 +988,9 @@ func (s) TestPickFirstLeaf_HappyEyeballs_TriggerConnectionDelay(t *testing.T) {
988988
}()
989989

990990
servers := newHangingServerGroup(t, 2)
991-
defer servers.close()
991+
addrs := resolverAddrsFromHangingServers(servers)
992992
rb := manual.NewBuilderWithScheme("whatever")
993-
rb.InitialState(resolver.State{Addresses: servers.addrs})
993+
rb.InitialState(resolver.State{Addresses: addrs})
994994
cc, err := grpc.NewClient("whatever:///this-gets-overwritten",
995995
grpc.WithTransportCredentials(insecure.NewCredentials()),
996996
grpc.WithDefaultServiceConfig(fmt.Sprintf(`{"loadBalancingConfig": [{"%s":{}}]}`, pickfirstleaf.Name)),
@@ -1005,21 +1005,21 @@ func (s) TestPickFirstLeaf_HappyEyeballs_TriggerConnectionDelay(t *testing.T) {
10051005
testutils.AwaitState(ctx, t, cc, connectivity.Connecting)
10061006

10071007
// Verify that the first server is contacted.
1008-
if err := servers.awaitContacted(ctx, 0); err != nil {
1009-
t.Fatalf("Server with address %q not contacted: %v", servers.addrs[0], err)
1008+
if err := servers[0].awaitContacted(ctx); err != nil {
1009+
t.Fatalf("Server with address %q not contacted: %v", addrs[0], err)
10101010
}
10111011

1012-
if got, want := servers.isContacted(1), false; got != want {
1013-
t.Fatalf("Servers.isContacted(%q) = %t, want %t", servers.addrs[1], got, want)
1012+
if got, want := servers[1].isContacted(), false; got != want {
1013+
t.Fatalf("Server[%q].isContacted() = %t, want %t", addrs[1], got, want)
10141014
}
10151015

10161016
timerCh <- struct{}{}
10171017

10181018
// Second connection attempt is successful.
1019-
if err := servers.awaitContacted(ctx, 1); err != nil {
1020-
t.Fatalf("Server with address %q not contacted: %v", servers.addrs[1], err)
1019+
if err := servers[1].awaitContacted(ctx); err != nil {
1020+
t.Fatalf("Server with address %q not contacted: %v", addrs[1], err)
10211021
}
1022-
servers.enterReady(1)
1022+
servers[1].enterReady()
10231023
testutils.AwaitState(ctx, t, cc, connectivity.Ready)
10241024
}
10251025

@@ -1055,9 +1055,9 @@ func (s) TestPickFirstLeaf_HappyEyeballs_TFThenTimerFires(t *testing.T) {
10551055
}()
10561056

10571057
servers := newHangingServerGroup(t, 3)
1058-
defer servers.close()
1058+
addrs := resolverAddrsFromHangingServers(servers)
10591059
rb := manual.NewBuilderWithScheme("whatever")
1060-
rb.InitialState(resolver.State{Addresses: servers.addrs})
1060+
rb.InitialState(resolver.State{Addresses: addrs})
10611061
cc, err := grpc.NewClient("whatever:///this-gets-overwritten",
10621062
grpc.WithTransportCredentials(insecure.NewCredentials()),
10631063
grpc.WithDefaultServiceConfig(fmt.Sprintf(`{"loadBalancingConfig": [{"%s":{}}]}`, pickfirstleaf.Name)),
@@ -1072,16 +1072,16 @@ func (s) TestPickFirstLeaf_HappyEyeballs_TFThenTimerFires(t *testing.T) {
10721072
testutils.AwaitState(ctx, t, cc, connectivity.Connecting)
10731073

10741074
// Verify that only the first server is contacted.
1075-
if err := servers.awaitContacted(ctx, 0); err != nil {
1076-
t.Fatalf("Server with address %q not contacted: %v", servers.addrs[0], err)
1075+
if err := servers[0].awaitContacted(ctx); err != nil {
1076+
t.Fatalf("Server with address %q not contacted: %v", addrs[0], err)
10771077
}
10781078

10791079
// Ensure no other servers are contacted.
1080-
if got, want := servers.isContacted(1), false; got != want {
1081-
t.Fatalf("Servers.isContacted(%q) = %t, want %t", servers.addrs[1], got, want)
1080+
if got, want := servers[1].isContacted(), false; got != want {
1081+
t.Fatalf("Server[%q].isContacted() = %t, want %t", addrs[1], got, want)
10821082
}
1083-
if got, want := servers.isContacted(2), false; got != want {
1084-
t.Fatalf("Servers.isContacted(%q) = %t, want %t", servers.addrs[2], got, want)
1083+
if got, want := servers[2].isContacted(), false; got != want {
1084+
t.Fatalf("Server[%q].isContacted() = %t, want %t", addrs[2], got, want)
10851085
}
10861086

10871087
// First SubConn Fails.
@@ -1090,29 +1090,29 @@ func (s) TestPickFirstLeaf_HappyEyeballs_TFThenTimerFires(t *testing.T) {
10901090
timerMu.Lock()
10911091
timerCh = make(chan struct{})
10921092
timerMu.Unlock()
1093-
servers.closeConn(0)
1093+
servers[0].closeConn()
10941094

10951095
// The second server is contacted.
10961096
// Verify that only the first server is contacted.
1097-
if err := servers.awaitContacted(ctx, 1); err != nil {
1098-
t.Fatalf("Server with address %q not contacted: %v", servers.addrs[1], err)
1097+
if err := servers[1].awaitContacted(ctx); err != nil {
1098+
t.Fatalf("Server with address %q not contacted: %v", addrs[1], err)
10991099
}
11001100

11011101
// Ensure no other servers are contacted.
1102-
if got, want := servers.isContacted(2), false; got != want {
1103-
t.Fatalf("Servers.isContacted(%q) = %t, want %t", servers.addrs[2], got, want)
1102+
if got, want := servers[2].isContacted(), false; got != want {
1103+
t.Fatalf("Server[%q].isContacted() = %t, want %t", addrs[2], got, want)
11041104
}
11051105

11061106
// The happy eyeballs timer expires, skipping server[1] and requesting the creation
11071107
// of a third SubConn.
11081108
timerCh <- struct{}{}
11091109

1110-
if err := servers.awaitContacted(ctx, 2); err != nil {
1111-
t.Fatalf("Server with address %q not contacted: %v", servers.addrs[2], err)
1110+
if err := servers[2].awaitContacted(ctx); err != nil {
1111+
t.Fatalf("Server with address %q not contacted: %v", addrs[2], err)
11121112
}
11131113

11141114
// Second SubConn connects.
1115-
servers.enterReady(1)
1115+
servers[1].enterReady()
11161116
testutils.AwaitState(ctx, t, cc, connectivity.Connecting)
11171117
}
11181118

@@ -1221,37 +1221,27 @@ func (c *ccStateSubscriber) OnMessage(msg any) {
12211221
c.transitions = append(c.transitions, msg.(connectivity.State))
12221222
}
12231223

1224-
// handingServerGroup is a group of servers that accept a TCP connection and
1225-
// remain idle until asked to close the connection. They can be used to control
1224+
// hangingServer is a server that accept a TCP connection and remains idle until
1225+
// asked to close or respond to the connection. They can be used to control
12261226
// how long it takes for a subchannel to report a TRANSIENT_FAILURE in tests.
1227-
type handingServerGroup struct {
1228-
addrs []resolver.Address
1229-
listeners []net.Listener
1230-
serverConnCloseFuncs []func()
1231-
serverContacted []*grpcsync.Event
1232-
readyChans []chan struct{}
1227+
type hangingServer struct {
1228+
addr resolver.Address
1229+
listener net.Listener
1230+
closeConn func()
1231+
contacted *grpcsync.Event
1232+
readyChan chan struct{}
12331233
}
12341234

1235-
func newHangingServerGroup(t *testing.T, count int) *handingServerGroup {
1236-
listeners := []net.Listener{}
1237-
closeFns := []func(){}
1238-
addrs := []resolver.Address{}
1239-
contactedEvents := []*grpcsync.Event{}
1240-
readyChans := []chan struct{}{}
1241-
1235+
func newHangingServerGroup(t *testing.T, count int) []*hangingServer {
1236+
servers := []*hangingServer{}
12421237
for i := 0; i < count; i++ {
12431238
lis, err := net.Listen("tcp", "localhost:0")
12441239
if err != nil {
12451240
t.Fatalf("Error while listening. Err: %v", err)
12461241
}
1247-
listeners = append(listeners, lis)
1248-
addrs = append(addrs, resolver.Address{Addr: lis.Addr().String()})
12491242
closeChan := make(chan struct{})
1250-
closeFns = append(closeFns, sync.OnceFunc(func() { close(closeChan) }))
12511243
contacted := grpcsync.NewEvent()
1252-
contactedEvents = append(contactedEvents, contacted)
12531244
readyChan := make(chan struct{})
1254-
readyChans = append(readyChans, readyChan)
12551245

12561246
go func() {
12571247
conn, err := lis.Accept()
@@ -1277,43 +1267,47 @@ func newHangingServerGroup(t *testing.T, count int) *handingServerGroup {
12771267
}
12781268
}
12791269
}()
1280-
}
1281-
1282-
return &handingServerGroup{
1283-
addrs: addrs,
1284-
listeners: listeners,
1285-
serverConnCloseFuncs: closeFns,
1286-
serverContacted: contactedEvents,
1287-
readyChans: readyChans,
1288-
}
1289-
}
1270+
server := &hangingServer{
1271+
addr: resolver.Address{Addr: lis.Addr().String()},
1272+
listener: lis,
1273+
closeConn: grpcsync.OnceFunc(func() {
1274+
close(closeChan)
1275+
}),
1276+
contacted: contacted,
1277+
readyChan: readyChan,
1278+
}
1279+
servers = append(servers, server)
12901280

1291-
func (hg *handingServerGroup) close() {
1292-
for _, fn := range hg.serverConnCloseFuncs {
1293-
fn()
1294-
}
1295-
for _, l := range hg.listeners {
1296-
l.Close()
1281+
t.Cleanup(func() {
1282+
server.closeConn()
1283+
server.listener.Close()
1284+
})
12971285
}
1298-
}
12991286

1300-
func (hg *handingServerGroup) closeConn(serverIdx int) {
1301-
hg.serverConnCloseFuncs[serverIdx]()
1287+
return servers
13021288
}
13031289

1304-
func (hg *handingServerGroup) isContacted(serverIdx int) bool {
1305-
return hg.serverContacted[serverIdx].HasFired()
1290+
func (s *hangingServer) isContacted() bool {
1291+
return s.contacted.HasFired()
13061292
}
13071293

1308-
func (hg *handingServerGroup) enterReady(serverIdx int) {
1309-
hg.readyChans[serverIdx] <- struct{}{}
1294+
func (s *hangingServer) enterReady() {
1295+
s.readyChan <- struct{}{}
13101296
}
13111297

1312-
func (hg *handingServerGroup) awaitContacted(ctx context.Context, serverIdx int) error {
1298+
func (s *hangingServer) awaitContacted(ctx context.Context) error {
13131299
select {
13141300
case <-ctx.Done():
13151301
return ctx.Err()
1316-
case <-hg.serverContacted[serverIdx].Done():
1302+
case <-s.contacted.Done():
13171303
}
13181304
return nil
13191305
}
1306+
1307+
func resolverAddrsFromHangingServers(servers []*hangingServer) []resolver.Address {
1308+
addrs := []resolver.Address{}
1309+
for _, srv := range servers {
1310+
addrs = append(addrs, srv.addr)
1311+
}
1312+
return addrs
1313+
}

0 commit comments

Comments
 (0)