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

Code before:
	cc := testutils.NewBalancerClientConn(t)
	bg := New(Options{
		CC:              cc,
		BuildOpts:       bOpts,
		StateAggregator: nil,
		Logger:          nil,
	})

	// Add the stub balancer build above as a child policy.
	balancerBuilder := balancer.Get(balancerName)
	bg.Add(testBalancerIDs[0], balancerBuilder)

	// Send an empty clientConn state change. This should trigger the
	// verification of the buildOptions being passed to the child policy.
	if err := bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{}); err != nil {
		t.Fatal(err)
	}
}

func (s) TestBalancerGroup_UpdateClientConnState_AfterClose(t *testing.T) {
	balancerName := "stub-balancer-test-update-client-state-after-close"
	called := false

	stub.Register(balancerName, stub.BalancerFuncs{
		UpdateClientConnState: func(_ *stub.BalancerData, _ balancer.ClientConnState) error {
			called = true
			return nil
		},
	})

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

	bg.Add(testBalancerIDs[0], balancer.Get(balancerName))
	bg.Close()

	err := bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{})
	if err != nil {
		t.Fatalf("Expected nil error, got %v", err)
	}
	if called {
		t.Fatalf("UpdateClientConnState was called after BalancerGroup was closed")
	}
}

func (s) TestBalancerGroup_ResolverError_AfterClose(t *testing.T) {
	balancerName := "stub-balancer-test-resolver-error-after-close"

Review comment (suggestion):
Change to use channel here too. In the check add a `defaultShortTestTimeout` case along with wait for channel in select block.

Code after:
	cc := testutils.NewBalancerClientConn(t)
	bg := New(Options{
		CC:              cc,
		BuildOpts:       bOpts,
		StateAggregator: nil,
		Logger:          nil,
	})

	// Add the stub balancer build above as a child policy.
	balancerBuilder := balancer.Get(balancerName)
	bg.Add(testBalancerIDs[0], balancerBuilder)

	// Send an empty clientConn state change. This should trigger the
	// verification of the buildOptions being passed to the child policy.
	if err := bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{}); err != nil {
		t.Fatal(err)
	}
}

func (s) TestBalancerGroup_UpdateClientConnState_AfterClose(t *testing.T) {
	balancerName := t.Name()
	clientConnStateCh := make(chan struct{}, 1)

	stub.Register(balancerName, stub.BalancerFuncs{
		UpdateClientConnState: func(_ *stub.BalancerData, _ balancer.ClientConnState) error {
			clientConnStateCh <- struct{}{}
			return nil
		},
	})

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

	bg.Add(testBalancerIDs[0], balancer.Get(balancerName))
	bg.Close()

	if err := bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{}); err != nil {
		t.Fatalf("Expected nil error, got %v", err)
	}

	select {
	case <-clientConnStateCh:
		t.Fatalf("UpdateClientConnState was called after BalancerGroup was closed")
	case <-time.After(defaultTestShortTimeout):
	}
}