Repository: grpc/grpc-go
File: internal/balancergroup/balancergroup_test.go (Go)

Code before:
				delete(scs, sc)
				continue
			} else {
				t.Fatalf("Shutdown called for wrong SubConn %v, want in %v", sc, scs)
			}
		}
	}
}

func (s) TestBalancerExitIdle_All(t *testing.T) {
	balancerOne := "stub-balancer-test-balancer-group-exit-idle-one"
	balancerTwo := "stub-balancer-test-balancer-group-exit-idle-two"
	balancerThree := "stub-balancer-test-balancer-group-exit-idle-three"

	balancerNames := []string{balancerOne, balancerTwo, balancerThree}
	testIDs := []string{testBalancerIDs[0], testBalancerIDs[1], testBalancerIDs[2]}

	exitIdleCh := make(chan string, len(balancerNames))

	for _, name := range balancerNames {
		stub.Register(name, stub.BalancerFuncs{
			ExitIdle: func(_ *stub.BalancerData) {
				exitIdleCh <- name
			},
		})
	}

	cc := testutils.NewBalancerClientConn(t)
	bg := New(Options{
		CC:              cc,
		BuildOpts:       balancer.BuildOptions{},
		StateAggregator: nil,
		Logger:          nil,
	})
	defer bg.Close()

	bg.Add(testIDs[0], balancer.Get(balancerOne))
	bg.Add(testIDs[1], balancer.Get(balancerTwo))
	bg.Add(testIDs[2], balancer.Get(balancerThree))

	bg.ExitIdle()

	called := make(map[string]bool)
	for i := 0; i < len(balancerNames); i++ {
		select {
		case name := <-exitIdleCh:
			if called[name] {
				t.Fatalf("ExitIdle was called multiple times for sub-balancer %q", name)
			}
			called[name] = true
		case <-time.After(defaultTestTimeout):

Review comment (suggestion):
I would get rid of the slices and the loops and instead keep this downright simple and dump.

```golang
  // Have two test IDs: testID1 and testID2
  // Have two channels, exitIdleCh1 and exitIdleCh2
  // Register two stub balancers, the first one for balancer1 that writes to exitIdleCh1 and the second one for balancer2 that writes to exitIdleCh2

  // Create the balancer group, and add the child balancers in there

  // Call ExitIdle

  // Validate that both exitIdleCh1 and exitIdleCh2 are written to. (I don't think it's an issue if one of those are written to more than once). In this case though, you might have to spawn goroutines to validate that both channels are written to, since we cannot guarantee the order in which they would be written to.
```

Code after:
			// SubConns, and both should be deleted.
			if ok := scs[sc]; ok {
				delete(scs, sc)
				continue
			} else {
				t.Fatalf("Shutdown called for wrong SubConn %v, want in %v", sc, scs)
			}
		}
	}
}

func (s) TestBalancerExitIdle_All(t *testing.T) {
	balancer1 := t.Name() + "-1"
	balancer2 := t.Name() + "-2"

	testID1, testID2 := testBalancerIDs[0], testBalancerIDs[1]

	exitIdleCh1, exitIdleCh2 := make(chan struct{}, 1), make(chan struct{}, 1)

	stub.Register(balancer1, stub.BalancerFuncs{
		ExitIdle: func(_ *stub.BalancerData) {
			exitIdleCh1 <- struct{}{}
		},
	})

	stub.Register(balancer2, stub.BalancerFuncs{
		ExitIdle: func(_ *stub.BalancerData) {
			exitIdleCh2 <- struct{}{}
		},
	})

	cc := testutils.NewBalancerClientConn(t)
	bg := New(Options{
		CC:              cc,
		BuildOpts:       balancer.BuildOptions{},
		StateAggregator: nil,
		Logger:          nil,
	})
	defer bg.Close()

	bg.Add(testID1, balancer.Get(balancer1))
	bg.Add(testID2, balancer.Get(balancer2))

	bg.ExitIdle()

	errCh := make(chan error, 2)

	go func() {
		select {
		case <-exitIdleCh1:
			errCh <- nil