snapshot of https://github.com/grpc/grpc-go.git (master, 2026-09-15; blobs > 131072B omitted)

internal/balancergroup/balancergroup_test.go

@@ -0,0 +1,930 @@
+/*
+ * Copyright 2019 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package balancergroup_test
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"strings"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/pickfirst"
+	"google.golang.org/grpc/balancer/roundrobin"
+	"google.golang.org/grpc/balancer/weightedtarget/weightedaggregator"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/internal/balancer/stub"
+	"google.golang.org/grpc/internal/balancergroup"
+	"google.golang.org/grpc/internal/channelz"
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/resolver"
+)
+
+const (
+	defaultTestTimeout      = 5 * time.Second
+	defaultTestShortTimeout = 100 * time.Millisecond
+)
+
+var (
+	rrBuilder            = balancer.Get(roundrobin.Name)
+	testBalancerIDs      = []string{"b1", "b2", "b3"}
+	testBackendAddrs     []resolver.Address
+	testBackendEndpoints []resolver.Endpoint
+)
+
+const testBackendAddrsCount = 12
+
+func init() {
+	for i := 0; i < testBackendAddrsCount; i++ {
+		addr := resolver.Address{Addr: fmt.Sprintf("%d.%d.%d.%d:%d", i, i, i, i, i)}
+		testBackendAddrs = append(testBackendAddrs, addr)
+		testBackendEndpoints = append(testBackendEndpoints, resolver.Endpoint{Addresses: []resolver.Address{addr}})
+	}
+}
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+// Create a new balancer group, add balancer and backends.
+// - b1, weight 2, backends [0,1]
+// - b2, weight 1, backends [2,3]
+// Start the balancer group and check behavior.
+//
+// Close the balancer group.
+func (s) TestBalancerGroup_start_close(t *testing.T) {
+	cc := testutils.NewBalancerClientConn(t)
+	gator := weightedaggregator.New(cc, nil, testutils.NewTestWRR)
+	gator.Start()
+	bg := balancergroup.New(balancergroup.Options{
+		CC:                      cc,
+		BuildOpts:               balancer.BuildOptions{},
+		StateAggregator:         gator,
+		Logger:                  nil,
+		SubBalancerCloseTimeout: time.Duration(0),
+	})
+
+	// Add two balancers to group and send two resolved addresses to both
+	// balancers.
+	gator.Add(testBalancerIDs[0], 2)
+	bg.Add(testBalancerIDs[0], rrBuilder)
+	bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[0:2]}})
+	gator.Add(testBalancerIDs[1], 1)
+	bg.Add(testBalancerIDs[1], rrBuilder)
+	bg.UpdateClientConnState(testBalancerIDs[1], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[2:4]}})
+
+	m1 := make(map[string]balancer.SubConn)
+	for i := 0; i < 4; i++ {
+		addrs := <-cc.NewSubConnAddrsCh
+		sc := <-cc.NewSubConnCh
+		m1[addrs[0].Addr] = sc
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	}
+
+	// Test roundrobin on the last picker.
+	p1 := <-cc.NewPickerCh
+	want := []balancer.SubConn{
+		m1[testBackendAddrs[0].Addr], m1[testBackendAddrs[0].Addr],
+		m1[testBackendAddrs[1].Addr], m1[testBackendAddrs[1].Addr],
+		m1[testBackendAddrs[2].Addr], m1[testBackendAddrs[3].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p1)); err != nil {
+		t.Fatalf("want %v, got %v", want, err)
+	}
+
+	gator.Stop()
+	bg.Close()
+	for i := 0; i < 4; i++ {
+		(<-cc.ShutdownSubConnCh).UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Shutdown})
+	}
+}
+
+// Test that balancer group start() doesn't deadlock if the balancer calls back
+// into balancer group inline when it gets an update.
+//
+// The potential deadlock can happen if we
+//   - hold a lock and send updates to balancer (e.g. update resolved addresses)
+//   - the balancer calls back (NewSubConn or update picker) in line
+//
+// The callback will try to hold the same lock again, which will cause a
+// deadlock.
+//
+// This test starts the balancer group with a test balancer, will update picker
+// whenever it gets an address update. It's expected that start() doesn't block
+// because of deadlock.
+func (s) TestBalancerGroup_start_close_deadlock(t *testing.T) {
+	const balancerName = "stub-TestBalancerGroup_start_close_deadlock"
+	stub.Register(balancerName, stub.BalancerFuncs{})
+	builder := balancer.Get(balancerName)
+
+	cc := testutils.NewBalancerClientConn(t)
+	gator := weightedaggregator.New(cc, nil, testutils.NewTestWRR)
+	gator.Start()
+	bg := balancergroup.New(balancergroup.Options{
+		CC:                      cc,
+		BuildOpts:               balancer.BuildOptions{},
+		StateAggregator:         gator,
+		Logger:                  nil,
+		SubBalancerCloseTimeout: time.Duration(0),
+	})
+
+	gator.Add(testBalancerIDs[0], 2)
+	bg.Add(testBalancerIDs[0], builder)
+	bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{ResolverState: resolver.State{Addresses: testBackendAddrs[0:2]}})
+	gator.Add(testBalancerIDs[1], 1)
+	bg.Add(testBalancerIDs[1], builder)
+	bg.UpdateClientConnState(testBalancerIDs[1], balancer.ClientConnState{ResolverState: resolver.State{Addresses: testBackendAddrs[2:4]}})
+}
+
+// initBalancerGroupForCachingTest creates a balancer group, and initialize it
+// to be ready for caching tests.
+//
+// Two rr balancers are added to bg, each with 2 ready subConns. A sub-balancer
+// is removed later, so the balancer group returned has one sub-balancer in its
+// own map, and one sub-balancer in cache.
+func initBalancerGroupForCachingTest(t *testing.T, idleCacheTimeout time.Duration) (*weightedaggregator.Aggregator, *balancergroup.BalancerGroup, *testutils.BalancerClientConn, map[string]*testutils.TestSubConn) {
+	cc := testutils.NewBalancerClientConn(t)
+	gator := weightedaggregator.New(cc, nil, testutils.NewTestWRR)
+	gator.Start()
+	bg := balancergroup.New(balancergroup.Options{
+		CC:                      cc,
+		BuildOpts:               balancer.BuildOptions{},
+		StateAggregator:         gator,
+		Logger:                  nil,
+		SubBalancerCloseTimeout: idleCacheTimeout,
+	})
+
+	// Add two balancers to group and send two resolved addresses to both
+	// balancers.
+	gator.Add(testBalancerIDs[0], 2)
+	bg.Add(testBalancerIDs[0], rrBuilder)
+	bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[0:2]}})
+	gator.Add(testBalancerIDs[1], 1)
+	bg.Add(testBalancerIDs[1], rrBuilder)
+	bg.UpdateClientConnState(testBalancerIDs[1], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[2:4]}})
+
+	m1 := make(map[string]*testutils.TestSubConn)
+	for i := 0; i < 4; i++ {
+		addrs := <-cc.NewSubConnAddrsCh
+		sc := <-cc.NewSubConnCh
+		m1[addrs[0].Addr] = sc
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	}
+
+	// Test roundrobin on the last picker.
+	p1 := <-cc.NewPickerCh
+	want := []balancer.SubConn{
+		m1[testBackendAddrs[0].Addr], m1[testBackendAddrs[0].Addr],
+		m1[testBackendAddrs[1].Addr], m1[testBackendAddrs[1].Addr],
+		m1[testBackendAddrs[2].Addr], m1[testBackendAddrs[3].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p1)); err != nil {
+		t.Fatalf("want %v, got %v", want, err)
+	}
+
+	gator.Remove(testBalancerIDs[1])
+	bg.Remove(testBalancerIDs[1])
+	// Don't wait for SubConns to be removed after close, because they are only
+	// removed after close timeout.
+	for i := 0; i < 10; i++ {
+		select {
+		case sc := <-cc.ShutdownSubConnCh:
+			t.Fatalf("Got request to shut down subconn %v, want no shut down subconn (because subconns were still in cache)", sc)
+		default:
+		}
+		time.Sleep(time.Millisecond)
+	}
+	// Test roundrobin on the with only sub-balancer0.
+	p2 := <-cc.NewPickerCh
+	want = []balancer.SubConn{
+		m1[testBackendAddrs[0].Addr], m1[testBackendAddrs[1].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p2)); err != nil {
+		t.Fatalf("want %v, got %v", want, err)
+	}
+
+	return gator, bg, cc, m1
+}
+
+// Test that if a sub-balancer is removed, and re-added within close timeout,
+// the subConns won't be re-created.
+func (s) TestBalancerGroup_locality_caching(t *testing.T) {
+	gator, bg, cc, addrToSC := initBalancerGroupForCachingTest(t, defaultTestTimeout)
+
+	// Turn down subconn for addr2, shouldn't get picker update because
+	// sub-balancer1 was removed.
+	addrToSC[testBackendAddrs[2].Addr].UpdateState(balancer.SubConnState{
+		ConnectivityState: connectivity.TransientFailure,
+		ConnectionError:   errors.New("test error"),
+	})
+	for i := 0; i < 10; i++ {
+		select {
+		case <-cc.NewPickerCh:
+			t.Fatalf("Got new picker, want no new picker (because the sub-balancer was removed)")
+		default:
+		}
+		time.Sleep(defaultTestShortTimeout)
+	}
+
+	// Re-add sub-balancer-1, because subconns were in cache, no new subconns
+	// should be created. But a new picker will still be generated, with subconn
+	// states update to date.
+	gator.Add(testBalancerIDs[1], 1)
+	bg.Add(testBalancerIDs[1], rrBuilder)
+
+	p3 := <-cc.NewPickerCh
+	want := []balancer.SubConn{
+		addrToSC[testBackendAddrs[0].Addr], addrToSC[testBackendAddrs[0].Addr],
+		addrToSC[testBackendAddrs[1].Addr], addrToSC[testBackendAddrs[1].Addr],
+		// addr2 is down, b2 only has addr3 in READY state.
+		addrToSC[testBackendAddrs[3].Addr], addrToSC[testBackendAddrs[3].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p3)); err != nil {
+		t.Fatalf("want %v, got %v", want, err)
+	}
+
+	for i := 0; i < 10; i++ {
+		select {
+		case <-cc.NewSubConnAddrsCh:
+			t.Fatalf("Got new subconn, want no new subconn (because subconns were still in cache)")
+		default:
+		}
+		time.Sleep(defaultTestShortTimeout)
+	}
+}
+
+// Sub-balancers are put in cache when they are shut down. If balancer group is
+// closed within close timeout, all subconns should still be removed
+// immediately.
+func (s) TestBalancerGroup_locality_caching_close_group(t *testing.T) {
+	_, bg, cc, addrToSC := initBalancerGroupForCachingTest(t, defaultTestTimeout)
+
+	bg.Close()
+	// The balancer group is closed. The subconns should be shutdown immediately.
+	shutdownTimeout := time.After(time.Millisecond * 500)
+	scToShutdown := map[balancer.SubConn]int{
+		addrToSC[testBackendAddrs[0].Addr]: 1,
+		addrToSC[testBackendAddrs[1].Addr]: 1,
+		addrToSC[testBackendAddrs[2].Addr]: 1,
+		addrToSC[testBackendAddrs[3].Addr]: 1,
+	}
+	for i := 0; i < len(scToShutdown); i++ {
+		select {
+		case sc := <-cc.ShutdownSubConnCh:
+			c := scToShutdown[sc]
+			if c == 0 {
+				t.Fatalf("Got Shutdown for %v when there's %d shutdown expected", sc, c)
+			}
+			scToShutdown[sc] = c - 1
+		case <-shutdownTimeout:
+			t.Fatalf("timeout waiting for subConns (from balancer in cache) to be shut down")
+		}
+	}
+}
+
+// Sub-balancers in cache will be closed if not re-added within timeout, and
+// subConns will be shut down.
+func (s) TestBalancerGroup_locality_caching_not_read_within_timeout(t *testing.T) {
+	_, _, cc, addrToSC := initBalancerGroupForCachingTest(t, time.Second)
+
+	// The sub-balancer is not re-added within timeout. The subconns should be
+	// shut down.
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	scToShutdown := map[balancer.SubConn]int{
+		addrToSC[testBackendAddrs[2].Addr]: 1,
+		addrToSC[testBackendAddrs[3].Addr]: 1,
+	}
+	for i := 0; i < len(scToShutdown); i++ {
+		select {
+		case sc := <-cc.ShutdownSubConnCh:
+			c := scToShutdown[sc]
+			if c == 0 {
+				t.Fatalf("Got Shutdown for %v when there's %d shutdown expected", sc, c)
+			}
+			scToShutdown[sc] = c - 1
+		case <-ctx.Done():
+			t.Fatalf("timeout waiting for subConns (from balancer in cache) to be shut down")
+		}
+	}
+}
+
+// Wrap the rr builder, so it behaves the same, but has a different name.
+type noopBalancerBuilderWrapper struct {
+	balancer.Builder
+}
+
+func init() {
+	balancer.Register(&noopBalancerBuilderWrapper{Builder: rrBuilder})
+}
+
+func (*noopBalancerBuilderWrapper) Name() string {
+	return "noopBalancerBuilderWrapper"
+}
+
+// After removing a sub-balancer, re-add with same ID, but different balancer
+// builder. Old subconns should be shut down, and new subconns should be created.
+func (s) TestBalancerGroup_locality_caching_read_with_different_builder(t *testing.T) {
+	gator, bg, cc, addrToSC := initBalancerGroupForCachingTest(t, defaultTestTimeout)
+
+	// Re-add sub-balancer-1, but with a different balancer builder. The
+	// sub-balancer was still in cache, but can't be reused. This should cause
+	// old sub-balancer's subconns to be shut down immediately, and new
+	// subconns to be created.
+	gator.Add(testBalancerIDs[1], 1)
+	bg.Add(testBalancerIDs[1], &noopBalancerBuilderWrapper{rrBuilder})
+
+	// The cached sub-balancer should be closed, and the subconns should be
+	// shut down immediately.
+	shutdownTimeout := time.After(time.Millisecond * 500)
+	scToShutdown := map[balancer.SubConn]int{
+		addrToSC[testBackendAddrs[2].Addr]: 1,
+		addrToSC[testBackendAddrs[3].Addr]: 1,
+	}
+	for i := 0; i < len(scToShutdown); i++ {
+		select {
+		case sc := <-cc.ShutdownSubConnCh:
+			c := scToShutdown[sc]
+			if c == 0 {
+				t.Fatalf("Got Shutdown for %v when there's %d shutdown expected", sc, c)
+			}
+			scToShutdown[sc] = c - 1
+		case <-shutdownTimeout:
+			t.Fatalf("timeout waiting for subConns (from balancer in cache) to be shut down")
+		}
+	}
+
+	bg.UpdateClientConnState(testBalancerIDs[1], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[4:6]}})
+
+	newSCTimeout := time.After(time.Millisecond * 500)
+	scToAdd := map[string]int{
+		testBackendAddrs[4].Addr: 1,
+		testBackendAddrs[5].Addr: 1,
+	}
+	for i := 0; i < len(scToAdd); i++ {
+		select {
+		case addr := <-cc.NewSubConnAddrsCh:
+			c := scToAdd[addr[0].Addr]
+			if c == 0 {
+				t.Fatalf("Got newSubConn for %v when there's %d new expected", addr, c)
+			}
+			scToAdd[addr[0].Addr] = c - 1
+			sc := <-cc.NewSubConnCh
+			addrToSC[addr[0].Addr] = sc
+			sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+			sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+		case <-newSCTimeout:
+			t.Fatalf("timeout waiting for subConns (from new sub-balancer) to be newed")
+		}
+	}
+
+	// Test roundrobin on the new picker.
+	p3 := <-cc.NewPickerCh
+	want := []balancer.SubConn{
+		addrToSC[testBackendAddrs[0].Addr], addrToSC[testBackendAddrs[0].Addr],
+		addrToSC[testBackendAddrs[1].Addr], addrToSC[testBackendAddrs[1].Addr],
+		addrToSC[testBackendAddrs[4].Addr], addrToSC[testBackendAddrs[5].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p3)); err != nil {
+		t.Fatalf("want %v, got %v", want, err)
+	}
+}
+
+// After removing a sub-balancer, it will be kept in cache. Make sure that this
+// sub-balancer's Close is called when the balancer group is closed.
+func (s) TestBalancerGroup_CloseStopsBalancerInCache(t *testing.T) {
+	const balancerName = "stub-TestBalancerGroup_check_close"
+	closed := make(chan struct{})
+	stub.Register(balancerName, stub.BalancerFuncs{Close: func(_ *stub.BalancerData) {
+		close(closed)
+	}})
+	builder := balancer.Get(balancerName)
+
+	gator, bg, _, _ := initBalancerGroupForCachingTest(t, time.Second)
+
+	// Add balancer, and remove
+	gator.Add(testBalancerIDs[2], 1)
+	bg.Add(testBalancerIDs[2], builder)
+	gator.Remove(testBalancerIDs[2])
+	bg.Remove(testBalancerIDs[2])
+
+	// Immediately close balancergroup, before the cache timeout.
+	bg.Close()
+
+	// Make sure the removed child balancer is closed eventually.
+	select {
+	case <-closed:
+	case <-time.After(time.Second * 2):
+		t.Fatalf("timeout waiting for the child balancer in cache to be closed")
+	}
+}
+
+// TestBalancerGroupBuildOptions verifies that the balancer.BuildOptions passed
+// to the balancergroup at creation time is passed to child policies.
+func (s) TestBalancerGroupBuildOptions(t *testing.T) {
+	const (
+		balancerName = "stubBalancer-TestBalancerGroupBuildOptions"
+		userAgent    = "ua"
+	)
+
+	// Setup the stub balancer such that we can read the build options passed to
+	// it in the UpdateClientConnState method.
+	bOpts := balancer.BuildOptions{
+		DialCreds:       insecure.NewCredentials(),
+		ChannelzParent:  channelz.RegisterChannel(nil, "test channel"),
+		CustomUserAgent: userAgent,
+	}
+	stub.Register(balancerName, stub.BalancerFuncs{
+		UpdateClientConnState: func(bd *stub.BalancerData, _ balancer.ClientConnState) error {
+			if bd.BuildOptions.DialCreds != bOpts.DialCreds || bd.BuildOptions.ChannelzParent != bOpts.ChannelzParent || bd.BuildOptions.CustomUserAgent != bOpts.CustomUserAgent {
+				return fmt.Errorf("buildOptions in child balancer: %v, want %v", bd, bOpts)
+			}
+			return nil
+		},
+	})
+	cc := testutils.NewBalancerClientConn(t)
+	bg := balancergroup.New(balancergroup.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 := balancergroup.New(balancergroup.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):
+	}
+}
+
+func (s) TestBalancerGroup_ResolverError_AfterClose(t *testing.T) {
+	balancerName := t.Name()
+	resolveErrorCh := make(chan struct{}, 1)
+
+	stub.Register(balancerName, stub.BalancerFuncs{
+		ResolverError: func(_ *stub.BalancerData, _ error) {
+			resolveErrorCh <- struct{}{}
+		},
+	})
+
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              testutils.NewBalancerClientConn(t),
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+
+	bg.Add(testBalancerIDs[0], balancer.Get(balancerName))
+	bg.Close()
+
+	bg.ResolverError(errors.New("test error"))
+
+	select {
+	case <-resolveErrorCh:
+		t.Fatalf("ResolverError was called on sub-balancer after BalancerGroup was closed")
+	case <-time.After(defaultTestShortTimeout):
+	}
+}
+
+func (s) TestBalancerExitIdleOne(t *testing.T) {
+	const balancerName = "stub-balancer-test-balancergroup-exit-idle-one"
+	exitIdleCh := make(chan struct{}, 1)
+	stub.Register(balancerName, stub.BalancerFuncs{
+		ExitIdle: func(*stub.BalancerData) {
+			exitIdleCh <- struct{}{}
+		},
+	})
+	cc := testutils.NewBalancerClientConn(t)
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              cc,
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+	defer bg.Close()
+
+	// Add the stub balancer build above as a child policy.
+	builder := balancer.Get(balancerName)
+	bg.Add(testBalancerIDs[0], builder)
+
+	// Call ExitIdleOne on the child policy.
+	bg.ExitIdleOne(testBalancerIDs[0])
+	select {
+	case <-time.After(time.Second):
+		t.Fatal("Timeout when waiting for ExitIdle to be invoked on child policy")
+	case <-exitIdleCh:
+	}
+}
+
+func (s) TestBalancerGroup_ExitIdleOne_AfterClose(t *testing.T) {
+	balancerName := t.Name()
+	exitIdleCh := make(chan struct{})
+
+	stub.Register(balancerName, stub.BalancerFuncs{
+		ExitIdle: func(_ *stub.BalancerData) {
+			close(exitIdleCh)
+		},
+	})
+
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              testutils.NewBalancerClientConn(t),
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+
+	bg.Add(testBalancerIDs[0], balancer.Get(balancerName))
+	bg.Close()
+	bg.ExitIdleOne(testBalancerIDs[0])
+
+	select {
+	case <-time.After(defaultTestShortTimeout):
+	case <-exitIdleCh:
+		t.Fatalf("ExitIdleOne called ExitIdle on sub-balancer after BalancerGroup was closed")
+	}
+}
+
+func (s) TestBalancerGroup_ExitIdleOne_NonExistentID(t *testing.T) {
+	balancerName := t.Name()
+	exitIdleCh := make(chan struct{}, 1)
+
+	stub.Register(balancerName, stub.BalancerFuncs{
+		ExitIdle: func(_ *stub.BalancerData) {
+			exitIdleCh <- struct{}{}
+		},
+	})
+
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              testutils.NewBalancerClientConn(t),
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+	defer bg.Close()
+
+	bg.Add(testBalancerIDs[0], balancer.Get(balancerName))
+	bg.ExitIdleOne("non-existent-id")
+
+	select {
+	case <-time.After(defaultTestShortTimeout):
+	case <-exitIdleCh:
+		t.Fatalf("ExitIdleOne called ExitIdle on wrong sub-balancer ID")
+	}
+}
+
+// TestBalancerGracefulSwitch tests the graceful switch functionality for a
+// child of the balancer group. At first, the child is configured as a round
+// robin load balancer, and thus should behave accordingly. The test then
+// gracefully switches this child to a custom type which only creates a SubConn
+// for the second passed in address and also only picks that created SubConn.
+// The new aggregated picker should reflect this change for the child.
+func (s) TestBalancerGracefulSwitch(t *testing.T) {
+	cc := testutils.NewBalancerClientConn(t)
+	gator := weightedaggregator.New(cc, nil, testutils.NewTestWRR)
+	gator.Start()
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              cc,
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: gator,
+		Logger:          nil,
+	})
+	gator.Add(testBalancerIDs[0], 1)
+	bg.Add(testBalancerIDs[0], rrBuilder)
+	bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{ResolverState: resolver.State{Endpoints: testBackendEndpoints[0:2]}})
+
+	defer bg.Close()
+
+	m1 := make(map[string]balancer.SubConn)
+	scs := make(map[balancer.SubConn]bool)
+	for i := 0; i < 2; i++ {
+		addrs := <-cc.NewSubConnAddrsCh
+		sc := <-cc.NewSubConnCh
+		m1[addrs[0].Addr] = sc
+		scs[sc] = true
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+		sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	}
+
+	p1 := <-cc.NewPickerCh
+	want := []balancer.SubConn{
+		m1[testBackendAddrs[0].Addr], m1[testBackendAddrs[1].Addr],
+	}
+	if err := testutils.IsRoundRobin(want, testutils.SubConnFromPicker(p1)); err != nil {
+		t.Fatal(err)
+	}
+
+	// The balancer type for testBalancersIDs[0] is currently Round Robin. Now,
+	// change it to a balancer that has separate behavior logically (creating
+	// SubConn for second address in address list and always picking that
+	// SubConn), and see if the downstream behavior reflects that change.
+	childPolicyName := t.Name()
+	stub.Register(childPolicyName, stub.BalancerFuncs{
+		Init: func(bd *stub.BalancerData) {
+			bd.ChildBalancer = balancer.Get(pickfirst.Name).Build(bd.ClientConn, bd.BuildOptions)
+		},
+		Close: func(bd *stub.BalancerData) {
+			bd.ChildBalancer.Close()
+		},
+		UpdateClientConnState: func(bd *stub.BalancerData, ccs balancer.ClientConnState) error {
+			ccs.ResolverState.Endpoints = ccs.ResolverState.Endpoints[1:]
+			return bd.ChildBalancer.UpdateClientConnState(ccs)
+		},
+	})
+	cfgJSON := json.RawMessage(fmt.Sprintf(`[{%q: {}}]`, t.Name()))
+	lbCfg, err := balancergroup.ParseConfig(cfgJSON)
+	if err != nil {
+		t.Fatalf("ParseConfig(%s) failed: %v", string(cfgJSON), err)
+	}
+	if err := bg.UpdateClientConnState(testBalancerIDs[0], balancer.ClientConnState{
+		ResolverState:  resolver.State{Endpoints: testBackendEndpoints[2:4]},
+		BalancerConfig: lbCfg,
+	}); err != nil {
+		t.Fatalf("error updating ClientConn state: %v", err)
+	}
+
+	addrs := <-cc.NewSubConnAddrsCh
+	if addrs[0].Addr != testBackendAddrs[3].Addr {
+		// Verifies forwarded to new created balancer, as the wrapped pick first
+		// balancer will delete first address.
+		t.Fatalf("newSubConn called with wrong address, want: %v, got : %v", testBackendAddrs[3].Addr, addrs[0].Addr)
+	}
+	sc := <-cc.NewSubConnCh
+
+	// Update the pick first balancers SubConn as CONNECTING. This will cause
+	// the pick first balancer to UpdateState() with CONNECTING, which shouldn't send
+	// a Picker update back, as the Graceful Switch process is not complete.
+	sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Connecting})
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestShortTimeout)
+	defer cancel()
+	select {
+	case <-cc.NewPickerCh:
+		t.Fatalf("No new picker should have been sent due to the Graceful Switch process not completing")
+	case <-ctx.Done():
+	}
+
+	// Update the pick first balancers SubConn as READY. This will cause
+	// the pick first balancer to UpdateState() with READY, which should send a
+	// Picker update back, as the Graceful Switch process is complete. This
+	// Picker should always pick the pick first's created SubConn which
+	// corresponds to address 3.
+	sc.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	p2 := <-cc.NewPickerCh
+	pr, err := p2.Pick(balancer.PickInfo{})
+	if err != nil {
+		t.Fatalf("error picking: %v", err)
+	}
+	if pr.SubConn != sc {
+		t.Fatalf("picker.Pick(), want %v, got %v", sc, pr.SubConn)
+	}
+
+	// The Graceful Switch process completing for the child should cause the
+	// SubConns for the balancer being gracefully switched from to get deleted.
+	ctx, cancel = context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+	for i := 0; i < 2; i++ {
+		select {
+		case <-ctx.Done():
+			t.Fatalf("error waiting for Shutdown()")
+		case sc := <-cc.ShutdownSubConnCh:
+			// The SubConn shut down should have been one of the two created
+			// 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 := balancergroup.New(balancergroup.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
+		case <-time.After(defaultTestTimeout):
+			errCh <- fmt.Errorf("timeout waiting for ExitIdle on balancer1")
+		}
+	}()
+
+	go func() {
+		select {
+		case <-exitIdleCh2:
+			errCh <- nil
+		case <-time.After(defaultTestTimeout):
+			errCh <- fmt.Errorf("timeout waiting for ExitIdle on balancer2")
+		}
+	}()
+
+	for i := 0; i < 2; i++ {
+		if err := <-errCh; err != nil {
+			t.Fatal(err)
+		}
+	}
+}
+
+func (s) TestBalancerGroup_ExitIdle_AfterClose(t *testing.T) {
+	balancerName := t.Name()
+	exitIdleCh := make(chan struct{}, 1)
+
+	stub.Register(balancerName, stub.BalancerFuncs{
+		ExitIdle: func(_ *stub.BalancerData) {
+			exitIdleCh <- struct{}{}
+		},
+	})
+
+	bg := balancergroup.New(balancergroup.Options{
+		CC:              testutils.NewBalancerClientConn(t),
+		BuildOpts:       balancer.BuildOptions{},
+		StateAggregator: nil,
+		Logger:          nil,
+	})
+
+	bg.Add(testBalancerIDs[0], balancer.Get(balancerName))
+	bg.Close()
+	bg.ExitIdle()
+
+	select {
+	case <-exitIdleCh:
+		t.Fatalf("ExitIdle was called on sub-balancer even after BalancerGroup was closed")
+	case <-time.After(defaultTestShortTimeout):
+	}
+}
+
+// Tests verifies that the RemoveImmediately method of balancergroup closes the
+// child balancer immediately, while the Remove method does not close the child
+// balancer immediately, when the cache is enabled.
+func (s) TestBalancerGroup_RemoveImmediately(t *testing.T) {
+	// Channels to track child balancer creation and closure.
+	childLBCreated := make(chan string, 1)
+	childLBClosed := make(chan string, 1)
+
+	childLBName1 := strings.ToLower(t.Name()) + "-child-1"
+	t.Logf("Registering a child balancer with name %q", childLBName1)
+	stub.Register(childLBName1, stub.BalancerFuncs{
+		Init: func(*stub.BalancerData) {
+			childLBCreated <- childLBName1
+		},
+		Close: func(*stub.BalancerData) {
+			childLBClosed <- childLBName1
+		},
+	})
+
+	childLBName2 := strings.ToLower(t.Name()) + "-child-2"
+	t.Logf("Registering a child balancer with name %q", childLBName2)
+	stub.Register(childLBName2, stub.BalancerFuncs{
+		Init: func(*stub.BalancerData) {
+			childLBCreated <- childLBName2
+		},
+		Close: func(*stub.BalancerData) {
+			childLBClosed <- childLBName2
+		},
+	})
+
+	t.Logf("Creating a balancergroup with cache enabled")
+	tcc := testutils.NewBalancerClientConn(t)
+	bg := balancergroup.New(balancergroup.Options{
+		CC:                      tcc,
+		BuildOpts:               balancer.BuildOptions{},
+		SubBalancerCloseTimeout: defaultTestTimeout,
+	})
+
+	t.Logf("Adding a child balancer with name %q to the group", "child-1")
+	bg.AddWithClientConn("child-1", childLBName1, tcc)
+	select {
+	case <-childLBCreated:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("Timeout when waiting for child LB to be created")
+	}
+
+	t.Logf("Adding a child balancer with name %q to the group", "child-2")
+	bg.AddWithClientConn("child-2", childLBName2, tcc)
+	select {
+	case <-childLBCreated:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("Timeout when waiting for child LB to be created")
+	}
+
+	t.Logf("Removing the child balancer with name %q from the group with immediate effect", "child-1")
+	bg.RemoveImmediately("child-1")
+	select {
+	case <-childLBClosed:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("Timeout when waiting for child LB to be closed")
+	}
+
+	t.Logf("Removing the child balancer with name %q from the group with caching", "child-2")
+	bg.Remove("child-2")
+	sCtx, sCancel := context.WithTimeout(context.Background(), defaultTestShortTimeout)
+	defer sCancel()
+	select {
+	case <-childLBClosed:
+		t.Fatalf("Child LB closed when expected to be cached")
+	case <-sCtx.Done():
+	}
+
+	t.Logf("Closing the balancergroup, which should close the second child balancer")
+	bg.Close()
+	select {
+	case <-childLBClosed:
+	case <-time.After(defaultTestTimeout):
+		t.Fatalf("Timeout when waiting for child LB to be closed")
+	}
+}