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

internal/testutils/balancer.go

@@ -0,0 +1,413 @@
+/*
+ *
+ * Copyright 2020 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 testutils
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"testing"
+
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/experimental/stats"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/grpcsync"
+	"google.golang.org/grpc/resolver"
+
+	istats "google.golang.org/grpc/internal/stats"
+)
+
+// TestSubConn implements the SubConn interface, to be used in tests.
+type TestSubConn struct {
+	balancer.SubConn
+	tcc            *BalancerClientConn // the CC that owns this SubConn
+	id             string
+	ConnectCh      chan struct{}
+	stateListener  func(balancer.SubConnState)
+	healthListener func(balancer.SubConnState)
+	connectCalled  *grpcsync.Event
+	Addresses      []resolver.Address
+}
+
+// NewTestSubConn returns a newly initialized SubConn.  Typically, subconns
+// should be created via TestClientConn.NewSubConn instead, but can be useful
+// for some tests.
+func NewTestSubConn(id string) *TestSubConn {
+	return &TestSubConn{
+		ConnectCh:     make(chan struct{}, 1),
+		connectCalled: grpcsync.NewEvent(),
+		id:            id,
+	}
+}
+
+// UpdateAddresses is a no-op.
+func (tsc *TestSubConn) UpdateAddresses([]resolver.Address) {}
+
+// Connect is a no-op.
+func (tsc *TestSubConn) Connect() {
+	tsc.connectCalled.Fire()
+	select {
+	case tsc.ConnectCh <- struct{}{}:
+	default:
+	}
+}
+
+// GetOrBuildProducer is a no-op.
+func (tsc *TestSubConn) GetOrBuildProducer(balancer.ProducerBuilder) (balancer.Producer, func()) {
+	return nil, nil
+}
+
+// UpdateState pushes the state to the listener, if one is registered.
+func (tsc *TestSubConn) UpdateState(state balancer.SubConnState) {
+	<-tsc.connectCalled.Done()
+	if tsc.stateListener != nil {
+		tsc.stateListener(state)
+	}
+	// pickfirst registers a health listener synchronously while handing updates
+	// to READY. It updates the balancing state only after receiving the health
+	// update. We update the health state here so callers of tsc.UpdateState
+	// can verify picker updates as soon as UpdateState returns.
+	if state.ConnectivityState == connectivity.Ready && tsc.healthListener != nil {
+		tsc.healthListener(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	}
+}
+
+// Shutdown pushes the SubConn to the ShutdownSubConn channel in the parent
+// TestClientConn.
+func (tsc *TestSubConn) Shutdown() {
+	tsc.tcc.logger.Logf("SubConn %s: Shutdown", tsc)
+	select {
+	case tsc.tcc.ShutdownSubConnCh <- tsc:
+	default:
+	}
+}
+
+// String implements stringer to print human friendly error message.
+func (tsc *TestSubConn) String() string {
+	return tsc.id
+}
+
+// RegisterHealthListener sends a READY update to mock a situation when no
+// health checking mechanisms are configured.
+func (tsc *TestSubConn) RegisterHealthListener(lis func(balancer.SubConnState)) {
+	tsc.healthListener = lis
+}
+
+// BalancerClientConn is a mock balancer.ClientConn used in tests.
+type BalancerClientConn struct {
+	internal.EnforceClientConnEmbedding
+	logger Logger
+
+	NewSubConnAddrsCh      chan []resolver.Address // the last 10 []Address to create subconn.
+	NewSubConnCh           chan *TestSubConn       // the last 10 subconn created.
+	ShutdownSubConnCh      chan *TestSubConn       // the last 10 subconn removed.
+	UpdateAddressesAddrsCh chan []resolver.Address // last updated address via UpdateAddresses().
+
+	NewPickerCh  chan balancer.Picker            // the last picker updated.
+	NewStateCh   chan connectivity.State         // the last state.
+	ResolveNowCh chan resolver.ResolveNowOptions // the last ResolveNow().
+
+	subConnIdx int
+}
+
+// NewBalancerClientConn creates a BalancerClientConn.
+func NewBalancerClientConn(t *testing.T) *BalancerClientConn {
+	return &BalancerClientConn{
+		logger: t,
+
+		NewSubConnAddrsCh:      make(chan []resolver.Address, 10),
+		NewSubConnCh:           make(chan *TestSubConn, 10),
+		ShutdownSubConnCh:      make(chan *TestSubConn, 10),
+		UpdateAddressesAddrsCh: make(chan []resolver.Address, 1),
+
+		NewPickerCh:  make(chan balancer.Picker, 1),
+		NewStateCh:   make(chan connectivity.State, 1),
+		ResolveNowCh: make(chan resolver.ResolveNowOptions, 1),
+	}
+}
+
+// NewSubConn creates a new SubConn.
+func (tcc *BalancerClientConn) NewSubConn(a []resolver.Address, o balancer.NewSubConnOptions) (balancer.SubConn, error) {
+	sc := &TestSubConn{
+		tcc:           tcc,
+		id:            fmt.Sprintf("sc%d", tcc.subConnIdx),
+		ConnectCh:     make(chan struct{}, 1),
+		stateListener: o.StateListener,
+		connectCalled: grpcsync.NewEvent(),
+		Addresses:     a,
+	}
+	tcc.subConnIdx++
+	tcc.logger.Logf("testClientConn: NewSubConn(%v, %+v) => %s", a, o, sc)
+	select {
+	case tcc.NewSubConnAddrsCh <- a:
+	default:
+	}
+
+	select {
+	case tcc.NewSubConnCh <- sc:
+	default:
+	}
+
+	return sc, nil
+}
+
+// MetricsRecorder returns an empty MetricsRecorderList.
+func (*BalancerClientConn) MetricsRecorder() stats.MetricsRecorder {
+	return istats.NewMetricsRecorderList(nil)
+}
+
+// RemoveSubConn is a nop; tests should all be updated to use sc.Shutdown()
+// instead.
+func (tcc *BalancerClientConn) RemoveSubConn(sc balancer.SubConn) {
+	tcc.logger.Errorf("RemoveSubConn(%v) called unexpectedly", sc)
+}
+
+// UpdateAddresses updates the addresses on the SubConn.
+func (tcc *BalancerClientConn) UpdateAddresses(sc balancer.SubConn, addrs []resolver.Address) {
+	tcc.logger.Logf("testutils.BalancerClientConn: UpdateAddresses(%v, %+v)", sc, addrs)
+	select {
+	case tcc.UpdateAddressesAddrsCh <- addrs:
+	default:
+	}
+}
+
+// UpdateState updates connectivity state and picker.
+func (tcc *BalancerClientConn) UpdateState(bs balancer.State) {
+	tcc.logger.Logf("testutils.BalancerClientConn: UpdateState(%v)", bs)
+	select {
+	case <-tcc.NewStateCh:
+	default:
+	}
+	tcc.NewStateCh <- bs.ConnectivityState
+
+	select {
+	case <-tcc.NewPickerCh:
+	default:
+	}
+	tcc.NewPickerCh <- bs.Picker
+}
+
+// ResolveNow panics.
+func (tcc *BalancerClientConn) ResolveNow(o resolver.ResolveNowOptions) {
+	select {
+	case <-tcc.ResolveNowCh:
+	default:
+	}
+	tcc.ResolveNowCh <- o
+}
+
+// Target panics.
+func (tcc *BalancerClientConn) Target() string {
+	panic("not implemented")
+}
+
+// WaitForErrPicker waits until an error picker is pushed to this ClientConn.
+// Returns error if the provided context expires or a non-error picker is pushed
+// to the ClientConn.
+func (tcc *BalancerClientConn) WaitForErrPicker(ctx context.Context) error {
+	select {
+	case <-ctx.Done():
+		return errors.New("timeout when waiting for an error picker")
+	case picker := <-tcc.NewPickerCh:
+		if _, perr := picker.Pick(balancer.PickInfo{Ctx: ctx}); perr == nil {
+			return fmt.Errorf("balancer returned a picker which is not an error picker")
+		}
+	}
+	return nil
+}
+
+// WaitForPickerWithErr waits until an error picker is pushed to this
+// ClientConn with the error matching the wanted error.  Returns an error if
+// the provided context expires, including the last received picker error (if
+// any).
+func (tcc *BalancerClientConn) WaitForPickerWithErr(ctx context.Context, want error) error {
+	lastErr := errors.New("received no picker")
+	for {
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("timeout when waiting for an error picker with %v; last picker error: %v", want, lastErr)
+		case picker := <-tcc.NewPickerCh:
+			if _, lastErr = picker.Pick(balancer.PickInfo{Ctx: ctx}); lastErr != nil && lastErr.Error() == want.Error() {
+				return nil
+			}
+		}
+	}
+}
+
+// WaitForConnectivityState waits until the state pushed to this ClientConn
+// matches the wanted state.  Returns an error if the provided context expires,
+// including the last received state (if any).
+func (tcc *BalancerClientConn) WaitForConnectivityState(ctx context.Context, want connectivity.State) error {
+	var lastState connectivity.State = -1
+	for {
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("timeout when waiting for state to be %s; last state: %s", want, lastState)
+		case s := <-tcc.NewStateCh:
+			if s == want {
+				return nil
+			}
+			lastState = s
+		}
+	}
+}
+
+// WaitForRoundRobinPicker waits for a picker that passes IsRoundRobin.  Also
+// drains the matching state channel and requires it to be READY (if an entry
+// is pending) to be considered.  Returns an error if the provided context
+// expires, including the last received error from IsRoundRobin or the picker
+// (if any).
+func (tcc *BalancerClientConn) WaitForRoundRobinPicker(ctx context.Context, want ...balancer.SubConn) error {
+	lastErr := errors.New("received no picker")
+	for {
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("timeout when waiting for round robin picker with %v; last error: %v", want, lastErr)
+		case p := <-tcc.NewPickerCh:
+			s := connectivity.Ready
+			select {
+			case s = <-tcc.NewStateCh:
+			default:
+			}
+			if s != connectivity.Ready {
+				lastErr = fmt.Errorf("received state %v instead of ready", s)
+				break
+			}
+			var pickerErr error
+			if err := IsRoundRobin(want, func() balancer.SubConn {
+				sc, err := p.Pick(balancer.PickInfo{Ctx: ctx})
+				if err != nil {
+					pickerErr = err
+				} else if sc.Done != nil {
+					sc.Done(balancer.DoneInfo{})
+				}
+				return sc.SubConn
+			}); pickerErr != nil {
+				lastErr = pickerErr
+				continue
+			} else if err != nil {
+				lastErr = err
+				continue
+			}
+			return nil
+		}
+	}
+}
+
+// WaitForPicker waits for a picker that results in f returning nil.  If the
+// context expires, returns the last error returned by f (if any).
+func (tcc *BalancerClientConn) WaitForPicker(ctx context.Context, f func(balancer.Picker) error) error {
+	lastErr := errors.New("received no picker")
+	for {
+		select {
+		case <-ctx.Done():
+			return fmt.Errorf("timeout when waiting for picker; last error: %v", lastErr)
+		case p := <-tcc.NewPickerCh:
+			if err := f(p); err != nil {
+				lastErr = err
+				continue
+			}
+			return nil
+		}
+	}
+}
+
+// IsRoundRobin checks whether f's return value is roundrobin of elements from
+// want. But it doesn't check for the order. Note that want can contain
+// duplicate items, which makes it weight-round-robin.
+//
+// Step 1. the return values of f should form a permutation of all elements in
+// want, but not necessary in the same order. E.g. if want is {a,a,b}, the check
+// fails if f returns:
+//   - {a,a,a}: third a is returned before b
+//   - {a,b,b}: second b is returned before the second a
+//
+// If error is found in this step, the returned error contains only the first
+// iteration until where it goes wrong.
+//
+// Step 2. the return values of f should be repetitions of the same permutation.
+// E.g. if want is {a,a,b}, the check fails if f returns:
+//   - {a,b,a,b,a,a}: though it satisfies step 1, the second iteration is not
+//     repeating the first iteration.
+//
+// If error is found in this step, the returned error contains the first
+// iteration + the second iteration until where it goes wrong.
+func IsRoundRobin(want []balancer.SubConn, f func() balancer.SubConn) error {
+	wantSet := make(map[balancer.SubConn]int) // SubConn -> count, for weighted RR.
+	for _, sc := range want {
+		wantSet[sc]++
+	}
+
+	// The first iteration: makes sure f's return values form a permutation of
+	// elements in want.
+	//
+	// Also keep the returns values in a slice, so we can compare the order in
+	// the second iteration.
+	gotSliceFirstIteration := make([]balancer.SubConn, 0, len(want))
+	for range want {
+		got := f()
+		gotSliceFirstIteration = append(gotSliceFirstIteration, got)
+		wantSet[got]--
+		if wantSet[got] < 0 {
+			return fmt.Errorf("non-roundrobin want: %v, result: %v", want, gotSliceFirstIteration)
+		}
+	}
+
+	// The second iteration should repeat the first iteration.
+	var gotSliceSecondIteration []balancer.SubConn
+	for i := 0; i < 2; i++ {
+		for _, w := range gotSliceFirstIteration {
+			g := f()
+			gotSliceSecondIteration = append(gotSliceSecondIteration, g)
+			if w != g {
+				return fmt.Errorf("non-roundrobin, first iter: %v, second iter: %v", gotSliceFirstIteration, gotSliceSecondIteration)
+			}
+		}
+	}
+
+	return nil
+}
+
+// SubConnFromPicker returns a function which returns a SubConn by calling the
+// Pick() method of the provided picker. There is no caching of SubConns here.
+// Every invocation of the returned function results in a new pick.
+func SubConnFromPicker(p balancer.Picker) func() balancer.SubConn {
+	return func() balancer.SubConn {
+		scst, _ := p.Pick(balancer.PickInfo{})
+		return scst.SubConn
+	}
+}
+
+// ErrTestConstPicker is error returned by test const picker.
+var ErrTestConstPicker = fmt.Errorf("const picker error")
+
+// TestConstPicker is a const picker for tests.
+type TestConstPicker struct {
+	Err error
+	SC  balancer.SubConn
+}
+
+// Pick returns the const SubConn or the error.
+func (tcp *TestConstPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
+	if tcp.Err != nil {
+		return balancer.PickResult{}, tcp.Err
+	}
+	return balancer.PickResult{SubConn: tcp.SC}, nil
+}

internal/testutils/blocking_context_dialer.go

@@ -0,0 +1,139 @@
+/*
+ *
+ * Copyright 2024 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 testutils
+
+import (
+	"context"
+	"net"
+	"sync"
+
+	"google.golang.org/grpc/grpclog"
+)
+
+var logger = grpclog.Component("testutils")
+
+// BlockingDialer is a dialer that waits for Resume() to be called before
+// dialing.
+type BlockingDialer struct {
+	// mu protects holds.
+	mu sync.Mutex
+	// holds maps network addresses to a list of holds for that address.
+	holds map[string][]*Hold
+}
+
+// NewBlockingDialer returns a dialer that waits for Resume() to be called
+// before dialing.
+func NewBlockingDialer() *BlockingDialer {
+	return &BlockingDialer{
+		holds: make(map[string][]*Hold),
+	}
+}
+
+// DialContext implements a context dialer for use with grpc.WithContextDialer
+// dial option for a BlockingDialer.
+func (d *BlockingDialer) DialContext(ctx context.Context, addr string) (net.Conn, error) {
+	d.mu.Lock()
+	holds := d.holds[addr]
+	if len(holds) == 0 {
+		// No hold for this addr.
+		d.mu.Unlock()
+		return (&net.Dialer{}).DialContext(ctx, "tcp", addr)
+	}
+	hold := holds[0]
+	d.holds[addr] = holds[1:]
+	d.mu.Unlock()
+
+	logger.Infof("Hold %p: Intercepted connection attempt to addr %q", hold, addr)
+	close(hold.waitCh)
+	select {
+	case err := <-hold.blockCh:
+		if err != nil {
+			return nil, err
+		}
+		return (&net.Dialer{}).DialContext(ctx, "tcp", addr)
+	case <-ctx.Done():
+		logger.Infof("Hold %p: Connection attempt to addr %q timed out", hold, addr)
+		return nil, ctx.Err()
+	}
+}
+
+// Hold is a handle to a single connection attempt. It can be used to block,
+// fail and succeed connection attempts.
+type Hold struct {
+	// dialer is the dialer that created this hold.
+	dialer *BlockingDialer
+	// waitCh is closed when a connection attempt is received.
+	waitCh chan struct{}
+	// blockCh receives the value to return from DialContext for this connection
+	// attempt (nil on resume, an error on fail). It receives at most 1 value.
+	blockCh chan error
+	// addr is the address that this hold is for.
+	addr string
+}
+
+// Hold blocks the dialer when a connection attempt is made to the given addr.
+// A hold is valid for exactly one connection attempt. Multiple holds for an
+// addr can be added, and they will apply in the order that the connections are
+// attempted.
+func (d *BlockingDialer) Hold(addr string) *Hold {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	h := Hold{dialer: d, blockCh: make(chan error, 1), waitCh: make(chan struct{}), addr: addr}
+	d.holds[addr] = append(d.holds[addr], &h)
+	return &h
+}
+
+// Wait blocks until there is a connection attempt on this Hold, or the context
+// expires. Return false if the context has expired, true otherwise.
+func (h *Hold) Wait(ctx context.Context) bool {
+	logger.Infof("Hold %p: Waiting for a connection attempt to addr %q", h, h.addr)
+	select {
+	case <-ctx.Done():
+		return false
+	case <-h.waitCh:
+		return true
+	}
+}
+
+// Resume unblocks the dialer for the given addr. Either Resume or Fail must be
+// called at most once on a hold. Otherwise, Resume panics.
+func (h *Hold) Resume() {
+	logger.Infof("Hold %p: Resuming connection attempt to addr %q", h, h.addr)
+	h.blockCh <- nil
+	close(h.blockCh)
+}
+
+// Fail fails the connection attempt. Either Resume or Fail must be
+// called at most once on a hold. Otherwise, Resume panics.
+func (h *Hold) Fail(err error) {
+	logger.Infof("Hold %p: Failing connection attempt to addr %q", h, h.addr)
+	h.blockCh <- err
+	close(h.blockCh)
+}
+
+// IsStarted returns true if this hold has received a connection attempt.
+func (h *Hold) IsStarted() bool {
+	select {
+	case <-h.waitCh:
+		return true
+	default:
+		return false
+	}
+}

internal/testutils/blocking_context_dialer_test.go

@@ -0,0 +1,214 @@
+/*
+ *
+ * Copyright 2024 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 testutils
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+)
+
+const (
+	testTimeout      = 5 * time.Second
+	testShortTimeout = 10 * time.Millisecond
+)
+
+func (s) TestBlockingDialer_NoHold(t *testing.T) {
+	lis, err := LocalTCPListener()
+	if err != nil {
+		t.Fatalf("Failed to listen: %v", err)
+	}
+	defer lis.Close()
+
+	d := NewBlockingDialer()
+
+	// This should not block.
+	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
+	defer cancel()
+	conn, err := d.DialContext(ctx, lis.Addr().String())
+	if err != nil {
+		t.Fatalf("Failed to dial: %v", err)
+	}
+	conn.Close()
+}
+
+func (s) TestBlockingDialer_HoldWaitResume(t *testing.T) {
+	lis, err := LocalTCPListener()
+	if err != nil {
+		t.Fatalf("Failed to listen: %v", err)
+	}
+	defer lis.Close()
+
+	d := NewBlockingDialer()
+	h := d.Hold(lis.Addr().String())
+
+	if h.IsStarted() {
+		t.Fatalf("hold.IsStarted() = true, want false")
+	}
+
+	done := make(chan struct{})
+	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
+	defer cancel()
+	go func() {
+		defer close(done)
+		conn, err := d.DialContext(ctx, lis.Addr().String())
+		if err != nil {
+			t.Errorf("BlockingDialer.DialContext() got error: %v, want success", err)
+			return
+		}
+
+		if !h.IsStarted() {
+			t.Errorf("hold.IsStarted() = false, want true")
+		}
+		conn.Close()
+	}()
+
+	// This should block until the goroutine above is scheduled.
+	if !h.Wait(ctx) {
+		t.Fatalf("Timeout while waiting for a connection attempt to %q", h.addr)
+	}
+
+	if !h.IsStarted() {
+		t.Errorf("hold.IsStarted() = false, want true")
+	}
+
+	select {
+	case <-done:
+		t.Fatalf("Expected dialer to be blocked.")
+	case <-time.After(testShortTimeout):
+	}
+
+	h.Resume() // Unblock the above goroutine.
+
+	select {
+	case <-done:
+	case <-ctx.Done():
+		t.Errorf("Timeout waiting for connection attempt to resume.")
+	}
+}
+
+func (s) TestBlockingDialer_HoldWaitFail(t *testing.T) {
+	lis, err := LocalTCPListener()
+	if err != nil {
+		t.Fatalf("Failed to listen: %v", err)
+	}
+	defer lis.Close()
+
+	d := NewBlockingDialer()
+	h := d.Hold(lis.Addr().String())
+
+	wantErr := errors.New("test error")
+
+	dialError := make(chan error)
+	ctx, cancel := context.WithTimeout(context.Background(), testTimeout)
+	defer cancel()
+	go func() {
+		_, err := d.DialContext(ctx, lis.Addr().String())
+		dialError <- err
+	}()
+
+	if !h.Wait(ctx) {
+		t.Fatal("Timeout while waiting for a connection attempt to " + h.addr)
+	}
+	select {
+	case err = <-dialError:
+		t.Errorf("DialContext got unblocked with err %v. Want DialContext to still be blocked after Wait()", err)
+	case <-time.After(testShortTimeout):
+	}
+
+	h.Fail(wantErr)
+
+	select {
+	case err = <-dialError:
+		if !errors.Is(err, wantErr) {
+			t.Errorf("BlockingDialer.DialContext() after Fail(): got error %v, want %v", err, wantErr)
+		}
+	case <-ctx.Done():
+		t.Errorf("Timeout waiting for connection attempt to fail.")
+	}
+}
+
+func (s) TestBlockingDialer_ContextCanceled(t *testing.T) {
+	lis, err := LocalTCPListener()
+	if err != nil {
+		t.Fatalf("Failed to listen: %v", err)
+	}
+	defer lis.Close()
+
+	d := NewBlockingDialer()
+	h := d.Hold(lis.Addr().String())
+
+	dialErr := make(chan error)
+	testCtx, cancel := context.WithTimeout(context.Background(), testTimeout)
+	defer cancel()
+
+	ctx, cancel := context.WithCancel(testCtx)
+	defer cancel()
+	go func() {
+		_, err := d.DialContext(ctx, lis.Addr().String())
+		dialErr <- err
+	}()
+	if !h.Wait(testCtx) {
+		t.Errorf("Timeout while waiting for a connection attempt to %q", h.addr)
+	}
+
+	cancel()
+
+	select {
+	case err = <-dialErr:
+		if !errors.Is(err, context.Canceled) {
+			t.Errorf("BlockingDialer.DialContext() after context cancel: got error %v, want %v", err, context.Canceled)
+		}
+	case <-testCtx.Done():
+		t.Errorf("Timeout while waiting for Wait to return.")
+	}
+
+	h.Resume() // noop, just make sure nothing bad happen.
+}
+
+func (s) TestBlockingDialer_CancelWait(t *testing.T) {
+	lis, err := LocalTCPListener()
+	if err != nil {
+		t.Fatalf("Failed to listen: %v", err)
+	}
+	defer lis.Close()
+
+	d := NewBlockingDialer()
+	h := d.Hold(lis.Addr().String())
+
+	testCtx, cancel := context.WithTimeout(context.Background(), testTimeout)
+	defer cancel()
+
+	ctx, cancel := context.WithCancel(testCtx)
+	cancel()
+	done := make(chan struct{})
+	go func() {
+		if h.Wait(ctx) {
+			t.Errorf("Expected cancel to return false when context expires")
+		}
+		done <- struct{}{}
+	}()
+
+	select {
+	case <-done:
+	case <-testCtx.Done():
+		t.Errorf("Timeout while waiting for Wait to return.")
+	}
+}

internal/testutils/channel.go

@@ -0,0 +1,118 @@
+/*
+ *
+ * Copyright 2020 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 testutils
+
+import (
+	"context"
+)
+
+// DefaultChanBufferSize is the default buffer size of the underlying channel.
+const DefaultChanBufferSize = 1
+
+// Channel wraps a generic channel and provides a timed receive operation.
+type Channel struct {
+	// C is the underlying channel on which values sent using the SendXxx()
+	// methods are delivered. Tests which cannot use ReceiveXxx() for whatever
+	// reasons can use C to read the values.
+	C chan any
+}
+
+// Send sends value on the underlying channel.
+func (c *Channel) Send(value any) {
+	c.C <- value
+}
+
+// SendContext sends value on the underlying channel, or returns an error if
+// the context expires.
+func (c *Channel) SendContext(ctx context.Context, value any) error {
+	select {
+	case c.C <- value:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+// SendOrFail attempts to send value on the underlying channel.  Returns true
+// if successful or false if the channel was full.
+func (c *Channel) SendOrFail(value any) bool {
+	select {
+	case c.C <- value:
+		return true
+	default:
+		return false
+	}
+}
+
+// ReceiveOrFail returns the value on the underlying channel and true, or nil
+// and false if the channel was empty.
+func (c *Channel) ReceiveOrFail() (any, bool) {
+	select {
+	case got := <-c.C:
+		return got, true
+	default:
+		return nil, false
+	}
+}
+
+// Drain drains the channel by repeatedly reading from it until it is empty.
+func (c *Channel) Drain() {
+	for {
+		select {
+		case <-c.C:
+		default:
+			return
+		}
+	}
+}
+
+// Receive returns the value received on the underlying channel, or the error
+// returned by ctx if it is closed or cancelled.
+func (c *Channel) Receive(ctx context.Context) (any, error) {
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case got := <-c.C:
+		return got, nil
+	}
+}
+
+// Replace clears the value on the underlying channel, and sends the new value.
+//
+// It's expected to be used with a size-1 channel, to only keep the most
+// up-to-date item. This method is inherently racy when invoked concurrently
+// from multiple goroutines.
+func (c *Channel) Replace(value any) {
+	for {
+		select {
+		case c.C <- value:
+			return
+		case <-c.C:
+		}
+	}
+}
+
+// NewChannel returns a new Channel.
+func NewChannel() *Channel {
+	return NewChannelWithSize(DefaultChanBufferSize)
+}
+
+// NewChannelWithSize returns a new Channel with a buffer of bufSize.
+func NewChannelWithSize(bufSize int) *Channel {
+	return &Channel{C: make(chan any, bufSize)}
+}

internal/testutils/envconfig.go

@@ -0,0 +1,33 @@
+package testutils
+
+/*
+ *
+ * Copyright 2025 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.
+ */
+
+import (
+	"testing"
+)
+
+// SetEnvConfig sets the value of the given variable to the specified value,
+// taking care of restoring the original value after the test completes.
+func SetEnvConfig[T any](t *testing.T, variable *T, value T) {
+	t.Helper()
+	old := *variable
+	t.Cleanup(func() {
+		*variable = old
+	})
+	*variable = value
+}

internal/testutils/fakegrpclb/server.go

@@ -0,0 +1,250 @@
+/*
+ *
+ * Copyright 2022 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 fakegrpclb provides a fake implementation of the grpclb server.
+package fakegrpclb
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"net/netip"
+	"strconv"
+	"sync"
+	"time"
+
+	"google.golang.org/grpc"
+	lbgrpc "google.golang.org/grpc/balancer/grpclb/grpc_lb_v1"
+	lbpb "google.golang.org/grpc/balancer/grpclb/grpc_lb_v1"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/internal/pretty"
+	"google.golang.org/grpc/status"
+)
+
+var logger = grpclog.Component("fake_grpclb")
+
+// ServerParams wraps options passed while creating a Server.
+type ServerParams struct {
+	ListenPort    int                 // Listening port for the balancer server.
+	ServerOptions []grpc.ServerOption // gRPC options for the balancer server.
+
+	LoadBalancedServiceName string   // Service name being load balanced for.
+	LoadBalancedServicePort int      // Service port being load balanced for.
+	BackendAddresses        []string // Service backends to balance load across.
+	ShortStream             bool     // End balancer stream after sending server list.
+}
+
+// Server is a fake implementation of the grpclb LoadBalancer service. It does
+// not support stats reporting from clients, and always sends back a static list
+// of backends to the client to balance load across.
+//
+// It is safe for concurrent access.
+type Server struct {
+	lbgrpc.UnimplementedLoadBalancerServer
+
+	// Options copied over from ServerParams passed to NewServer.
+	sOpts       []grpc.ServerOption // gRPC server options.
+	serviceName string              // Service name being load balanced for.
+	servicePort int                 // Service port being load balanced for.
+	shortStream bool                // End balancer stream after sending server list.
+
+	// Values initialized using ServerParams passed to NewServer.
+	backends []*lbpb.Server // Service backends to balance load across.
+	lis      net.Listener   // Listener for grpc connections to the LoadBalancer service.
+
+	// mu guards access to below fields.
+	mu         sync.Mutex
+	grpcServer *grpc.Server // Underlying grpc server.
+	address    string       // Actual listening address.
+
+	stopped chan struct{} // Closed when Stop() is called.
+}
+
+// NewServer creates a new Server with passed in params. Returns a non-nil error
+// if the params are invalid.
+func NewServer(params ServerParams) (*Server, error) {
+	var servers []*lbpb.Server
+	for _, addr := range params.BackendAddresses {
+		ipStr, portStr, err := net.SplitHostPort(addr)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse list of backend address %q: %v", addr, err)
+		}
+		ip, err := netip.ParseAddr(ipStr)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse ip %q: %v", ipStr, err)
+		}
+		port, err := strconv.Atoi(portStr)
+		if err != nil {
+			return nil, fmt.Errorf("failed to convert port %q to int", portStr)
+		}
+		logger.Infof("Adding backend ip: %q, port: %d to server list", ip.String(), port)
+		servers = append(servers, &lbpb.Server{
+			IpAddress: ip.AsSlice(),
+			Port:      int32(port),
+		})
+	}
+
+	lis, err := net.Listen("tcp", "localhost:"+strconv.Itoa(params.ListenPort))
+	if err != nil {
+		return nil, fmt.Errorf("failed to listen on port %q: %v", params.ListenPort, err)
+	}
+
+	return &Server{
+		sOpts:       params.ServerOptions,
+		serviceName: params.LoadBalancedServiceName,
+		servicePort: params.LoadBalancedServicePort,
+		shortStream: params.ShortStream,
+		backends:    servers,
+		lis:         lis,
+		address:     lis.Addr().String(),
+		stopped:     make(chan struct{}),
+	}, nil
+}
+
+// Serve starts serving the LoadBalancer service on a gRPC server.
+//
+// It returns early with a non-nil error if it is unable to start serving.
+// Otherwise, it blocks until Stop() is called, at which point it returns the
+// error returned by the underlying grpc.Server's Serve() method.
+func (s *Server) Serve() error {
+	s.mu.Lock()
+	if s.grpcServer != nil {
+		s.mu.Unlock()
+		return errors.New("Serve() called multiple times")
+	}
+
+	server := grpc.NewServer(s.sOpts...)
+	s.grpcServer = server
+	s.mu.Unlock()
+
+	logger.Infof("Begin listening on %s", s.lis.Addr().String())
+	lbgrpc.RegisterLoadBalancerServer(server, s)
+	return server.Serve(s.lis) // This call will block.
+}
+
+// Stop stops serving the LoadBalancer service and unblocks the preceding call
+// to Serve().
+func (s *Server) Stop() {
+	defer close(s.stopped)
+	s.mu.Lock()
+	if s.grpcServer != nil {
+		s.grpcServer.Stop()
+		s.grpcServer = nil
+	}
+	s.mu.Unlock()
+}
+
+// Address returns the host:port on which the LoadBalancer service is serving.
+func (s *Server) Address() string {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.address
+}
+
+// BalanceLoad provides a fake implementation of the LoadBalancer service.
+func (s *Server) BalanceLoad(stream lbgrpc.LoadBalancer_BalanceLoadServer) error {
+	logger.Info("New BalancerLoad stream started")
+
+	req, err := stream.Recv()
+	if err == io.EOF {
+		logger.Warning("Received EOF when reading from the stream")
+		return nil
+	}
+	if err != nil {
+		logger.Warningf("Failed to read LoadBalanceRequest from stream: %v", err)
+		return err
+	}
+	logger.Infof("Received LoadBalancerRequest:\n%s", pretty.ToJSON(req))
+
+	// Initial request contains the service being load balanced for.
+	initialReq := req.GetInitialRequest()
+	if initialReq == nil {
+		logger.Info("First message on the stream does not contain an InitialLoadBalanceRequest")
+		return status.Error(codes.Unknown, "First request not an InitialLoadBalanceRequest")
+	}
+
+	// Basic validation of the service name and port from the incoming request.
+	//
+	// Clients targeting service:port can sometimes include the ":port" suffix in
+	// their requested names; handle this case.
+	serviceName, port, err := net.SplitHostPort(initialReq.Name)
+	if err != nil {
+		// Requested name did not contain a port. So, use the name as is.
+		serviceName = initialReq.Name
+	} else {
+		p, err := strconv.Atoi(port)
+		if err != nil {
+			logger.Infof("Failed to parse requested service port %q to integer", port)
+			return status.Error(codes.Unknown, "Bad requested service port number")
+		}
+		if p != s.servicePort {
+			logger.Infof("Requested service port number %d does not match expected %d", p, s.servicePort)
+			return status.Error(codes.Unknown, "Bad requested service port number")
+		}
+	}
+	if serviceName != s.serviceName {
+		logger.Infof("Requested service name %q does not match expected %q", serviceName, s.serviceName)
+		return status.Error(codes.NotFound, "Bad requested service name")
+	}
+
+	// Empty initial response disables stats reporting from the client. Stats
+	// reporting from the client is used to determine backend load and is not
+	// required for the purposes of this fake.
+	initResp := &lbpb.LoadBalanceResponse{
+		LoadBalanceResponseType: &lbpb.LoadBalanceResponse_InitialResponse{
+			InitialResponse: &lbpb.InitialLoadBalanceResponse{},
+		},
+	}
+	if err := stream.Send(initResp); err != nil {
+		logger.Warningf("Failed to send InitialLoadBalanceResponse on the stream: %v", err)
+		return err
+	}
+
+	resp := &lbpb.LoadBalanceResponse{
+		LoadBalanceResponseType: &lbpb.LoadBalanceResponse_ServerList{
+			ServerList: &lbpb.ServerList{Servers: s.backends},
+		},
+	}
+	logger.Infof("Sending response with server list: %s", pretty.ToJSON(resp))
+	if err := stream.Send(resp); err != nil {
+		logger.Warningf("Failed to send InitialLoadBalanceResponse on the stream: %v", err)
+		return err
+	}
+
+	if s.shortStream {
+		logger.Info("Ending stream early as the short stream option was set")
+		return nil
+	}
+
+	for {
+		select {
+		case <-stream.Context().Done():
+			return nil
+		case <-s.stopped:
+			return nil
+		case <-time.After(10 * time.Second):
+			logger.Infof("Sending response with server list: %s", pretty.ToJSON(resp))
+			if err := stream.Send(resp); err != nil {
+				logger.Warningf("Failed to send InitialLoadBalanceResponse on the stream: %v", err)
+				return err
+			}
+		}
+	}
+}

internal/testutils/http_client.go

@@ -0,0 +1,63 @@
+/*
+ *
+ * Copyright 2020 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 testutils
+
+import (
+	"context"
+	"net/http"
+	"time"
+)
+
+// DefaultHTTPRequestTimeout is the default timeout value for the amount of time
+// this client waits for a response to be pushed on RespChan before it fails the
+// Do() call.
+const DefaultHTTPRequestTimeout = 1 * time.Second
+
+// FakeHTTPClient helps mock out HTTP calls made by the code under test. It
+// makes HTTP requests made by the code under test available through a channel,
+// and makes it possible to inject various responses.
+type FakeHTTPClient struct {
+	// ReqChan exposes the HTTP.Request made by the code under test.
+	ReqChan *Channel
+	// RespChan is a channel on which this fake client accepts responses to be
+	// sent to the code under test.
+	RespChan *Channel
+	// Err, if set, is returned by Do().
+	Err error
+	// RecvTimeout is the amount of the time this client waits for a response to
+	// be pushed on RespChan before it fails the Do() call. If this field is
+	// left unspecified, DefaultHTTPRequestTimeout is used.
+	RecvTimeout time.Duration
+}
+
+// Do pushes req on ReqChan and returns the response available on RespChan.
+func (fc *FakeHTTPClient) Do(req *http.Request) (*http.Response, error) {
+	fc.ReqChan.Send(req)
+
+	timeout := fc.RecvTimeout
+	if timeout == 0 {
+		timeout = DefaultHTTPRequestTimeout
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+	val, err := fc.RespChan.Receive(ctx)
+	if err != nil {
+		return nil, err
+	}
+	return val.(*http.Response), fc.Err
+}

internal/testutils/local_listener.go

@@ -0,0 +1,26 @@
+/*
+ *
+ * Copyright 2020 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 testutils
+
+import "net"
+
+// LocalTCPListener returns a net.Listener listening on local address and port.
+func LocalTCPListener() (net.Listener, error) {
+	return net.Listen("tcp", "localhost:0")
+}

internal/testutils/marshal_any.go

@@ -0,0 +1,37 @@
+/*
+ *
+ * Copyright 2021 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 testutils
+
+import (
+	"testing"
+
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/anypb"
+)
+
+// MarshalAny is a convenience function to marshal protobuf messages into any
+// protos. function will fail the test with a fatal error if the marshaling fails.
+func MarshalAny(t *testing.T, m proto.Message) *anypb.Any {
+	t.Helper()
+
+	a, err := anypb.New(m)
+	if err != nil {
+		t.Fatalf("Failed to marshal proto %+v into an Any: %v", m, err)
+	}
+	return a
+}

internal/testutils/parse_port.go

@@ -0,0 +1,39 @@
+/*
+ *
+ * Copyright 2023 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 testutils
+
+import (
+	"net"
+	"strconv"
+	"testing"
+)
+
+// ParsePort returns the port from the given address string, as a unit32.
+func ParsePort(t *testing.T, addr string) uint32 {
+	t.Helper()
+
+	_, p, err := net.SplitHostPort(addr)
+	if err != nil {
+		t.Fatalf("Invalid serving address: %v", err)
+	}
+	port, err := strconv.ParseUint(p, 10, 32)
+	if err != nil {
+		t.Fatalf("Invalid serving port: %v", err)
+	}
+	return uint32(port)
+}

internal/testutils/parse_url.go

@@ -0,0 +1,34 @@
+/*
+ *
+ * Copyright 2023 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 testutils
+
+import (
+	"fmt"
+	"net/url"
+)
+
+// MustParseURL attempts to parse the provided target using url.Parse()
+// and panics if parsing fails.
+func MustParseURL(target string) *url.URL {
+	u, err := url.Parse(target)
+	if err != nil {
+		panic(fmt.Sprintf("Error parsing target(%s): %v", target, err))
+	}
+	return u
+}

internal/testutils/pickfirst/pickfirst.go

@@ -0,0 +1,82 @@
+/*
+ *
+ * Copyright 2022 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 pickfirst contains helper functions to check for pickfirst load
+// balancing of RPCs in tests.
+package pickfirst
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/resolver"
+
+	testgrpc "google.golang.org/grpc/interop/grpc_testing"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
+)
+
+// CheckRPCsToBackend makes a bunch of RPCs on the given ClientConn and verifies
+// if the RPCs are routed to a peer matching wantAddr.
+//
+// Returns a non-nil error if context deadline expires before all RPCs begin to
+// be routed to the peer matching wantAddr, or if the backend returns RPC errors.
+func CheckRPCsToBackend(ctx context.Context, cc *grpc.ClientConn, wantAddr resolver.Address) error {
+	client := testgrpc.NewTestServiceClient(cc)
+	peer := &peer.Peer{}
+	// Make sure that 20 RPCs in a row reach the expected backend. Some
+	// tests switch from round_robin back to pick_first and call this
+	// function. None of our tests spin up more than 10 backends. So,
+	// waiting for 20 RPCs to reach a single backend would a decent
+	// indicator of having switched to pick_first.
+	count := 0
+	for {
+		time.Sleep(time.Millisecond)
+		if ctx.Err() != nil {
+			return fmt.Errorf("timeout waiting for RPC to be routed to %s", wantAddr.Addr)
+		}
+		if _, err := client.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(peer)); err != nil {
+			// Some tests remove backends and check if pick_first is happening across
+			// the remaining backends. In such cases, RPCs can initially fail on the
+			// connection using the removed backend. Just keep retrying and eventually
+			// the connection using the removed backend will shutdown and will be
+			// removed.
+			continue
+		}
+		if peer.Addr.String() != wantAddr.Addr {
+			count = 0
+			continue
+		}
+		count++
+		if count > 20 {
+			break
+		}
+	}
+	// Make sure subsequent RPCs are all routed to the same backend.
+	for i := 0; i < 10; i++ {
+		if _, err := client.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(peer)); err != nil {
+			return fmt.Errorf("EmptyCall() = %v, want <nil>", err)
+		}
+		if gotAddr := peer.Addr.String(); gotAddr != wantAddr.Addr {
+			return fmt.Errorf("rpc sent to peer %q, want peer %q", gotAddr, wantAddr)
+		}
+	}
+	return nil
+}

internal/testutils/pipe_listener.go

@@ -0,0 +1,96 @@
+/*
+ *
+ * Copyright 2018 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 testutils contains testing helpers.
+package testutils
+
+import (
+	"errors"
+	"net"
+	"time"
+)
+
+var errClosed = errors.New("closed")
+
+type pipeAddr struct{}
+
+func (p pipeAddr) Network() string { return "pipe" }
+func (p pipeAddr) String() string  { return "pipe" }
+
+// PipeListener is a listener with an unbuffered pipe. Each write will complete only once the other side reads. It
+// should only be created using NewPipeListener.
+type PipeListener struct {
+	c    chan chan<- net.Conn
+	done chan struct{}
+}
+
+// NewPipeListener creates a new pipe listener.
+func NewPipeListener() *PipeListener {
+	return &PipeListener{
+		c:    make(chan chan<- net.Conn),
+		done: make(chan struct{}),
+	}
+}
+
+// Accept accepts a connection.
+func (p *PipeListener) Accept() (net.Conn, error) {
+	var connChan chan<- net.Conn
+	select {
+	case <-p.done:
+		return nil, errClosed
+	case connChan = <-p.c:
+		select {
+		case <-p.done:
+			close(connChan)
+			return nil, errClosed
+		default:
+		}
+	}
+	c1, c2 := net.Pipe()
+	connChan <- c1
+	close(connChan)
+	return c2, nil
+}
+
+// Close closes the listener.
+func (p *PipeListener) Close() error {
+	close(p.done)
+	return nil
+}
+
+// Addr returns a pipe addr.
+func (p *PipeListener) Addr() net.Addr {
+	return pipeAddr{}
+}
+
+// Dialer dials a connection.
+func (p *PipeListener) Dialer() func(string, time.Duration) (net.Conn, error) {
+	return func(string, time.Duration) (net.Conn, error) {
+		connChan := make(chan net.Conn)
+		select {
+		case p.c <- connChan:
+		case <-p.done:
+			return nil, errClosed
+		}
+		conn, ok := <-connChan
+		if !ok {
+			return nil, errClosed
+		}
+		return conn, nil
+	}
+}

internal/testutils/pipe_listener_test.go

@@ -0,0 +1,172 @@
+/*
+ *
+ * Copyright 2018 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 testutils_test
+
+import (
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/internal/testutils"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+func (s) TestPipeListener(t *testing.T) {
+	pl := testutils.NewPipeListener()
+	recvdBytes := make(chan []byte, 1)
+	const want = "hello world"
+
+	go func() {
+		c, err := pl.Accept()
+		if err != nil {
+			t.Error(err)
+		}
+
+		read := make([]byte, len(want))
+		_, err = c.Read(read)
+		if err != nil {
+			t.Error(err)
+		}
+		recvdBytes <- read
+	}()
+
+	dl := pl.Dialer()
+	conn, err := dl("", time.Duration(0))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = conn.Write([]byte(want))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	select {
+	case gotBytes := <-recvdBytes:
+		got := string(gotBytes)
+		if got != want {
+			t.Fatalf("expected to get %s, got %s", got, want)
+		}
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("timed out waiting for server to receive bytes")
+	}
+}
+
+func (s) TestUnblocking(t *testing.T) {
+	for _, test := range []struct {
+		desc                 string
+		blockFuncShouldError bool
+		blockFunc            func(*testutils.PipeListener, chan struct{}) error
+		unblockFunc          func(*testutils.PipeListener) error
+	}{
+		{
+			desc: "Accept unblocks Dial",
+			blockFunc: func(pl *testutils.PipeListener, done chan struct{}) error {
+				dl := pl.Dialer()
+				_, err := dl("", time.Duration(0))
+				close(done)
+				return err
+			},
+			unblockFunc: func(pl *testutils.PipeListener) error {
+				_, err := pl.Accept()
+				return err
+			},
+		},
+		{
+			desc:                 "Close unblocks Dial",
+			blockFuncShouldError: true, // because pl.Close will be called
+			blockFunc: func(pl *testutils.PipeListener, done chan struct{}) error {
+				dl := pl.Dialer()
+				_, err := dl("", time.Duration(0))
+				close(done)
+				return err
+			},
+			unblockFunc: func(pl *testutils.PipeListener) error {
+				return pl.Close()
+			},
+		},
+		{
+			desc: "Dial unblocks Accept",
+			blockFunc: func(pl *testutils.PipeListener, done chan struct{}) error {
+				_, err := pl.Accept()
+				close(done)
+				return err
+			},
+			unblockFunc: func(pl *testutils.PipeListener) error {
+				dl := pl.Dialer()
+				_, err := dl("", time.Duration(0))
+				return err
+			},
+		},
+		{
+			desc:                 "Close unblocks Accept",
+			blockFuncShouldError: true, // because pl.Close will be called
+			blockFunc: func(pl *testutils.PipeListener, done chan struct{}) error {
+				_, err := pl.Accept()
+				close(done)
+				return err
+			},
+			unblockFunc: func(pl *testutils.PipeListener) error {
+				return pl.Close()
+			},
+		},
+	} {
+		t.Log(test.desc)
+		testUnblocking(t, test.blockFunc, test.unblockFunc, test.blockFuncShouldError)
+	}
+}
+
+func testUnblocking(t *testing.T, blockFunc func(*testutils.PipeListener, chan struct{}) error, unblockFunc func(*testutils.PipeListener) error, blockFuncShouldError bool) {
+	pl := testutils.NewPipeListener()
+	dialFinished := make(chan struct{})
+
+	go func() {
+		err := blockFunc(pl, dialFinished)
+		if blockFuncShouldError && err == nil {
+			t.Error("expected blocking func to return error because pl.Close was called, but got nil")
+		}
+
+		if !blockFuncShouldError && err != nil {
+			t.Error(err)
+		}
+	}()
+
+	select {
+	case <-dialFinished:
+		t.Fatal("expected Dial to block until pl.Close or pl.Accept")
+	default:
+	}
+
+	if err := unblockFunc(pl); err != nil {
+		t.Fatal(err)
+	}
+
+	select {
+	case <-dialFinished:
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("expected Accept to unblock after pl.Accept was called")
+	}
+}

internal/testutils/proxyserver/proxyserver.go

@@ -0,0 +1,134 @@
+/*
+ *
+ * Copyright 2024 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 proxyserver provides an implementation of a proxy server for testing purposes.
+// The server supports only a single incoming connection at a time and is not concurrent.
+// It handles only HTTP CONNECT requests; other HTTP methods are not supported.
+package proxyserver
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+	"net"
+	"net/http"
+	"testing"
+	"time"
+
+	"google.golang.org/grpc/internal/testutils"
+)
+
+// ProxyServer represents a test proxy server.
+type ProxyServer struct {
+	lis       net.Listener
+	in        net.Conn            // Connection from the client to the proxy.
+	out       net.Conn            // Connection from the proxy to the backend.
+	onRequest func(*http.Request) // Function to check the request sent to proxy.
+	Addr      string              // Address of the proxy
+}
+
+const defaultTestTimeout = 10 * time.Second
+
+// Stop closes the ProxyServer and its connections to client and server.
+func (p *ProxyServer) stop() {
+	p.lis.Close()
+	if p.in != nil {
+		p.in.Close()
+	}
+	if p.out != nil {
+		p.out.Close()
+	}
+}
+
+func (p *ProxyServer) handleRequest(t *testing.T, in net.Conn, waitForServerHello bool) {
+	req, err := http.ReadRequest(bufio.NewReader(in))
+	if err != nil {
+		t.Errorf("failed to read CONNECT req: %v", err)
+		return
+	}
+	if req.Method != http.MethodConnect {
+		t.Errorf("unexpected Method %q, want %q", req.Method, http.MethodConnect)
+	}
+	p.onRequest(req)
+
+	t.Logf("Dialing to %s", req.URL.Host)
+	out, err := net.Dial("tcp", req.URL.Host)
+	if err != nil {
+		in.Close()
+		t.Logf("failed to dial to server: %v", err)
+		return
+	}
+	out.SetDeadline(time.Now().Add(defaultTestTimeout))
+	resp := http.Response{StatusCode: http.StatusOK, Proto: "HTTP/1.0"}
+	var buf bytes.Buffer
+	resp.Write(&buf)
+
+	if waitForServerHello {
+		// Batch the first message from the server with the http connect
+		// response. This is done to test the cases in which the grpc client has
+		// the response to the connect request and proxied packets from the
+		// destination server when it reads the transport.
+		b := make([]byte, 50)
+		bytesRead, err := out.Read(b)
+		if err != nil {
+			t.Errorf("Got error while reading server hello: %v", err)
+			in.Close()
+			out.Close()
+			return
+		}
+		buf.Write(b[0:bytesRead])
+	}
+	p.in = in
+	p.in.Write(buf.Bytes())
+	p.out = out
+
+	go io.Copy(p.in, p.out)
+	go io.Copy(p.out, p.in)
+}
+
+// New initializes and starts a proxy server, registers a cleanup to
+// stop it, and returns a ProxyServer.
+func New(t *testing.T, reqCheck func(*http.Request), waitForServerHello bool) *ProxyServer {
+	t.Helper()
+	pLis, err := testutils.LocalTCPListener()
+	if err != nil {
+		t.Fatalf("failed to listen: %v", err)
+	}
+
+	p := &ProxyServer{
+		lis:       pLis,
+		onRequest: reqCheck,
+		Addr:      pLis.Addr().String(),
+	}
+
+	// Start the proxy server.
+	go func() {
+		for {
+			in, err := p.lis.Accept()
+			if err != nil {
+				return
+			}
+			// p.handleRequest is not invoked in a goroutine because the test
+			// proxy currently supports handling only one connection at a time.
+			p.handleRequest(t, in, waitForServerHello)
+		}
+	}()
+	t.Logf("Started proxy at: %q", pLis.Addr().String())
+	t.Cleanup(p.stop)
+	return p
+}

internal/testutils/resolver.go

@@ -0,0 +1,70 @@
+/*
+ *
+ * Copyright 2023 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 testutils
+
+import (
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/pretty"
+	"google.golang.org/grpc/resolver"
+	"google.golang.org/grpc/serviceconfig"
+)
+
+// Logger wraps the logging methods from testing.T.
+type Logger interface {
+	Log(args ...any)
+	Logf(format string, args ...any)
+	Errorf(format string, args ...any)
+}
+
+// ResolverClientConn is a fake implementation of the resolver.ClientConn
+// interface to be used in tests.
+type ResolverClientConn struct {
+	resolver.ClientConn // Embedding the interface to avoid implementing deprecated methods.
+
+	Logger       Logger                     // Tests should pass testing.T for this.
+	UpdateStateF func(resolver.State) error // Invoked when resolver pushes a state update.
+	ReportErrorF func(err error)            // Invoked when resolver pushes an error.
+}
+
+// UpdateState invokes the test specified callback with the update received from
+// the resolver. If the callback returns a non-nil error, the same will be
+// propagated to the resolver.
+func (t *ResolverClientConn) UpdateState(s resolver.State) error {
+	t.Logger.Logf("testutils.ResolverClientConn: UpdateState(%s)", pretty.ToJSON(s))
+
+	if t.UpdateStateF != nil {
+		return t.UpdateStateF(s)
+	}
+	return nil
+}
+
+// ReportError pushes the error received from the resolver on to ErrorCh.
+func (t *ResolverClientConn) ReportError(err error) {
+	t.Logger.Logf("testutils.ResolverClientConn: ReportError(%v)", err)
+
+	if t.ReportErrorF != nil {
+		t.ReportErrorF(err)
+	}
+}
+
+// ParseServiceConfig parses the provided service by delegating the work to the
+// implementation in the grpc package.
+func (t *ResolverClientConn) ParseServiceConfig(jsonSC string) *serviceconfig.ParseResult {
+	return internal.ParseServiceConfig.(func(string) *serviceconfig.ParseResult)(jsonSC)
+}

internal/testutils/restartable_listener.go

@@ -0,0 +1,102 @@
+/*
+ *
+ * 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 testutils
+
+import (
+	"net"
+	"sync"
+)
+
+type tempError struct{}
+
+func (*tempError) Error() string {
+	return "restartable listener temporary error"
+}
+func (*tempError) Temporary() bool {
+	return true
+}
+
+// RestartableListener wraps a net.Listener and supports stopping and restarting
+// the latter.
+type RestartableListener struct {
+	lis net.Listener
+
+	mu      sync.Mutex
+	stopped bool
+	conns   []net.Conn
+}
+
+// NewRestartableListener returns a new RestartableListener wrapping l.
+func NewRestartableListener(l net.Listener) *RestartableListener {
+	return &RestartableListener{lis: l}
+}
+
+// Accept waits for and returns the next connection to the listener.
+//
+// If the listener is currently not accepting new connections, because `Stop`
+// was called on it, the connection is immediately closed after accepting
+// without any bytes being sent on it.
+func (l *RestartableListener) Accept() (net.Conn, error) {
+	conn, err := l.lis.Accept()
+	if err != nil {
+		return nil, err
+	}
+
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	if l.stopped {
+		conn.Close()
+		return nil, &tempError{}
+	}
+	l.conns = append(l.conns, conn)
+	return conn, nil
+}
+
+// Close closes the listener.
+func (l *RestartableListener) Close() error {
+	return l.lis.Close()
+}
+
+// Addr returns the listener's network address.
+func (l *RestartableListener) Addr() net.Addr {
+	return l.lis.Addr()
+}
+
+// Stop closes existing connections on the listener and prevents new connections
+// from being accepted.
+func (l *RestartableListener) Stop() {
+	logger.Infof("Stopping restartable listener %q", l.Addr())
+
+	l.mu.Lock()
+	l.stopped = true
+	for _, conn := range l.conns {
+		conn.Close()
+	}
+	l.conns = nil
+	l.mu.Unlock()
+}
+
+// Restart gets a previously stopped listener to start accepting connections.
+func (l *RestartableListener) Restart() {
+	logger.Infof("Restarting listener %q", l.Addr())
+
+	l.mu.Lock()
+	l.stopped = false
+	l.mu.Unlock()
+}

internal/testutils/rls/fake_rls_server.go

@@ -0,0 +1,134 @@
+/*
+ *
+ * Copyright 2022 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 rls contains utilities for RouteLookupService e2e tests.
+package rls
+
+import (
+	"context"
+	"net"
+	"sync"
+	"testing"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	rlsgrpc "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+	rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/status"
+)
+
+// RouteLookupResponse wraps an RLS response and the associated error to be sent
+// to a client when the RouteLookup RPC is invoked.
+type RouteLookupResponse struct {
+	Resp *rlspb.RouteLookupResponse
+	Err  error
+}
+
+// SetupFakeRLSServer starts and returns a fake RouteLookupService server
+// listening on the given listener or on a random local port. Also returns a
+// channel for tests to get notified whenever the RouteLookup RPC is invoked on
+// the fake server.
+//
+// This function sets up the fake server to respond with an empty response for
+// the RouteLookup RPCs. Tests can override this by calling the
+// SetResponseCallback() method on the returned fake server.
+func SetupFakeRLSServer(t *testing.T, lis net.Listener, opts ...grpc.ServerOption) (*FakeRouteLookupServer, chan struct{}) {
+	s, cancel := StartFakeRouteLookupServer(t, lis, opts...)
+	t.Logf("Started fake RLS server at %q", s.Address)
+
+	ch := make(chan struct{}, 1)
+	s.SetRequestCallback(func(*rlspb.RouteLookupRequest) {
+		select {
+		case ch <- struct{}{}:
+		default:
+		}
+	})
+	t.Cleanup(cancel)
+	return s, ch
+}
+
+// FakeRouteLookupServer is a fake implementation of the RouteLookupService.
+//
+// It is safe for concurrent use.
+type FakeRouteLookupServer struct {
+	rlsgrpc.UnimplementedRouteLookupServiceServer
+	Address string
+
+	mu     sync.Mutex
+	respCb func(context.Context, *rlspb.RouteLookupRequest) *RouteLookupResponse
+	reqCb  func(*rlspb.RouteLookupRequest)
+}
+
+// StartFakeRouteLookupServer starts a fake RLS server listening for requests on
+// lis. If lis is nil, it creates a new listener on a random local port. The
+// returned cancel function should be invoked by the caller upon completion of
+// the test.
+func StartFakeRouteLookupServer(t *testing.T, lis net.Listener, opts ...grpc.ServerOption) (*FakeRouteLookupServer, func()) {
+	t.Helper()
+
+	if lis == nil {
+		var err error
+		lis, err = testutils.LocalTCPListener()
+		if err != nil {
+			t.Fatalf("net.Listen() failed: %v", err)
+		}
+	}
+
+	s := &FakeRouteLookupServer{Address: lis.Addr().String()}
+	server := grpc.NewServer(opts...)
+	rlsgrpc.RegisterRouteLookupServiceServer(server, s)
+	go server.Serve(lis)
+	return s, func() { server.Stop() }
+}
+
+// RouteLookup implements the RouteLookupService.
+func (s *FakeRouteLookupServer) RouteLookup(ctx context.Context, req *rlspb.RouteLookupRequest) (*rlspb.RouteLookupResponse, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.reqCb != nil {
+		s.reqCb(req)
+	}
+	if err := ctx.Err(); err != nil {
+		return nil, status.Error(codes.DeadlineExceeded, err.Error())
+	}
+	if s.respCb == nil {
+		return &rlspb.RouteLookupResponse{}, nil
+	}
+	resp := s.respCb(ctx, req)
+	return resp.Resp, resp.Err
+}
+
+// SetResponseCallback sets a callback to be invoked on every RLS request. If
+// this callback is set, the response returned by the fake server depends on the
+// value returned by the callback. If this callback is not set, the fake server
+// responds with an empty response.
+func (s *FakeRouteLookupServer) SetResponseCallback(f func(context.Context, *rlspb.RouteLookupRequest) *RouteLookupResponse) {
+	s.mu.Lock()
+	s.respCb = f
+	s.mu.Unlock()
+}
+
+// SetRequestCallback sets a callback to be invoked on every RLS request. The
+// callback is given the incoming request, and tests can use this to verify that
+// the request matches its expectations.
+func (s *FakeRouteLookupServer) SetRequestCallback(f func(*rlspb.RouteLookupRequest)) {
+	s.mu.Lock()
+	s.reqCb = f
+	s.mu.Unlock()
+}

internal/testutils/roundrobin/roundrobin.go

@@ -0,0 +1,286 @@
+/*
+ *
+ * Copyright 2022 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 roundrobin contains helper functions to check for roundrobin and
+// weighted-roundrobin load balancing of RPCs in tests.
+package roundrobin
+
+import (
+	"context"
+	"fmt"
+	"math"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"gonum.org/v1/gonum/stat/distuv"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/grpclog"
+	"google.golang.org/grpc/peer"
+	"google.golang.org/grpc/resolver"
+
+	testgrpc "google.golang.org/grpc/interop/grpc_testing"
+	testpb "google.golang.org/grpc/interop/grpc_testing"
+)
+
+var logger = grpclog.Component("testutils-roundrobin")
+
+// waitForTrafficToReachBackends repeatedly makes RPCs using the provided
+// TestServiceClient until RPCs reach all backends specified in addrs, or the
+// context expires, in which case a non-nil error is returned.
+func waitForTrafficToReachBackends(ctx context.Context, client testgrpc.TestServiceClient, addrs []resolver.Address) error {
+	// Make sure connections to all backends are up. We need to do this two
+	// times (to be sure that round_robin has kicked in) because the channel
+	// could have been configured with a different LB policy before the switch
+	// to round_robin. And the previous LB policy could be sharing backends with
+	// round_robin, and therefore in the first iteration of this loop, RPCs
+	// could land on backends owned by the previous LB policy.
+	for j := 0; j < 2; j++ {
+		for i := 0; i < len(addrs); i++ {
+			for {
+				time.Sleep(time.Millisecond)
+				if ctx.Err() != nil {
+					return fmt.Errorf("timeout waiting for connection to %q to be up", addrs[i].Addr)
+				}
+				var peer peer.Peer
+				if _, err := client.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer)); err != nil {
+					// Some tests remove backends and check if round robin is
+					// happening across the remaining backends. In such cases,
+					// RPCs can initially fail on the connection using the
+					// removed backend. Just keep retrying and eventually the
+					// connection using the removed backend will shutdown and
+					// will be removed.
+					continue
+				}
+				if peer.Addr.String() == addrs[i].Addr {
+					break
+				}
+			}
+		}
+	}
+	return nil
+}
+
+// CheckRoundRobinRPCs verifies that EmptyCall RPCs on the given ClientConn,
+// connected to a server exposing the test.grpc_testing.TestService, are
+// roundrobined across the given backend addresses.
+//
+// Returns a non-nil error if context deadline expires before RPCs start to get
+// roundrobined across the given backends.
+func CheckRoundRobinRPCs(ctx context.Context, client testgrpc.TestServiceClient, addrs []resolver.Address) error {
+	if err := waitForTrafficToReachBackends(ctx, client, addrs); err != nil {
+		return err
+	}
+
+	// At this point, RPCs are getting successfully executed at the backends
+	// that we care about. To support duplicate addresses (in addrs) and
+	// backends being removed from the list of addresses passed to the
+	// roundrobin LB, we do the following:
+	// 1. Determine the count of RPCs that we expect each of our backends to
+	//    receive per iteration.
+	// 2. Wait until the same pattern repeats a few times, or the context
+	//    deadline expires.
+	wantAddrCount := make(map[string]int)
+	for _, addr := range addrs {
+		wantAddrCount[addr.Addr]++
+	}
+	for ; ctx.Err() == nil; <-time.After(time.Millisecond) {
+		// Perform 3 more iterations.
+		var iterations [][]string
+		for i := 0; i < 3; i++ {
+			iteration := make([]string, len(addrs))
+			for c := 0; c < len(addrs); c++ {
+				var peer peer.Peer
+				if _, err := client.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer)); err != nil {
+					return fmt.Errorf("EmptyCall() = %v, want <nil>", err)
+				}
+				iteration[c] = peer.Addr.String()
+			}
+			iterations = append(iterations, iteration)
+		}
+		// Ensure the first iteration contains all addresses in addrs.
+		gotAddrCount := make(map[string]int)
+		for _, addr := range iterations[0] {
+			gotAddrCount[addr]++
+		}
+		if diff := cmp.Diff(gotAddrCount, wantAddrCount); diff != "" {
+			logger.Infof("non-roundrobin, got address count in one iteration: %v, want: %v, Diff: %s", gotAddrCount, wantAddrCount, diff)
+			continue
+		}
+		// Ensure all three iterations contain the same addresses.
+		if !cmp.Equal(iterations[0], iterations[1]) || !cmp.Equal(iterations[0], iterations[2]) {
+			logger.Infof("non-roundrobin, first iter: %v, second iter: %v, third iter: %v", iterations[0], iterations[1], iterations[2])
+			continue
+		}
+		return nil
+	}
+	return fmt.Errorf("timeout when waiting for roundrobin distribution of RPCs across addresses: %v", addrs)
+}
+
+// CheckWeightedRoundRobinRPCs verifies that EmptyCall RPCs on the given
+// ClientConn, connected to a server exposing the test.grpc_testing.TestService,
+// are weighted roundrobined (with randomness) across the given backend
+// addresses.
+//
+// Returns a non-nil error if context deadline expires before RPCs start to get
+// roundrobined across the given backends.
+func CheckWeightedRoundRobinRPCs(ctx context.Context, t *testing.T, client testgrpc.TestServiceClient, addrs []resolver.Address) error {
+	if err := waitForTrafficToReachBackends(ctx, client, addrs); err != nil {
+		return err
+	}
+
+	// At this point, RPCs are getting successfully executed at the backends
+	// that we care about. To take the randomness of the WRR into account, we
+	// look for approximate distribution instead of exact.
+	wantAddrCount := make(map[string]int)
+	for _, addr := range addrs {
+		wantAddrCount[addr.Addr]++
+	}
+	attemptCount := attemptCounts(wantAddrCount)
+
+	expectedCount := make(map[string]float64)
+	for addr, count := range wantAddrCount {
+		expectedCount[addr] = float64(count) / float64(len(addrs)) * float64(attemptCount)
+	}
+
+	// There is a small possibility that RPCs are reaching backends that we
+	// don't expect them to reach here. The can happen because:
+	// - at time T0, the list of backends [A, B, C, D].
+	// - at time T1, the test updates the list of backends to [A, B, C], and
+	//   immediately starts attempting to check the distribution of RPCs to the
+	//   new backends.
+	// - there is no way for the test to wait for a new picker to be pushed on
+	//   to the channel (which contains the updated list of backends) before
+	//   starting to attempt the RPC distribution checks.
+	// - This is usually a transitory state and will eventually fix itself when
+	//   the new picker is pushed on the channel, and RPCs will start getting
+	//   routed to only backends that we care about.
+	//
+	// We work around this situation by using two loops. The inner loop contains
+	// the meat of the calculations, and includes the logic which factors out
+	// the randomness in weighted roundrobin. If we ever see an RPCs getting
+	// routed to a backend that we don't expect it to get routed to, we break
+	// from the inner loop thereby resetting all state and start afresh.
+	for {
+		observedCount := make(map[string]float64)
+	InnerLoop:
+		for {
+			if ctx.Err() != nil {
+				return fmt.Errorf("timeout when waiting for roundrobin distribution of RPCs across addresses: %v", addrs)
+			}
+			for i := 0; i < attemptCount; i++ {
+				var peer peer.Peer
+				if _, err := client.EmptyCall(ctx, &testpb.Empty{}, grpc.Peer(&peer)); err != nil {
+					return fmt.Errorf("EmptyCall() = %v, want <nil>", err)
+				}
+				if addr := peer.Addr.String(); wantAddrCount[addr] == 0 {
+					break InnerLoop
+				}
+				observedCount[peer.Addr.String()]++
+			}
+
+			return pearsonsChiSquareTest(t, observedCount, expectedCount)
+		}
+		<-time.After(time.Millisecond)
+	}
+}
+
+// attemptCounts returns the number of attempts needed to verify the expected
+// distribution of RPCs. Having more attempts per category will give the test
+// a greater ability to detect a small but real deviation from the expected
+// distribution.
+func attemptCounts(wantAddrWeights map[string]int) int {
+	if len(wantAddrWeights) == 0 {
+		return 0
+	}
+
+	totalWeight := 0
+	minWeight := -1
+
+	for _, weight := range wantAddrWeights {
+		totalWeight += weight
+		if minWeight == -1 || weight < minWeight {
+			minWeight = weight
+		}
+	}
+
+	minRatio := float64(minWeight) / float64(totalWeight)
+
+	// We want the expected count for the smallest category to be at least 500.
+	// ExpectedCount = TotalAttempts * minRatio
+	// So, 500 <= TotalAttempts * minRatio
+	// which means TotalAttempts >= 500 / minRatio
+	const minExpectedPerCategory = 500.0
+	requiredAttempts := minExpectedPerCategory / minRatio
+
+	return int(math.Ceil(requiredAttempts))
+}
+
+// pearsonsChiSquareTest checks if the observed counts match the expected
+// counts.
+// Pearson's Chi-Squared Test Formula:
+//
+//	χ² = ∑ (Oᵢ - Eᵢ)² / Eᵢ
+//
+// Where:
+//   - χ² is the chi-square statistic
+//   - Oᵢ is the observed count for category i
+//   - Eᵢ is the expected count for category i
+//   - The sum is over all categories (i = 1 to k)
+//
+// This tests how well the observed distribution matches the expected one.
+// Larger χ² values indicate a greater difference between observed and expected
+// counts. The p-value is computed as P(χ² ≥ computed value) under the
+// chi-square distribution with degrees of freedom:
+// df = number of categories - 1
+// See https://en.wikipedia.org/wiki/Pearson%27s_chi-squared_test for more
+// details.
+func pearsonsChiSquareTest(t *testing.T, observedCounts, expectedCounts map[string]float64) error {
+	chiSquaredStat := 0.0
+	for addr, want := range expectedCounts {
+		got := observedCounts[addr]
+		chiSquaredStat += (got - want) * (got - want) / want
+	}
+	degreesOfFreedom := len(expectedCounts) - 1
+	const alpha = 1e-6
+	chiSquareDist := distuv.ChiSquared{K: float64(degreesOfFreedom)}
+	pValue := chiSquareDist.Survival(chiSquaredStat)
+	t.Logf("Observed ratio: %v", observedCounts)
+	t.Logf("Expected ratio: %v", expectedCounts)
+	t.Logf("χ² statistic: %.4f", chiSquaredStat)
+	t.Logf("Degrees of freedom: %d\n", degreesOfFreedom)
+	t.Logf("p-value: %.10f\n", pValue)
+	// Alpha (α) is the threshold we use to decide if the test "fails".
+	// It controls how sensitive the chi-square test is.
+	//
+	// A smaller alpha means we require stronger evidence to say the load
+	// balancing is wrong. A larger alpha makes the test more likely to fail,
+	// even for small random fluctuations.
+	//
+	// For CI, we set alpha = 1e-6 to avoid flaky test failures.
+	// That means:
+	//  - There's only a 1-in-a-million chance the test fails due to random
+	//    variation, assuming the load balancer is working correctly.
+	//  - If the test *does* fail, it's very likely there's a real bug.
+	//
+	// TL;DR: smaller alpha = stricter test, fewer false alarms.
+	if pValue > alpha {
+		return nil
+	}
+	return fmt.Errorf("observed distribution significantly differs from expectations, observeredCounts: %v, expectedCounts: %v", observedCounts, expectedCounts)
+}

internal/testutils/state.go

@@ -0,0 +1,85 @@
+/*
+ *
+ * Copyright 2023 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 testutils
+
+import (
+	"context"
+	"testing"
+
+	"google.golang.org/grpc/connectivity"
+)
+
+// A StateChanger reports state changes, e.g. a grpc.ClientConn.
+type StateChanger interface {
+	// Connect begins connecting the StateChanger.
+	Connect()
+	// GetState returns the current state of the StateChanger.
+	GetState() connectivity.State
+	// WaitForStateChange returns true when the state becomes s, or returns
+	// false if ctx is canceled first.
+	WaitForStateChange(ctx context.Context, s connectivity.State) bool
+}
+
+// StayConnected makes sc stay connected by repeatedly calling sc.Connect()
+// until the state becomes Shutdown or until the context expires.
+func StayConnected(ctx context.Context, sc StateChanger) {
+	for {
+		state := sc.GetState()
+		switch state {
+		case connectivity.Idle:
+			sc.Connect()
+		case connectivity.Shutdown:
+			return
+		}
+		if !sc.WaitForStateChange(ctx, state) {
+			return
+		}
+	}
+}
+
+// AwaitState waits for sc to enter stateWant or fatal errors if it doesn't
+// happen before ctx expires.
+func AwaitState(ctx context.Context, t *testing.T, sc StateChanger, stateWant connectivity.State) {
+	t.Helper()
+	for state := sc.GetState(); state != stateWant; state = sc.GetState() {
+		if !sc.WaitForStateChange(ctx, state) {
+			t.Fatalf("Timed out waiting for state change.  got %v; want %v", state, stateWant)
+		}
+	}
+}
+
+// AwaitNotState waits for sc to leave stateDoNotWant or fatal errors if it
+// doesn't happen before ctx expires.
+func AwaitNotState(ctx context.Context, t *testing.T, sc StateChanger, stateDoNotWant connectivity.State) {
+	t.Helper()
+	for state := sc.GetState(); state == stateDoNotWant; state = sc.GetState() {
+		if !sc.WaitForStateChange(ctx, state) {
+			t.Fatalf("Timed out waiting for state change.  got %v; want NOT %v", state, stateDoNotWant)
+		}
+	}
+}
+
+// AwaitNoStateChange expects ctx to be canceled before sc's state leaves
+// currState, and fatal errors otherwise.
+func AwaitNoStateChange(ctx context.Context, t *testing.T, sc StateChanger, currState connectivity.State) {
+	t.Helper()
+	if sc.WaitForStateChange(ctx, currState) {
+		t.Fatalf("State changed from %q to %q when no state change was expected", currState, sc.GetState())
+	}
+}

internal/testutils/stats/test_metrics_recorder.go

@@ -0,0 +1,323 @@
+/*
+ *
+ * Copyright 2024 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 stats implements a TestMetricsRecorder utility.
+package stats
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	"github.com/google/go-cmp/cmp"
+	estats "google.golang.org/grpc/experimental/stats"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/stats"
+)
+
+// TestMetricsRecorder is a MetricsRecorder to be used in tests. It sends
+// recording events on channels and provides helpers to check if certain events
+// have taken place. It also persists metrics data keyed on the metrics
+// descriptor.
+type TestMetricsRecorder struct {
+	estats.UnimplementedMetricsRecorder
+	intCountCh       *testutils.Channel
+	floatCountCh     *testutils.Channel
+	intHistoCh       *testutils.Channel
+	floatHistoCh     *testutils.Channel
+	intGaugeCh       *testutils.Channel
+	intUpDownCountCh *testutils.Channel
+
+	// mu protects data.
+	mu sync.Mutex
+	// data is the most recent update for each metric name.
+	data map[string]float64
+}
+
+// NewTestMetricsRecorder returns a new TestMetricsRecorder.
+func NewTestMetricsRecorder() *TestMetricsRecorder {
+	return &TestMetricsRecorder{
+		intCountCh:       testutils.NewChannelWithSize(10),
+		floatCountCh:     testutils.NewChannelWithSize(10),
+		intHistoCh:       testutils.NewChannelWithSize(10),
+		floatHistoCh:     testutils.NewChannelWithSize(10),
+		intGaugeCh:       testutils.NewChannelWithSize(10),
+		intUpDownCountCh: testutils.NewChannelWithSize(10),
+
+		data: make(map[string]float64),
+	}
+}
+
+// Metric returns the most recent data for a metric, and whether this recorder
+// has received data for a metric.
+func (r *TestMetricsRecorder) Metric(name string) (float64, bool) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	data, ok := r.data[name]
+	return data, ok
+}
+
+// ClearMetrics clears the metrics data store of the test metrics recorder.
+func (r *TestMetricsRecorder) ClearMetrics() {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.data = make(map[string]float64)
+}
+
+// MetricsData represents data associated with a metric.
+type MetricsData struct {
+	Handle *estats.MetricDescriptor
+
+	// Only set based on the type of metric. So only one of IntIncr or FloatIncr
+	// is set.
+	IntIncr   int64
+	FloatIncr float64
+
+	LabelKeys []string
+	LabelVals []string
+}
+
+// WaitForInt64Count waits for an int64 count metric to be recorded and verifies
+// that the recorded metrics data matches the expected metricsDataWant. Returns
+// an error if failed to wait or received wrong data.
+func (r *TestMetricsRecorder) WaitForInt64Count(ctx context.Context, metricsDataWant MetricsData) error {
+	got, err := r.intCountCh.Receive(ctx)
+	if err != nil {
+		return fmt.Errorf("timeout waiting for int64Count")
+	}
+	metricsDataGot := got.(MetricsData)
+	if diff := cmp.Diff(metricsDataGot, metricsDataWant); diff != "" {
+		return fmt.Errorf("int64count metricsData received unexpected value (-got, +want): %v", diff)
+	}
+	return nil
+}
+
+// WaitForInt64CountIncr waits for an int64 count metric to be recorded and
+// verifies that the recorded metrics data incr matches the expected incr.
+// Returns an error if failed to wait or received wrong data.
+func (r *TestMetricsRecorder) WaitForInt64CountIncr(ctx context.Context, incrWant int64) error {
+	got, err := r.intCountCh.Receive(ctx)
+	if err != nil {
+		return fmt.Errorf("timeout waiting for int64Count")
+	}
+	metricsDataGot := got.(MetricsData)
+	if diff := cmp.Diff(metricsDataGot.IntIncr, incrWant); diff != "" {
+		return fmt.Errorf("int64count metricsData received unexpected value (-got, +want): %v", diff)
+	}
+	return nil
+}
+
+// RecordInt64Count sends the metrics data to the intCountCh channel and updates
+// the internal data map with the recorded value.
+func (r *TestMetricsRecorder) RecordInt64Count(handle *estats.Int64CountHandle, incr int64, labels ...string) {
+	r.intCountCh.ReceiveOrFail()
+	r.intCountCh.Send(MetricsData{
+		Handle:    handle.Descriptor(),
+		IntIncr:   incr,
+		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
+		LabelVals: labels,
+	})
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.data[handle.Name] = float64(incr)
+}
+
+// RecordInt64UpDownCount sends the metrics data to the intUpDownCountCh channel and updates
+// the internal data map with the recorded value.
+func (r *TestMetricsRecorder) RecordInt64UpDownCount(handle *estats.Int64UpDownCountHandle, incr int64, labels ...string) {
+	r.intUpDownCountCh.ReceiveOrFail()
+	r.intUpDownCountCh.Send(MetricsData{
+		Handle:    handle.Descriptor(),
+		IntIncr:   incr,
+		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
+		LabelVals: labels,
+	})
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.data[handle.Name] = float64(incr)
+}
+
+// WaitForFloat64Count waits for a float count metric to be recorded and
+// verifies that the recorded metrics data matches the expected metricsDataWant.
+// Returns an error if failed to wait or received wrong data.
+func (r *TestMetricsRecorder) WaitForFloat64Count(ctx context.Context, metricsDataWant MetricsData) error {
+	got, err := r.floatCountCh.Receive(ctx)
+	if err != nil {
+		return fmt.Errorf("timeout waiting for float64Count")
+	}
+	metricsDataGot := got.(MetricsData)
+	if diff := cmp.Diff(metricsDataGot, metricsDataWant); diff != "" {
+		return fmt.Errorf("float64count metricsData received unexpected value (-got, +want): %v", diff)
+	}
+	return nil
+}
+
+// RecordFloat64Count sends the metrics data to the floatCountCh channel and
+// updates the internal data map with the recorded value.
+func (r *TestMetricsRecorder) RecordFloat64Count(handle *estats.Float64CountHandle, incr float64, labels ...string) {
+	r.floatCountCh.ReceiveOrFail()
+	r.floatCountCh.Send(MetricsData{
+		Handle:    handle.Descriptor(),
+		FloatIncr: incr,
+		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
+		LabelVals: labels,
+	})
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.data[handle.Name] = incr
+}
+
+// WaitForInt64Histo waits for an int histo metric to be recorded and verifies
+// that the recorded metrics data matches the expected metricsDataWant. Returns
+// an error if failed to wait or received wrong data.
+func (r *TestMetricsRecorder) WaitForInt64Histo(ctx context.Context, metricsDataWant MetricsData) error {
+	got, err := r.intHistoCh.Receive(ctx)
+	if err != nil {
+		return fmt.Errorf("timeout waiting for int64Histo")
+	}
+	metricsDataGot := got.(MetricsData)
+	if diff := cmp.Diff(metricsDataGot, metricsDataWant); diff != "" {
+		return fmt.Errorf("int64Histo metricsData received unexpected value (-got, +want): %v", diff)
+	}
+	return nil
+}
+
+// RecordInt64Histo sends the metrics data to the intHistoCh channel and updates
+// the internal data map with the recorded value.
+func (r *TestMetricsRecorder) RecordInt64Histo(handle *estats.Int64HistoHandle, incr int64, labels ...string) {
+	r.intHistoCh.ReceiveOrFail()
+	r.intHistoCh.Send(MetricsData{
+		Handle:    handle.Descriptor(),
+		IntIncr:   incr,
+		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
+		LabelVals: labels,
+	})
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.data[handle.Name] = float64(incr)
+}
+
+// WaitForFloat64Histo waits for a float histo metric to be recorded and
+// verifies that the recorded metrics data matches the expected metricsDataWant.
+// Returns an error if failed to wait or received wrong data.
+func (r *TestMetricsRecorder) WaitForFloat64Histo(ctx context.Context, metricsDataWant MetricsData) error {
+	got, err := r.floatHistoCh.Receive(ctx)
+	if err != nil {
+		return fmt.Errorf("timeout waiting for float64Histo")
+	}
+	metricsDataGot := got.(MetricsData)
+	if diff := cmp.Diff(metricsDataGot, metricsDataWant); diff != "" {
+		return fmt.Errorf("float64Histo metricsData received unexpected value (-got, +want): %v", diff)
+	}
+	return nil
+}
+
+// RecordFloat64Histo sends the metrics data to the floatHistoCh channel and
+// updates the internal data map with the recorded value.
+func (r *TestMetricsRecorder) RecordFloat64Histo(handle *estats.Float64HistoHandle, incr float64, labels ...string) {
+	r.floatHistoCh.ReceiveOrFail()
+	r.floatHistoCh.Send(MetricsData{
+		Handle:    handle.Descriptor(),
+		FloatIncr: incr,
+		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
+		LabelVals: labels,
+	})
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.data[handle.Name] = incr
+}
+
+// WaitForInt64Gauge waits for a int gauge metric to be recorded and verifies
+// that the recorded metrics data matches the expected metricsDataWant.
+func (r *TestMetricsRecorder) WaitForInt64Gauge(ctx context.Context, metricsDataWant MetricsData) error {
+	got, err := r.intGaugeCh.Receive(ctx)
+	if err != nil {
+		return fmt.Errorf("timeout waiting for int64Gauge")
+	}
+	metricsDataGot := got.(MetricsData)
+	if diff := cmp.Diff(metricsDataGot, metricsDataWant); diff != "" {
+		return fmt.Errorf("int64Gauge metricsData received unexpected value (-got, +want): %v", diff)
+	}
+	return nil
+}
+
+// RecordInt64Gauge sends the metrics data to the intGaugeCh channel and updates
+// the internal data map with the recorded value.
+func (r *TestMetricsRecorder) RecordInt64Gauge(handle *estats.Int64GaugeHandle, incr int64, labels ...string) {
+	r.intGaugeCh.ReceiveOrFail()
+	r.intGaugeCh.Send(MetricsData{
+		Handle:    handle.Descriptor(),
+		IntIncr:   incr,
+		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
+		LabelVals: labels,
+	})
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.data[handle.Name] = float64(incr)
+}
+
+// To implement a estats.AsyncMetricsRecorder, which allows it to be used in async metrics:
+
+// RecordInt64AsyncGauge sends the metrics data to the intGaugeCh channel and updates
+// the internal data map with the recorded value.
+func (r *TestMetricsRecorder) RecordInt64AsyncGauge(handle *estats.Int64AsyncGaugeHandle, incr int64, labels ...string) {
+	r.intGaugeCh.ReceiveOrFail()
+	r.intGaugeCh.Send(MetricsData{
+		Handle:    handle.Descriptor(),
+		IntIncr:   incr,
+		LabelKeys: append(handle.Labels, handle.OptionalLabels...),
+		LabelVals: labels,
+	})
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.data[handle.Name] = float64(incr)
+}
+
+// ReadFloat64Histo waits for a float64 histogram metric to be recorded and returns its data.
+func (r *TestMetricsRecorder) ReadFloat64Histo(ctx context.Context) (MetricsData, error) {
+	got, err := r.floatHistoCh.Receive(ctx)
+	if err != nil {
+		return MetricsData{}, err
+	}
+	return got.(MetricsData), nil
+}
+
+// To implement a stats.Handler, which allows it to be set as a dial option:
+
+// TagRPC is TestMetricsRecorder's implementation of TagRPC.
+func (r *TestMetricsRecorder) TagRPC(ctx context.Context, _ *stats.RPCTagInfo) context.Context {
+	return ctx
+}
+
+// HandleRPC is TestMetricsRecorder's implementation of HandleRPC.
+func (r *TestMetricsRecorder) HandleRPC(context.Context, stats.RPCStats) {}
+
+// TagConn is TestMetricsRecorder's implementation of TagConn.
+func (r *TestMetricsRecorder) TagConn(ctx context.Context, _ *stats.ConnTagInfo) context.Context {
+	return ctx
+}
+
+// HandleConn is TestMetricsRecorder's implementation of HandleConn.
+func (r *TestMetricsRecorder) HandleConn(context.Context, stats.ConnStats) {}

internal/testutils/status_equal.go

@@ -0,0 +1,38 @@
+/*
+ *
+ * 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 testutils
+
+import (
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/proto"
+)
+
+// StatusErrEqual returns true iff both err1 and err2 wrap status.Status errors
+// and their underlying status protos are equal.
+func StatusErrEqual(err1, err2 error) bool {
+	status1, ok := status.FromError(err1)
+	if !ok {
+		return false
+	}
+	status2, ok := status.FromError(err2)
+	if !ok {
+		return false
+	}
+	return proto.Equal(status1.Proto(), status2.Proto())
+}

internal/testutils/status_equal_test.go

@@ -0,0 +1,66 @@
+/*
+ *
+ * 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 testutils
+
+import (
+	"testing"
+
+	spb "google.golang.org/genproto/googleapis/rpc/status"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/internal/grpctest"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/types/known/anypb"
+)
+
+type s struct {
+	grpctest.Tester
+}
+
+func Test(t *testing.T) {
+	grpctest.RunSubTests(t, s{})
+}
+
+var statusErr = status.ErrorProto(&spb.Status{
+	Code:    int32(codes.DataLoss),
+	Message: "error for testing",
+	Details: []*anypb.Any{{
+		TypeUrl: "url",
+		Value:   []byte{6, 0, 0, 6, 1, 3},
+	}},
+})
+
+func (s) TestStatusErrEqual(t *testing.T) {
+	tests := []struct {
+		name      string
+		err1      error
+		err2      error
+		wantEqual bool
+	}{
+		{"nil errors", nil, nil, true},
+		{"equal OK status", status.New(codes.OK, "").Err(), status.New(codes.OK, "").Err(), true},
+		{"equal status errors", statusErr, statusErr, true},
+		{"different status errors", statusErr, status.New(codes.OK, "").Err(), false},
+	}
+
+	for _, test := range tests {
+		if gotEqual := StatusErrEqual(test.err1, test.err2); gotEqual != test.wantEqual {
+			t.Errorf("%v: StatusErrEqual(%v, %v) = %v, want %v", test.name, test.err1, test.err2, gotEqual, test.wantEqual)
+		}
+	}
+}

internal/testutils/stubstatshandler.go

@@ -0,0 +1,65 @@
+/*
+ *
+ * Copyright 2023 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 testutils
+
+import (
+	"context"
+
+	"google.golang.org/grpc/stats"
+)
+
+// StubStatsHandler is a stats handler that is easy to customize within
+// individual test cases. It is a stubbable implementation of
+// google.golang.org/grpc/stats.Handler for testing purposes.
+type StubStatsHandler struct {
+	TagRPCF     func(ctx context.Context, info *stats.RPCTagInfo) context.Context
+	HandleRPCF  func(ctx context.Context, info stats.RPCStats)
+	TagConnF    func(ctx context.Context, info *stats.ConnTagInfo) context.Context
+	HandleConnF func(ctx context.Context, info stats.ConnStats)
+}
+
+// TagRPC calls the StubStatsHandler's TagRPCF, if set.
+func (ssh *StubStatsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
+	if ssh.TagRPCF != nil {
+		return ssh.TagRPCF(ctx, info)
+	}
+	return ctx
+}
+
+// HandleRPC calls the StubStatsHandler's HandleRPCF, if set.
+func (ssh *StubStatsHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
+	if ssh.HandleRPCF != nil {
+		ssh.HandleRPCF(ctx, rs)
+	}
+}
+
+// TagConn calls the StubStatsHandler's TagConnF, if set.
+func (ssh *StubStatsHandler) TagConn(ctx context.Context, info *stats.ConnTagInfo) context.Context {
+	if ssh.TagConnF != nil {
+		return ssh.TagConnF(ctx, info)
+	}
+	return ctx
+}
+
+// HandleConn calls the StubStatsHandler's HandleConnF, if set.
+func (ssh *StubStatsHandler) HandleConn(ctx context.Context, cs stats.ConnStats) {
+	if ssh.HandleConnF != nil {
+		ssh.HandleConnF(ctx, cs)
+	}
+}

internal/testutils/tls_creds.go

@@ -0,0 +1,156 @@
+/*
+ *
+ * Copyright 2024 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 testutils
+
+import (
+	"crypto/tls"
+	"crypto/x509"
+	"os"
+	"testing"
+
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/testdata"
+)
+
+// CreateClientTLSCredentials creates client-side TLS transport credentials
+// using certificate and key files from testdata/x509 directory.
+func CreateClientTLSCredentials(t *testing.T) credentials.TransportCredentials {
+	t.Helper()
+
+	cert, err := tls.LoadX509KeyPair(testdata.Path("x509/client1_cert.pem"), testdata.Path("x509/client1_key.pem"))
+	if err != nil {
+		t.Fatalf("tls.LoadX509KeyPair(x509/client1_cert.pem, x509/client1_key.pem) failed: %v", err)
+	}
+	b, err := os.ReadFile(testdata.Path("x509/server_ca_cert.pem"))
+	if err != nil {
+		t.Fatalf("os.ReadFile(x509/server_ca_cert.pem) failed: %v", err)
+	}
+	roots := x509.NewCertPool()
+	if !roots.AppendCertsFromPEM(b) {
+		t.Fatal("Failed to append certificates")
+	}
+	return credentials.NewTLS(&tls.Config{
+		Certificates: []tls.Certificate{cert},
+		RootCAs:      roots,
+		ServerName:   "x.test.example.com",
+	})
+}
+
+// CreateServerTLSCredentials creates server-side TLS transport credentials
+// using certificate and key files from testdata/x509 directory.
+func CreateServerTLSCredentials(t *testing.T, clientAuth tls.ClientAuthType) credentials.TransportCredentials {
+	t.Helper()
+
+	cert, err := tls.LoadX509KeyPair(testdata.Path("x509/server1_cert.pem"), testdata.Path("x509/server1_key.pem"))
+	if err != nil {
+		t.Fatalf("tls.LoadX509KeyPair(x509/server1_cert.pem, x509/server1_key.pem) failed: %v", err)
+	}
+	b, err := os.ReadFile(testdata.Path("x509/client_ca_cert.pem"))
+	if err != nil {
+		t.Fatalf("os.ReadFile(x509/client_ca_cert.pem) failed: %v", err)
+	}
+	ca := x509.NewCertPool()
+	if !ca.AppendCertsFromPEM(b) {
+		t.Fatal("Failed to append certificates")
+	}
+	return credentials.NewTLS(&tls.Config{
+		ClientAuth:   clientAuth,
+		Certificates: []tls.Certificate{cert},
+		ClientCAs:    ca,
+	})
+}
+
+// CreateServerTLSCredentialsCompatibleWithSPIFFE creates server-side TLS
+// transport credentials using certificate and key files from the
+// testdata/spiffe_end2end directory. These credentials are compatible with the
+// SPIFFE trust bundles used on the client side.
+func CreateServerTLSCredentialsCompatibleWithSPIFFE(t *testing.T, clientAuth tls.ClientAuthType) credentials.TransportCredentials {
+	t.Helper()
+
+	cert, err := tls.LoadX509KeyPair(testdata.Path("spiffe_end2end/server_spiffe.pem"), testdata.Path("spiffe_end2end/server.key"))
+	if err != nil {
+		t.Fatalf("tls.LoadX509KeyPair(spiffe_end2end/server_spiffe.pem, spiffe_end2end/server.key) failed: %v", err)
+	}
+	b, err := os.ReadFile(testdata.Path("spiffe_end2end/ca.pem"))
+	if err != nil {
+		t.Fatalf("os.ReadFile(spiffe_end2end/ca.pem) failed: %v", err)
+	}
+	ca := x509.NewCertPool()
+	if !ca.AppendCertsFromPEM(b) {
+		t.Fatal("Failed to append certificates")
+	}
+	return credentials.NewTLS(&tls.Config{
+		ClientAuth:   clientAuth,
+		Certificates: []tls.Certificate{cert},
+		ClientCAs:    ca,
+	})
+}
+
+// CreateServerTLSCredentialsCompatibleWithSPIFFEChain creates server-side TLS
+// transport credentials using a certificate chain and key files from the
+// testdata/spiffe_end2end directory. These credentials are compatible with the
+// SPIFFE trust bundles used on the client side.
+func CreateServerTLSCredentialsCompatibleWithSPIFFEChain(t *testing.T, clientAuth tls.ClientAuthType) credentials.TransportCredentials {
+	t.Helper()
+
+	certs, err := tls.LoadX509KeyPair(testdata.Path("spiffe_end2end/leaf_and_intermediate_chain.pem"), testdata.Path("spiffe_end2end/leaf_signed_by_intermediate.key"))
+	if err != nil {
+		t.Fatalf("tls.LoadX509KeyPair(spiffe_end2end/leaf_and_intermediate_chain.pem, spiffe_end2end/leaf_signed_by_intermediate.key) failed: %v", err)
+	}
+	b, err := os.ReadFile(testdata.Path("spiffe_end2end/ca.pem"))
+	if err != nil {
+		t.Fatalf("os.ReadFile(spiffe_end2end/ca.pem) failed: %v", err)
+	}
+	ca := x509.NewCertPool()
+	if !ca.AppendCertsFromPEM(b) {
+		t.Fatal("Failed to append certificates")
+	}
+	return credentials.NewTLS(&tls.Config{
+		ClientAuth:   clientAuth,
+		Certificates: []tls.Certificate{certs},
+		ClientCAs:    ca,
+	})
+}
+
+// CreateServerTLSCredentialsValidSPIFFEButWrongCA creates server-side TLS
+// transport credentials using certificate and key files from the
+// testdata/spiffe directory rather than the testdata/spiffe_end2end directory.
+// These credentials have the expected trust domains and SPIFFE IDs that are
+// compatible with testdata/spiffe_end2end client files, but they are signed by
+// a different CA and will thus fail the connection.
+func CreateServerTLSCredentialsValidSPIFFEButWrongCA(t *testing.T, clientAuth tls.ClientAuthType) credentials.TransportCredentials {
+	t.Helper()
+
+	cert, err := tls.LoadX509KeyPair(testdata.Path("spiffe/server1_spiffe.pem"), testdata.Path("server1.key"))
+	if err != nil {
+		t.Fatalf("tls.LoadX509KeyPair(spiffe/server1_spiffe.pem, spiffe/server.key) failed: %v", err)
+	}
+	b, err := os.ReadFile(testdata.Path("spiffe_end2end/ca.pem"))
+	if err != nil {
+		t.Fatalf("os.ReadFile(spiffe_end2end/ca.pem) failed: %v", err)
+	}
+	ca := x509.NewCertPool()
+	if !ca.AppendCertsFromPEM(b) {
+		t.Fatal("Failed to append certificates")
+	}
+	return credentials.NewTLS(&tls.Config{
+		ClientAuth:   clientAuth,
+		Certificates: []tls.Certificate{cert},
+		ClientCAs:    ca,
+	})
+}

internal/testutils/wrappers.go

@@ -0,0 +1,74 @@
+/*
+ *
+ * Copyright 2022 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 testutils
+
+import (
+	"net"
+	"testing"
+)
+
+// ConnWrapper wraps a net.Conn and pushes on a channel when closed.
+type ConnWrapper struct {
+	net.Conn
+	CloseCh *Channel
+}
+
+// Close closes the connection and sends a value on the close channel.
+func (cw *ConnWrapper) Close() error {
+	err := cw.Conn.Close()
+	cw.CloseCh.Replace(nil)
+	return err
+}
+
+// ListenerWrapper wraps a net.Listener and the returned net.Conn.
+//
+// It pushes on a channel whenever it accepts a new connection.
+type ListenerWrapper struct {
+	net.Listener
+	NewConnCh *Channel
+}
+
+// Accept wraps the Listener Accept and sends the accepted connection on a
+// channel.
+func (l *ListenerWrapper) Accept() (net.Conn, error) {
+	c, err := l.Listener.Accept()
+	if err != nil {
+		return nil, err
+	}
+	closeCh := NewChannel()
+	conn := &ConnWrapper{Conn: c, CloseCh: closeCh}
+	l.NewConnCh.Replace(conn)
+	return conn, nil
+}
+
+// NewListenerWrapper returns a ListenerWrapper.
+func NewListenerWrapper(t *testing.T, lis net.Listener) *ListenerWrapper {
+	if lis == nil {
+		var err error
+		lis, err = LocalTCPListener()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	return &ListenerWrapper{
+		Listener:  lis,
+		NewConnCh: NewChannel(),
+	}
+}

internal/testutils/wrr.go

@@ -0,0 +1,73 @@
+/*
+ *
+ * Copyright 2020 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 testutils
+
+import (
+	"fmt"
+	"sync"
+
+	"google.golang.org/grpc/internal/wrr"
+)
+
+// testWRR is a deterministic WRR implementation.
+//
+// The real implementation does random WRR. testWRR makes the balancer behavior
+// deterministic and easier to test.
+//
+// With {a: 2, b: 3}, the Next() results will be {a, a, b, b, b}.
+type testWRR struct {
+	itemsWithWeight []struct {
+		item   any
+		weight int64
+	}
+	length int
+
+	mu    sync.Mutex
+	idx   int   // The index of the item that will be picked
+	count int64 // The number of times the current item has been picked.
+}
+
+// NewTestWRR return a WRR for testing. It's deterministic instead of random.
+func NewTestWRR() wrr.WRR {
+	return &testWRR{}
+}
+
+func (twrr *testWRR) Add(item any, weight int64) {
+	twrr.itemsWithWeight = append(twrr.itemsWithWeight, struct {
+		item   any
+		weight int64
+	}{item: item, weight: weight})
+	twrr.length++
+}
+
+func (twrr *testWRR) Next() any {
+	twrr.mu.Lock()
+	iww := twrr.itemsWithWeight[twrr.idx]
+	twrr.count++
+	if twrr.count >= iww.weight {
+		twrr.idx = (twrr.idx + 1) % twrr.length
+		twrr.count = 0
+	}
+	twrr.mu.Unlock()
+	return iww.item
+}
+
+func (twrr *testWRR) String() string {
+	return fmt.Sprint(twrr.itemsWithWeight)
+}

internal/testutils/xds/e2e/bootstrap.go

@@ -0,0 +1,196 @@
+/*
+ *
+ * Copyright 2020 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 e2e
+
+import (
+	"encoding/json"
+	"fmt"
+	"os"
+	"path"
+	"testing"
+
+	"google.golang.org/grpc/internal/xds/bootstrap"
+	"google.golang.org/grpc/testdata"
+)
+
+// DefaultFileWatcherConfig is a helper function to create a default certificate
+// provider plugin configuration. The test is expected to have setup the files
+// appropriately before this configuration is used to instantiate providers.
+func DefaultFileWatcherConfig(certPath, keyPath, caPath string) json.RawMessage {
+	return json.RawMessage(fmt.Sprintf(`{
+			"plugin_name": "file_watcher",
+			"config": {
+				"certificate_file": %q,
+				"private_key_file": %q,
+				"ca_certificate_file": %q,
+				"refresh_interval": "600s"
+			}
+		}`, certPath, keyPath, caPath))
+}
+
+// SPIFFEFileWatcherConfig is a helper function to create a default certificate
+// provider plugin configuration. The test is expected to have setup the files
+// appropriately before this configuration is used to instantiate providers.
+func SPIFFEFileWatcherConfig(certPath, keyPath, caPath, spiffeBundleMapPath string) json.RawMessage {
+	return json.RawMessage(fmt.Sprintf(`{
+			"plugin_name": "file_watcher",
+			"config": {
+				"certificate_file": %q,
+				"private_key_file": %q,
+				"ca_certificate_file": %q,
+				"spiffe_trust_bundle_map_file": %q,
+				"refresh_interval": "600s"
+			}
+		}`, certPath, keyPath, caPath, spiffeBundleMapPath))
+}
+
+// SPIFFEBootstrapContents creates a bootstrap configuration with the given node
+// ID and server URI. It also creates certificate provider configuration using
+// SPIFFE certificates and sets the listener resource name template to be used
+// on the server side.
+func SPIFFEBootstrapContents(t *testing.T, nodeID, serverURI string) []byte {
+	t.Helper()
+
+	// Create a directory to hold certs and key files used on the server side.
+	serverDir, err := createTmpDirWithCerts("testServerSideXDSSPIFFE*", "spiffe_end2end/server_spiffe.pem", "spiffe_end2end/server.key", "spiffe_end2end/ca.pem", "spiffe_end2end/server_spiffebundle.json")
+	if err != nil {
+		t.Fatalf("Failed to create bootstrap configuration: %v", err)
+	}
+
+	// Create a directory to hold certs and key files used on the client side.
+	clientDir, err := createTmpDirWithCerts("testClientSideXDSSPIFFE*", "spiffe_end2end/client_spiffe.pem", "spiffe_end2end/client.key", "spiffe_end2end/ca.pem", "spiffe_end2end/client_spiffebundle.json")
+	if err != nil {
+		t.Fatalf("Failed to create bootstrap configuration: %v", err)
+	}
+
+	// Create certificate providers section of the bootstrap config with entries
+	// for both the client and server sides.
+	cpc := map[string]json.RawMessage{
+		ServerSideCertProviderInstance: SPIFFEFileWatcherConfig(path.Join(serverDir, certFile), path.Join(serverDir, keyFile), path.Join(serverDir, rootFile), path.Join(serverDir, spiffeBundleMapFile)),
+		ClientSideCertProviderInstance: SPIFFEFileWatcherConfig(path.Join(clientDir, certFile), path.Join(clientDir, keyFile), path.Join(clientDir, rootFile), path.Join(clientDir, spiffeBundleMapFile)),
+	}
+
+	// Create the bootstrap configuration.
+	bs, err := bootstrap.NewContentsForTesting(bootstrap.ConfigOptionsForTesting{
+		Servers: []byte(fmt.Sprintf(`[{
+			"server_uri": "passthrough:///%s",
+			"channel_creds": [{"type": "insecure"}]
+		}]`, serverURI)),
+		Node:                               []byte(fmt.Sprintf(`{"id": "%s"}`, nodeID)),
+		CertificateProviders:               cpc,
+		ServerListenerResourceNameTemplate: ServerListenerResourceNameTemplate,
+	})
+	if err != nil {
+		t.Fatalf("Failed to create bootstrap configuration: %v", err)
+	}
+	return bs
+
+}
+
+// DefaultBootstrapContents creates a default bootstrap configuration with the
+// given node ID and server URI. It also creates certificate provider
+// configuration and sets the listener resource name template to be used on the
+// server side.
+func DefaultBootstrapContents(t *testing.T, nodeID, serverURI string) []byte {
+	t.Helper()
+
+	// Create a directory to hold certs and key files used on the server side.
+	serverDir, err := createTmpDirWithCerts("testServerSideXDS*", "x509/server1_cert.pem", "x509/server1_key.pem", "x509/client_ca_cert.pem", "")
+	if err != nil {
+		t.Fatalf("Failed to create bootstrap configuration: %v", err)
+	}
+
+	// Create a directory to hold certs and key files used on the client side.
+	clientDir, err := createTmpDirWithCerts("testClientSideXDS*", "x509/client1_cert.pem", "x509/client1_key.pem", "x509/server_ca_cert.pem", "")
+	if err != nil {
+		t.Fatalf("Failed to create bootstrap configuration: %v", err)
+	}
+
+	// Create certificate providers section of the bootstrap config with entries
+	// for both the client and server sides.
+	cpc := map[string]json.RawMessage{
+		ServerSideCertProviderInstance: DefaultFileWatcherConfig(path.Join(serverDir, certFile), path.Join(serverDir, keyFile), path.Join(serverDir, rootFile)),
+		ClientSideCertProviderInstance: DefaultFileWatcherConfig(path.Join(clientDir, certFile), path.Join(clientDir, keyFile), path.Join(clientDir, rootFile)),
+	}
+
+	// Create the bootstrap configuration.
+	bs, err := bootstrap.NewContentsForTesting(bootstrap.ConfigOptionsForTesting{
+		Servers: []byte(fmt.Sprintf(`[{
+			"server_uri": "passthrough:///%s",
+			"channel_creds": [{"type": "insecure"}],
+			"server_features": ["trusted_xds_server"]
+		}]`, serverURI)),
+		Node:                               []byte(fmt.Sprintf(`{"id": "%s"}`, nodeID)),
+		CertificateProviders:               cpc,
+		ServerListenerResourceNameTemplate: ServerListenerResourceNameTemplate,
+	})
+	if err != nil {
+		t.Fatalf("Failed to create bootstrap configuration: %v", err)
+	}
+	return bs
+}
+
+const (
+	// Names of files inside tempdir, for certprovider plugin to watch.
+	certFile            = "cert.pem"
+	keyFile             = "key.pem"
+	rootFile            = "ca.pem"
+	spiffeBundleMapFile = "spiffe_bundle_map.json"
+)
+
+func createTmpFile(src, dst string) error {
+	data, err := os.ReadFile(src)
+	if err != nil {
+		return fmt.Errorf("os.ReadFile(%q) failed: %v", src, err)
+	}
+	if err := os.WriteFile(dst, data, os.ModePerm); err != nil {
+		return fmt.Errorf("os.WriteFile(%q) failed: %v", dst, err)
+	}
+	return nil
+}
+
+// createTmpDirWithCerts creates a temporary directory under the system default
+// tempDir with the given dirPattern. It also reads from certSrc, keySrc and
+// rootSrc files and creates appropriate files under the newly create tempDir.
+// Returns the path of the created tempDir if successful, and an error
+// otherwise.
+func createTmpDirWithCerts(dirPattern, certSrc, keySrc, rootSrc, spiffeBundleMapSrc string) (string, error) {
+	// Create a temp directory. Passing an empty string for the first argument
+	// uses the system temp directory.
+	dir, err := os.MkdirTemp("", dirPattern)
+	if err != nil {
+		return "", fmt.Errorf("os.MkdirTemp() failed: %v", err)
+	}
+
+	if err := createTmpFile(testdata.Path(certSrc), path.Join(dir, certFile)); err != nil {
+		return "", err
+	}
+	if err := createTmpFile(testdata.Path(keySrc), path.Join(dir, keyFile)); err != nil {
+		return "", err
+	}
+	if err := createTmpFile(testdata.Path(rootSrc), path.Join(dir, rootFile)); err != nil {
+		return "", err
+	}
+	if spiffeBundleMapSrc != "" {
+		if err := createTmpFile(testdata.Path(spiffeBundleMapSrc), path.Join(dir, spiffeBundleMapFile)); err != nil {
+			return "", err
+		}
+	}
+	return dir, nil
+}

internal/testutils/xds/e2e/clientresources.go

@@ -0,0 +1,846 @@
+/*
+ *
+ * Copyright 2021 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 e2e
+
+import (
+	"fmt"
+	"net"
+	"strconv"
+
+	"github.com/envoyproxy/go-control-plane/pkg/wellknown"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/anypb"
+	"google.golang.org/protobuf/types/known/structpb"
+	"google.golang.org/protobuf/types/known/wrapperspb"
+
+	v3clusterpb "github.com/envoyproxy/go-control-plane/envoy/config/cluster/v3"
+	v3corepb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
+	v3endpointpb "github.com/envoyproxy/go-control-plane/envoy/config/endpoint/v3"
+	v3listenerpb "github.com/envoyproxy/go-control-plane/envoy/config/listener/v3"
+	v3routepb "github.com/envoyproxy/go-control-plane/envoy/config/route/v3"
+	v3aggregateclusterpb "github.com/envoyproxy/go-control-plane/envoy/extensions/clusters/aggregate/v3"
+	v3routerpb "github.com/envoyproxy/go-control-plane/envoy/extensions/filters/http/router/v3"
+	v3httppb "github.com/envoyproxy/go-control-plane/envoy/extensions/filters/network/http_connection_manager/v3"
+	v3tlspb "github.com/envoyproxy/go-control-plane/envoy/extensions/transport_sockets/tls/v3"
+	v3typepb "github.com/envoyproxy/go-control-plane/envoy/type/v3"
+)
+
+const (
+	// ServerListenerResourceNameTemplate is the Listener resource name template
+	// used on the server side.
+	ServerListenerResourceNameTemplate = "grpc/server?xds.resource.listening_address=%s"
+	// ClientSideCertProviderInstance is the certificate provider instance name
+	// used in the Cluster resource on the client side.
+	ClientSideCertProviderInstance = "client-side-certificate-provider-instance"
+	// ServerSideCertProviderInstance is the certificate provider instance name
+	// used in the Listener resource on the server side.
+	ServerSideCertProviderInstance = "server-side-certificate-provider-instance"
+)
+
+// SecurityLevel allows the test to control the security level to be used in the
+// resource returned by this package.
+type SecurityLevel int
+
+const (
+	// SecurityLevelNone is used when no security configuration is required.
+	SecurityLevelNone SecurityLevel = iota
+	// SecurityLevelTLS is used when security configuration corresponding to TLS
+	// is required. Only the server presents an identity certificate in this
+	// configuration.
+	SecurityLevelTLS
+	// SecurityLevelMTLS is used when security configuration corresponding to
+	// mTLS is required. Both client and server present identity certificates in
+	// this configuration.
+	SecurityLevelMTLS
+	// SecurityLevelTLSWithSystemRootCerts is used when security configuration
+	// corresponding to TLS is required. Only the server presents an identity
+	// certificate in this configuration and the client uses system root certs
+	// to validate the server certificate.
+	SecurityLevelTLSWithSystemRootCerts
+)
+
+// ResourceParams wraps the arguments to be passed to DefaultClientResources.
+type ResourceParams struct {
+	// DialTarget is the client's dial target. This is used as the name of the
+	// Listener resource.
+	DialTarget string
+	// NodeID is the id of the xdsClient to which this update is to be pushed.
+	NodeID string
+	// Host is the host of the default Endpoint resource.
+	Host string
+	// port is the port of the default Endpoint resource.
+	Port uint32
+	// SecLevel controls the security configuration in the Cluster resource.
+	SecLevel SecurityLevel
+}
+
+// DefaultClientResources returns a set of resources (LDS, RDS, CDS, EDS) for a
+// client to generically connect to one server.
+func DefaultClientResources(params ResourceParams) UpdateOptions {
+	routeConfigName := "route-" + params.DialTarget
+	clusterName := "cluster-" + params.DialTarget
+	endpointsName := "endpoints-" + params.DialTarget
+	return UpdateOptions{
+		NodeID:    params.NodeID,
+		Listeners: []*v3listenerpb.Listener{DefaultClientListener(params.DialTarget, routeConfigName)},
+		Routes:    []*v3routepb.RouteConfiguration{DefaultRouteConfig(routeConfigName, params.DialTarget, clusterName)},
+		Clusters:  []*v3clusterpb.Cluster{DefaultCluster(clusterName, endpointsName, params.SecLevel)},
+		Endpoints: []*v3endpointpb.ClusterLoadAssignment{DefaultEndpoint(endpointsName, params.Host, []uint32{params.Port})},
+	}
+}
+
+// RouterHTTPFilter is the HTTP Filter configuration for the Router filter.
+var RouterHTTPFilter = HTTPFilter("router", &v3routerpb.Router{})
+
+// DefaultClientListener returns a basic xds Listener resource to be used on
+// the client side.
+func DefaultClientListener(target, routeName string) *v3listenerpb.Listener {
+	hcm := marshalAny(&v3httppb.HttpConnectionManager{
+		RouteSpecifier: &v3httppb.HttpConnectionManager_Rds{Rds: &v3httppb.Rds{
+			ConfigSource: &v3corepb.ConfigSource{
+				ConfigSourceSpecifier: &v3corepb.ConfigSource_Ads{Ads: &v3corepb.AggregatedConfigSource{}},
+			},
+			RouteConfigName: routeName,
+		}},
+		HttpFilters: []*v3httppb.HttpFilter{HTTPFilter("router", &v3routerpb.Router{})}, // router fields are unused by grpc
+	})
+	return &v3listenerpb.Listener{
+		Name:        target,
+		ApiListener: &v3listenerpb.ApiListener{ApiListener: hcm},
+		FilterChains: []*v3listenerpb.FilterChain{{
+			Name: "filter-chain-name",
+			Filters: []*v3listenerpb.Filter{{
+				Name:       wellknown.HTTPConnectionManager,
+				ConfigType: &v3listenerpb.Filter_TypedConfig{TypedConfig: hcm},
+			}},
+		}},
+	}
+}
+
+func marshalAny(m proto.Message) *anypb.Any {
+	a, err := anypb.New(m)
+	if err != nil {
+		panic(fmt.Sprintf("anypb.New(%+v) failed: %v", m, err))
+	}
+	return a
+}
+
+// DefaultServerListener returns a basic xds Listener resource to be used on the
+// server side. The returned Listener resource contains an inline route
+// configuration with the name of routeName.
+func DefaultServerListener(host string, port uint32, secLevel SecurityLevel, routeName string) *v3listenerpb.Listener {
+	return defaultServerListenerCommon(host, port, secLevel, routeName, true)
+}
+
+func defaultServerListenerCommon(host string, port uint32, secLevel SecurityLevel, routeName string, inlineRouteConfig bool) *v3listenerpb.Listener {
+	var hcm *v3httppb.HttpConnectionManager
+	if inlineRouteConfig {
+		hcm = &v3httppb.HttpConnectionManager{
+			RouteSpecifier: &v3httppb.HttpConnectionManager_RouteConfig{
+				RouteConfig: &v3routepb.RouteConfiguration{
+					Name: routeName,
+					VirtualHosts: []*v3routepb.VirtualHost{{
+						// This "*" string matches on any incoming authority. This is to ensure any
+						// incoming RPC matches to Route_NonForwardingAction and will proceed as
+						// normal.
+						Domains: []string{"*"},
+						Routes: []*v3routepb.Route{{
+							Match: &v3routepb.RouteMatch{
+								PathSpecifier: &v3routepb.RouteMatch_Prefix{Prefix: "/"},
+							},
+							Action: &v3routepb.Route_NonForwardingAction{},
+						}}}}},
+			},
+			HttpFilters: []*v3httppb.HttpFilter{RouterHTTPFilter},
+		}
+	} else {
+		hcm = &v3httppb.HttpConnectionManager{
+			RouteSpecifier: &v3httppb.HttpConnectionManager_Rds{
+				Rds: &v3httppb.Rds{
+					ConfigSource: &v3corepb.ConfigSource{
+						ConfigSourceSpecifier: &v3corepb.ConfigSource_Ads{Ads: &v3corepb.AggregatedConfigSource{}},
+					},
+					RouteConfigName: routeName,
+				},
+			},
+			HttpFilters: []*v3httppb.HttpFilter{RouterHTTPFilter},
+		}
+	}
+
+	var tlsContext *v3tlspb.DownstreamTlsContext
+	switch secLevel {
+	case SecurityLevelNone:
+	case SecurityLevelTLS:
+		tlsContext = &v3tlspb.DownstreamTlsContext{
+			CommonTlsContext: &v3tlspb.CommonTlsContext{
+				TlsCertificateCertificateProviderInstance: &v3tlspb.CommonTlsContext_CertificateProviderInstance{
+					InstanceName: ServerSideCertProviderInstance,
+				},
+			},
+		}
+	case SecurityLevelMTLS:
+		tlsContext = &v3tlspb.DownstreamTlsContext{
+			RequireClientCertificate: &wrapperspb.BoolValue{Value: true},
+			CommonTlsContext: &v3tlspb.CommonTlsContext{
+				TlsCertificateCertificateProviderInstance: &v3tlspb.CommonTlsContext_CertificateProviderInstance{
+					InstanceName: ServerSideCertProviderInstance,
+				},
+				ValidationContextType: &v3tlspb.CommonTlsContext_ValidationContextCertificateProviderInstance{
+					ValidationContextCertificateProviderInstance: &v3tlspb.CommonTlsContext_CertificateProviderInstance{
+						InstanceName: ServerSideCertProviderInstance,
+					},
+				},
+			},
+		}
+	}
+
+	var ts *v3corepb.TransportSocket
+	if tlsContext != nil {
+		ts = &v3corepb.TransportSocket{
+			Name: "envoy.transport_sockets.tls",
+			ConfigType: &v3corepb.TransportSocket_TypedConfig{
+				TypedConfig: marshalAny(tlsContext),
+			},
+		}
+	}
+	return &v3listenerpb.Listener{
+		Name: fmt.Sprintf(ServerListenerResourceNameTemplate, net.JoinHostPort(host, strconv.Itoa(int(port)))),
+		Address: &v3corepb.Address{
+			Address: &v3corepb.Address_SocketAddress{
+				SocketAddress: &v3corepb.SocketAddress{
+					Address: host,
+					PortSpecifier: &v3corepb.SocketAddress_PortValue{
+						PortValue: port,
+					},
+				},
+			},
+		},
+		FilterChains: []*v3listenerpb.FilterChain{
+			{
+				Name: "v4-wildcard",
+				FilterChainMatch: &v3listenerpb.FilterChainMatch{
+					PrefixRanges: []*v3corepb.CidrRange{
+						{
+							AddressPrefix: "0.0.0.0",
+							PrefixLen: &wrapperspb.UInt32Value{
+								Value: uint32(0),
+							},
+						},
+					},
+					SourceType: v3listenerpb.FilterChainMatch_SAME_IP_OR_LOOPBACK,
+					SourcePrefixRanges: []*v3corepb.CidrRange{
+						{
+							AddressPrefix: "0.0.0.0",
+							PrefixLen: &wrapperspb.UInt32Value{
+								Value: uint32(0),
+							},
+						},
+					},
+				},
+				Filters: []*v3listenerpb.Filter{
+					{
+						Name:       "filter-1",
+						ConfigType: &v3listenerpb.Filter_TypedConfig{TypedConfig: marshalAny(hcm)},
+					},
+				},
+				TransportSocket: ts,
+			},
+			{
+				Name: "v6-wildcard",
+				FilterChainMatch: &v3listenerpb.FilterChainMatch{
+					PrefixRanges: []*v3corepb.CidrRange{
+						{
+							AddressPrefix: "::",
+							PrefixLen: &wrapperspb.UInt32Value{
+								Value: uint32(0),
+							},
+						},
+					},
+					SourceType: v3listenerpb.FilterChainMatch_SAME_IP_OR_LOOPBACK,
+					SourcePrefixRanges: []*v3corepb.CidrRange{
+						{
+							AddressPrefix: "::",
+							PrefixLen: &wrapperspb.UInt32Value{
+								Value: uint32(0),
+							},
+						},
+					},
+				},
+				Filters: []*v3listenerpb.Filter{
+					{
+						Name:       "filter-1",
+						ConfigType: &v3listenerpb.Filter_TypedConfig{TypedConfig: marshalAny(hcm)},
+					},
+				},
+				TransportSocket: ts,
+			},
+		},
+	}
+}
+
+// HTTPFilter constructs an xds HttpFilter with the provided name and config.
+func HTTPFilter(name string, config proto.Message) *v3httppb.HttpFilter {
+	return &v3httppb.HttpFilter{
+		Name: name,
+		ConfigType: &v3httppb.HttpFilter_TypedConfig{
+			TypedConfig: marshalAny(config),
+		},
+	}
+}
+
+// DefaultRouteConfig returns a basic xds RouteConfig resource.
+func DefaultRouteConfig(routeName, vhDomain, clusterName string) *v3routepb.RouteConfiguration {
+	return &v3routepb.RouteConfiguration{
+		Name: routeName,
+		VirtualHosts: []*v3routepb.VirtualHost{{
+			Domains: []string{vhDomain},
+			Routes: []*v3routepb.Route{{
+				Match: &v3routepb.RouteMatch{PathSpecifier: &v3routepb.RouteMatch_Prefix{Prefix: "/"}},
+				Action: &v3routepb.Route_Route{Route: &v3routepb.RouteAction{
+					ClusterSpecifier: &v3routepb.RouteAction_WeightedClusters{WeightedClusters: &v3routepb.WeightedCluster{
+						Clusters: []*v3routepb.WeightedCluster_ClusterWeight{
+							{
+								Name:   clusterName,
+								Weight: &wrapperspb.UInt32Value{Value: 100},
+							},
+						},
+					}},
+				}},
+			}},
+		}},
+	}
+}
+
+// RouteConfigClusterSpecifierType determines the cluster specifier type for the
+// route actions configured in the returned RouteConfiguration resource.
+type RouteConfigClusterSpecifierType int
+
+const (
+	// RouteConfigClusterSpecifierTypeCluster results in the cluster specifier
+	// being set to a RouteAction_Cluster.
+	RouteConfigClusterSpecifierTypeCluster RouteConfigClusterSpecifierType = iota
+	// RouteConfigClusterSpecifierTypeWeightedCluster results in the cluster
+	// specifier being set to RouteAction_WeightedClusters.
+	RouteConfigClusterSpecifierTypeWeightedCluster
+	// RouteConfigClusterSpecifierTypeClusterSpecifierPlugin results in the
+	// cluster specifier being set to a RouteAction_ClusterSpecifierPlugin.
+	RouteConfigClusterSpecifierTypeClusterSpecifierPlugin
+)
+
+// RouteConfigOptions contains options to configure a RouteConfiguration
+// resource.
+type RouteConfigOptions struct {
+	// RouteConfigName is the name of the RouteConfiguration resource.
+	RouteConfigName string
+	// ListenerName is the name of the Listener resource which uses this
+	// RouteConfiguration.
+	ListenerName string
+	// ClusterSpecifierType determines the cluster specifier type.
+	ClusterSpecifierType RouteConfigClusterSpecifierType
+	// ClusterName is name of the cluster resource used when the cluster
+	// specifier type is set to RouteConfigClusterSpecifierTypeCluster.
+	//
+	// Default value of "A" is used if left unspecified.
+	ClusterName string
+	// WeightedClusters is a map from cluster name to weights, and is used when
+	// the cluster specifier type is set to
+	// RouteConfigClusterSpecifierTypeWeightedCluster.
+	//
+	// Default value of {"A": 75, "B": 25} is used if left unspecified.
+	WeightedClusters map[string]int
+	// The below two fields specify the name of the cluster specifier plugin and
+	// its configuration, and are used when the cluster specifier type is set to
+	// RouteConfigClusterSpecifierTypeClusterSpecifierPlugin. Tests are expected
+	// to provide valid values for these fields when appropriate.
+	ClusterSpecifierPluginName   string
+	ClusterSpecifierPluginConfig *anypb.Any
+}
+
+// RouteConfigResourceWithOptions returns a RouteConfiguration resource
+// configured with the provided options.
+func RouteConfigResourceWithOptions(opts RouteConfigOptions) *v3routepb.RouteConfiguration {
+	switch opts.ClusterSpecifierType {
+	case RouteConfigClusterSpecifierTypeCluster:
+		clusterName := opts.ClusterName
+		if clusterName == "" {
+			clusterName = "A"
+		}
+		return &v3routepb.RouteConfiguration{
+			Name: opts.RouteConfigName,
+			VirtualHosts: []*v3routepb.VirtualHost{{
+				Domains: []string{opts.ListenerName},
+				Routes: []*v3routepb.Route{{
+					Match: &v3routepb.RouteMatch{PathSpecifier: &v3routepb.RouteMatch_Prefix{Prefix: "/"}},
+					Action: &v3routepb.Route_Route{Route: &v3routepb.RouteAction{
+						ClusterSpecifier: &v3routepb.RouteAction_Cluster{Cluster: clusterName},
+					}},
+				}},
+			}},
+		}
+	case RouteConfigClusterSpecifierTypeWeightedCluster:
+		weightedClusters := opts.WeightedClusters
+		if weightedClusters == nil {
+			weightedClusters = map[string]int{"A": 75, "B": 25}
+		}
+		clusters := []*v3routepb.WeightedCluster_ClusterWeight{}
+		for name, weight := range weightedClusters {
+			clusters = append(clusters, &v3routepb.WeightedCluster_ClusterWeight{
+				Name:   name,
+				Weight: &wrapperspb.UInt32Value{Value: uint32(weight)},
+			})
+		}
+		return &v3routepb.RouteConfiguration{
+			Name: opts.RouteConfigName,
+			VirtualHosts: []*v3routepb.VirtualHost{{
+				Domains: []string{opts.ListenerName},
+				Routes: []*v3routepb.Route{{
+					Match: &v3routepb.RouteMatch{PathSpecifier: &v3routepb.RouteMatch_Prefix{Prefix: "/"}},
+					Action: &v3routepb.Route_Route{Route: &v3routepb.RouteAction{
+						ClusterSpecifier: &v3routepb.RouteAction_WeightedClusters{WeightedClusters: &v3routepb.WeightedCluster{Clusters: clusters}},
+					}},
+				}},
+			}},
+		}
+	case RouteConfigClusterSpecifierTypeClusterSpecifierPlugin:
+		return &v3routepb.RouteConfiguration{
+			Name: opts.RouteConfigName,
+			ClusterSpecifierPlugins: []*v3routepb.ClusterSpecifierPlugin{{
+				Extension: &v3corepb.TypedExtensionConfig{
+					Name:        opts.ClusterSpecifierPluginName,
+					TypedConfig: opts.ClusterSpecifierPluginConfig,
+				}},
+			},
+			VirtualHosts: []*v3routepb.VirtualHost{{
+				Domains: []string{opts.ListenerName},
+				Routes: []*v3routepb.Route{{
+					Match: &v3routepb.RouteMatch{PathSpecifier: &v3routepb.RouteMatch_Prefix{Prefix: "/"}},
+					Action: &v3routepb.Route_Route{Route: &v3routepb.RouteAction{
+						ClusterSpecifier: &v3routepb.RouteAction_ClusterSpecifierPlugin{ClusterSpecifierPlugin: opts.ClusterSpecifierPluginName},
+					}},
+				}},
+			}},
+		}
+	default:
+		panic(fmt.Sprintf("unsupported cluster specifier plugin type: %v", opts.ClusterSpecifierType))
+	}
+}
+
+// DefaultCluster returns a basic xds Cluster resource.
+func DefaultCluster(clusterName, edsServiceName string, secLevel SecurityLevel) *v3clusterpb.Cluster {
+	return ClusterResourceWithOptions(ClusterOptions{
+		ClusterName:   clusterName,
+		ServiceName:   edsServiceName,
+		Policy:        LoadBalancingPolicyRoundRobin,
+		SecurityLevel: secLevel,
+	})
+}
+
+// LoadBalancingPolicy determines the policy used for balancing load across
+// endpoints in the Cluster.
+type LoadBalancingPolicy int
+
+const (
+	// LoadBalancingPolicyRoundRobin results in the use of the weighted_target
+	// LB policy to balance load across localities and endpoints in the cluster.
+	LoadBalancingPolicyRoundRobin LoadBalancingPolicy = iota
+	// LoadBalancingPolicyRingHash results in the use of the ring_hash LB policy
+	// as the leaf policy.
+	LoadBalancingPolicyRingHash
+)
+
+// ClusterType specifies the type of the Cluster resource.
+type ClusterType int
+
+const (
+	// ClusterTypeEDS specifies a Cluster that uses EDS to resolve endpoints.
+	ClusterTypeEDS ClusterType = iota
+	// ClusterTypeLogicalDNS specifies a Cluster that uses DNS to resolve
+	// endpoints.
+	ClusterTypeLogicalDNS
+	// ClusterTypeAggregate specifies a Cluster that is made up of child
+	// clusters.
+	ClusterTypeAggregate
+)
+
+// ClusterOptions contains options to configure a Cluster resource.
+type ClusterOptions struct {
+	Type ClusterType
+	// ClusterName is the name of the Cluster resource.
+	ClusterName string
+	// ServiceName is the EDS service name of the Cluster. Applicable only when
+	// cluster type is EDS.
+	ServiceName string
+	// ChildNames is the list of child Cluster names. Applicable only when
+	// cluster type is Aggregate.
+	ChildNames []string
+	// DNSHostName is the dns host name of the Cluster. Applicable only when the
+	// cluster type is DNS.
+	DNSHostName string
+	// DNSPort is the port number of the Cluster. Applicable only when the
+	// cluster type is DNS.
+	DNSPort uint32
+	// Policy is the LB policy to be used.
+	Policy LoadBalancingPolicy
+	// SecurityLevel determines the security configuration for the Cluster.
+	SecurityLevel SecurityLevel
+	// EnableLRS adds a load reporting configuration with a config source
+	// pointing to self.
+	EnableLRS bool
+}
+
+// ClusterResourceWithOptions returns an xDS Cluster resource configured with
+// the provided options.
+func ClusterResourceWithOptions(opts ClusterOptions) *v3clusterpb.Cluster {
+	var tlsContext *v3tlspb.UpstreamTlsContext
+	switch opts.SecurityLevel {
+	case SecurityLevelNone:
+	case SecurityLevelTLS:
+		tlsContext = &v3tlspb.UpstreamTlsContext{
+			CommonTlsContext: &v3tlspb.CommonTlsContext{
+				ValidationContextType: &v3tlspb.CommonTlsContext_ValidationContextCertificateProviderInstance{
+					ValidationContextCertificateProviderInstance: &v3tlspb.CommonTlsContext_CertificateProviderInstance{
+						InstanceName: ClientSideCertProviderInstance,
+					},
+				},
+			},
+		}
+	case SecurityLevelMTLS:
+		tlsContext = &v3tlspb.UpstreamTlsContext{
+			CommonTlsContext: &v3tlspb.CommonTlsContext{
+				ValidationContextType: &v3tlspb.CommonTlsContext_ValidationContextCertificateProviderInstance{
+					ValidationContextCertificateProviderInstance: &v3tlspb.CommonTlsContext_CertificateProviderInstance{
+						InstanceName: ClientSideCertProviderInstance,
+					},
+				},
+				TlsCertificateCertificateProviderInstance: &v3tlspb.CommonTlsContext_CertificateProviderInstance{
+					InstanceName: ClientSideCertProviderInstance,
+				},
+			},
+		}
+	case SecurityLevelTLSWithSystemRootCerts:
+		tlsContext = &v3tlspb.UpstreamTlsContext{
+			CommonTlsContext: &v3tlspb.CommonTlsContext{
+				ValidationContextType: &v3tlspb.CommonTlsContext_ValidationContext{
+					ValidationContext: &v3tlspb.CertificateValidationContext{
+						SystemRootCerts: &v3tlspb.CertificateValidationContext_SystemRootCerts{},
+					},
+				},
+			},
+		}
+	}
+
+	var lbPolicy v3clusterpb.Cluster_LbPolicy
+	switch opts.Policy {
+	case LoadBalancingPolicyRoundRobin:
+		lbPolicy = v3clusterpb.Cluster_ROUND_ROBIN
+	case LoadBalancingPolicyRingHash:
+		lbPolicy = v3clusterpb.Cluster_RING_HASH
+	}
+	cluster := &v3clusterpb.Cluster{
+		Name:     opts.ClusterName,
+		LbPolicy: lbPolicy,
+	}
+	switch opts.Type {
+	case ClusterTypeEDS:
+		cluster.ClusterDiscoveryType = &v3clusterpb.Cluster_Type{Type: v3clusterpb.Cluster_EDS}
+		cluster.EdsClusterConfig = &v3clusterpb.Cluster_EdsClusterConfig{
+			EdsConfig: &v3corepb.ConfigSource{
+				ConfigSourceSpecifier: &v3corepb.ConfigSource_Ads{
+					Ads: &v3corepb.AggregatedConfigSource{},
+				},
+			},
+			ServiceName: opts.ServiceName,
+		}
+	case ClusterTypeLogicalDNS:
+		cluster.ClusterDiscoveryType = &v3clusterpb.Cluster_Type{Type: v3clusterpb.Cluster_LOGICAL_DNS}
+		cluster.LoadAssignment = &v3endpointpb.ClusterLoadAssignment{
+			Endpoints: []*v3endpointpb.LocalityLbEndpoints{{
+				LbEndpoints: []*v3endpointpb.LbEndpoint{{
+					HostIdentifier: &v3endpointpb.LbEndpoint_Endpoint{
+						Endpoint: &v3endpointpb.Endpoint{
+							Address: &v3corepb.Address{
+								Address: &v3corepb.Address_SocketAddress{
+									SocketAddress: &v3corepb.SocketAddress{
+										Address: opts.DNSHostName,
+										PortSpecifier: &v3corepb.SocketAddress_PortValue{
+											PortValue: opts.DNSPort,
+										},
+									},
+								},
+							},
+						},
+					},
+				}},
+			}},
+		}
+	case ClusterTypeAggregate:
+		cluster.ClusterDiscoveryType = &v3clusterpb.Cluster_ClusterType{
+			ClusterType: &v3clusterpb.Cluster_CustomClusterType{
+				Name: "envoy.clusters.aggregate",
+				TypedConfig: marshalAny(&v3aggregateclusterpb.ClusterConfig{
+					Clusters: opts.ChildNames,
+				}),
+			},
+		}
+	}
+	if tlsContext != nil {
+		cluster.TransportSocket = &v3corepb.TransportSocket{
+			Name: "envoy.transport_sockets.tls",
+			ConfigType: &v3corepb.TransportSocket_TypedConfig{
+				TypedConfig: marshalAny(tlsContext),
+			},
+		}
+	}
+	if opts.EnableLRS {
+		cluster.LrsServer = &v3corepb.ConfigSource{
+			ConfigSourceSpecifier: &v3corepb.ConfigSource_Self{
+				Self: &v3corepb.SelfConfigSource{},
+			},
+		}
+	}
+	return cluster
+}
+
+// LocalityID represents a locality identifier.
+type LocalityID struct {
+	Region  string
+	Zone    string
+	SubZone string
+}
+
+// LocalityOptions contains options to configure a Locality.
+type LocalityOptions struct {
+	// Name is the unique locality name.
+	Name string
+	// Weight is the weight of the locality, used for load balancing.
+	Weight uint32
+	// Backends is a set of backends belonging to this locality.
+	Backends []BackendOptions
+	// Priority is the priority of the locality. Defaults to 0.
+	Priority uint32
+	// Locality is the locality identifier. If not specified, a random
+	// identifier is generated.
+	Locality LocalityID
+}
+
+// BackendOptions contains options to configure individual backends in a
+// locality.
+type BackendOptions struct {
+	// Ports on which the backend is accepting connections. All backends
+	// are expected to run on localhost, hence host name is not stored here.
+	Ports []uint32
+	// Health status of the backend. Default is UNKNOWN which is treated the
+	// same as HEALTHY.
+	HealthStatus v3corepb.HealthStatus
+	// Weight sets the backend weight. Defaults to 1.
+	Weight uint32
+	// Hostname sets the endpoint hostname for authority rewriting.
+	Hostname string
+	// Metadata sets the LB endpoint metadata (envoy.lb FilterMetadata field).
+	// See https://www.envoyproxy.io/docs/envoy/latest/api-v3/config/core/v3/base.proto#envoy-v3-api-msg-config-core-v3-metadata
+	Metadata map[string]any
+}
+
+// EndpointOptions contains options to configure an Endpoint (or
+// ClusterLoadAssignment) resource.
+type EndpointOptions struct {
+	// ClusterName is the name of the Cluster resource (or EDS service name)
+	// containing the endpoints specified below.
+	ClusterName string
+	// Host is the hostname of the endpoints. In our e2e tests, hostname must
+	// always be "localhost".
+	Host string
+	// Localities is a set of localities belonging to this resource.
+	Localities []LocalityOptions
+	// DropPercents is a map from drop category to a drop percentage. If unset,
+	// no drops are configured.
+	DropPercents map[string]int
+}
+
+// DefaultEndpoint returns a basic xds Endpoint resource.
+func DefaultEndpoint(clusterName string, host string, ports []uint32) *v3endpointpb.ClusterLoadAssignment {
+	var bOpts []BackendOptions
+	for _, p := range ports {
+		bOpts = append(bOpts, BackendOptions{Ports: []uint32{p}, Weight: 1})
+	}
+	return EndpointResourceWithOptions(EndpointOptions{
+		ClusterName: clusterName,
+		Host:        host,
+		Localities: []LocalityOptions{
+			{
+				Backends: bOpts,
+				Weight:   1,
+			},
+		},
+	})
+}
+
+// EndpointResourceWithOptions returns an xds Endpoint resource configured with
+// the provided options.
+func EndpointResourceWithOptions(opts EndpointOptions) *v3endpointpb.ClusterLoadAssignment {
+	var endpoints []*v3endpointpb.LocalityLbEndpoints
+	for i, locality := range opts.Localities {
+		var lbEndpoints []*v3endpointpb.LbEndpoint
+		for _, b := range locality.Backends {
+			// Weight defaults to 1.
+			if b.Weight == 0 {
+				b.Weight = 1
+			}
+			additionalAddresses := make([]*v3endpointpb.Endpoint_AdditionalAddress, len(b.Ports)-1)
+			for i, p := range b.Ports[1:] {
+				additionalAddresses[i] = &v3endpointpb.Endpoint_AdditionalAddress{
+					Address: &v3corepb.Address{Address: &v3corepb.Address_SocketAddress{
+						SocketAddress: &v3corepb.SocketAddress{
+							Protocol:      v3corepb.SocketAddress_TCP,
+							Address:       opts.Host,
+							PortSpecifier: &v3corepb.SocketAddress_PortValue{PortValue: p},
+						}},
+					},
+				}
+			}
+			lbEndpoint := &v3endpointpb.LbEndpoint{
+				HostIdentifier: &v3endpointpb.LbEndpoint_Endpoint{Endpoint: &v3endpointpb.Endpoint{
+					Address: &v3corepb.Address{Address: &v3corepb.Address_SocketAddress{
+						SocketAddress: &v3corepb.SocketAddress{
+							Protocol:      v3corepb.SocketAddress_TCP,
+							Address:       opts.Host,
+							PortSpecifier: &v3corepb.SocketAddress_PortValue{PortValue: b.Ports[0]},
+						},
+					}},
+					Hostname:            b.Hostname,
+					AdditionalAddresses: additionalAddresses,
+				}},
+				HealthStatus:        b.HealthStatus,
+				LoadBalancingWeight: &wrapperspb.UInt32Value{Value: b.Weight},
+			}
+
+			if b.Metadata != nil {
+				metadata, err := structpb.NewStruct(b.Metadata)
+				if err != nil {
+					panic(fmt.Sprintf("failed to marshal metadata: %v", err))
+				}
+				lbEndpoint.Metadata = &v3corepb.Metadata{
+					FilterMetadata: map[string]*structpb.Struct{
+						"envoy.lb": metadata,
+					},
+				}
+			}
+			lbEndpoints = append(lbEndpoints, lbEndpoint)
+		}
+
+		l := locality.Locality
+		if l == (LocalityID{}) {
+			l = LocalityID{
+				Region:  fmt.Sprintf("region-%d", i+1),
+				Zone:    fmt.Sprintf("zone-%d", i+1),
+				SubZone: fmt.Sprintf("subzone-%d", i+1),
+			}
+		}
+		endpoints = append(endpoints, &v3endpointpb.LocalityLbEndpoints{
+			Locality:            &v3corepb.Locality{Region: l.Region, Zone: l.Zone, SubZone: l.SubZone},
+			LbEndpoints:         lbEndpoints,
+			LoadBalancingWeight: &wrapperspb.UInt32Value{Value: locality.Weight},
+			Priority:            locality.Priority,
+		})
+	}
+
+	cla := &v3endpointpb.ClusterLoadAssignment{
+		ClusterName: opts.ClusterName,
+		Endpoints:   endpoints,
+	}
+
+	var drops []*v3endpointpb.ClusterLoadAssignment_Policy_DropOverload
+	for category, val := range opts.DropPercents {
+		drops = append(drops, &v3endpointpb.ClusterLoadAssignment_Policy_DropOverload{
+			Category: category,
+			DropPercentage: &v3typepb.FractionalPercent{
+				Numerator:   uint32(val),
+				Denominator: v3typepb.FractionalPercent_HUNDRED,
+			},
+		})
+	}
+	if len(drops) != 0 {
+		cla.Policy = &v3endpointpb.ClusterLoadAssignment_Policy{
+			DropOverloads: drops,
+		}
+	}
+	return cla
+}
+
+// DefaultServerListenerWithRouteConfigName returns a basic xds Listener
+// resource to be used on the server side. The returned Listener resource
+// contains a RouteConfiguration resource name that needs to be resolved.
+func DefaultServerListenerWithRouteConfigName(host string, port uint32, secLevel SecurityLevel, routeName string) *v3listenerpb.Listener {
+	return defaultServerListenerCommon(host, port, secLevel, routeName, false)
+}
+
+// RouteConfigNoRouteMatch returns an xDS RouteConfig resource which a route
+// with no route match. This will be NACKed by the xDS Client.
+func RouteConfigNoRouteMatch(routeName string) *v3routepb.RouteConfiguration {
+	return &v3routepb.RouteConfiguration{
+		Name: routeName,
+		VirtualHosts: []*v3routepb.VirtualHost{{
+			// This "*" string matches on any incoming authority. This is to ensure any
+			// incoming RPC matches to Route_NonForwardingAction and will proceed as
+			// normal.
+			Domains: []string{"*"},
+			Routes: []*v3routepb.Route{{
+				Action: &v3routepb.Route_NonForwardingAction{},
+			}}}}}
+}
+
+// RouteConfigNonForwardingAction returns an xDS RouteConfig resource which
+// specifies to route to a route specifying non forwarding action. This is
+// intended to be used on the server side for RDS requests, and corresponds to
+// the inline route configuration in DefaultServerListener.
+func RouteConfigNonForwardingAction(routeName string) *v3routepb.RouteConfiguration {
+	return &v3routepb.RouteConfiguration{
+		Name: routeName,
+		VirtualHosts: []*v3routepb.VirtualHost{{
+			// This "*" string matches on any incoming authority. This is to ensure any
+			// incoming RPC matches to Route_NonForwardingAction and will proceed as
+			// normal.
+			Domains: []string{"*"},
+			Routes: []*v3routepb.Route{{
+				Match: &v3routepb.RouteMatch{
+					PathSpecifier: &v3routepb.RouteMatch_Prefix{Prefix: "/"},
+				},
+				Action: &v3routepb.Route_NonForwardingAction{},
+			}}}}}
+}
+
+// RouteConfigFilterAction returns an xDS RouteConfig resource which specifies
+// to route to a route specifying route filter action. Since this is not type
+// non forwarding action, this should fail requests that match to this server
+// side.
+func RouteConfigFilterAction(routeName string) *v3routepb.RouteConfiguration {
+	return &v3routepb.RouteConfiguration{
+		Name: routeName,
+		VirtualHosts: []*v3routepb.VirtualHost{{
+			// This "*" string matches on any incoming authority. This is to
+			// ensure any incoming RPC matches to Route_Route and will fail with
+			// UNAVAILABLE.
+			Domains: []string{"*"},
+			Routes: []*v3routepb.Route{{
+				Match: &v3routepb.RouteMatch{
+					PathSpecifier: &v3routepb.RouteMatch_Prefix{Prefix: "/"},
+				},
+				Action: &v3routepb.Route_FilterAction{},
+			}}}}}
+}

internal/testutils/xds/e2e/logging.go

@@ -0,0 +1,40 @@
+/*
+ *
+ * Copyright 2022 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 e2e
+
+// serverLogger implements the Logger interface defined at
+// envoyproxy/go-control-plane/pkg/log. This is passed to the Snapshot cache.
+type serverLogger struct {
+	logger interface {
+		Logf(format string, args ...any)
+	}
+}
+
+func (l serverLogger) Debugf(format string, args ...any) {
+	l.logger.Logf(format, args...)
+}
+func (l serverLogger) Infof(format string, args ...any) {
+	l.logger.Logf(format, args...)
+}
+func (l serverLogger) Warnf(format string, args ...any) {
+	l.logger.Logf(format, args...)
+}
+func (l serverLogger) Errorf(format string, args ...any) {
+	l.logger.Logf(format, args...)
+}

internal/testutils/xds/e2e/server.go

@@ -0,0 +1,256 @@
+/*
+ *
+ * Copyright 2020 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 e2e provides utilities for end2end testing of xDS functionality.
+package e2e
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"reflect"
+	"strconv"
+	"testing"
+
+	"github.com/envoyproxy/go-control-plane/pkg/cache/types"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/internal/testutils/xds/fakeserver"
+
+	v3clusterpb "github.com/envoyproxy/go-control-plane/envoy/config/cluster/v3"
+	v3corepb "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
+	v3endpointpb "github.com/envoyproxy/go-control-plane/envoy/config/endpoint/v3"
+	v3listenerpb "github.com/envoyproxy/go-control-plane/envoy/config/listener/v3"
+	v3routepb "github.com/envoyproxy/go-control-plane/envoy/config/route/v3"
+	v3discoverygrpc "github.com/envoyproxy/go-control-plane/envoy/service/discovery/v3"
+	v3discoverypb "github.com/envoyproxy/go-control-plane/envoy/service/discovery/v3"
+	v3lrsgrpc "github.com/envoyproxy/go-control-plane/envoy/service/load_stats/v3"
+	v3cache "github.com/envoyproxy/go-control-plane/pkg/cache/v3"
+	v3resource "github.com/envoyproxy/go-control-plane/pkg/resource/v3"
+	v3server "github.com/envoyproxy/go-control-plane/pkg/server/v3"
+)
+
+// ManagementServer is a thin wrapper around the xDS control plane
+// implementation provided by envoyproxy/go-control-plane.
+type ManagementServer struct {
+	// Address is the host:port on which the management server is listening for
+	// new connections.
+	Address string
+
+	// LRSServer points to the fake LRS server implementation. Set only if the
+	// SupportLoadReportingService option was set to true when creating this
+	// management server.
+	LRSServer *fakeserver.Server
+
+	cancel  context.CancelFunc    // To stop the v3 ADS service.
+	xs      v3server.Server       // v3 implementation of ADS.
+	gs      *grpc.Server          // gRPC server which exports the ADS service.
+	cache   v3cache.SnapshotCache // Resource snapshot.
+	version int                   // Version of resource snapshot.
+
+	// A logging interface, usually supplied from *testing.T.
+	logger interface {
+		Logf(format string, args ...any)
+	}
+}
+
+// ManagementServerOptions contains options to be passed to the management
+// server during creation.
+type ManagementServerOptions struct {
+	// Listener to accept connections on. If nil, a TPC listener on a local port
+	// will be created and used.
+	Listener net.Listener
+
+	// SupportLoadReportingService, if set, results in the load reporting
+	// service being registered on the same port as that of ADS.
+	SupportLoadReportingService bool
+
+	// AllowResourceSubSet allows the management server to respond to requests
+	// before all configured resources are explicitly named in the request. The
+	// default behavior that we want is for the management server to wait for
+	// all configured resources to be requested before responding to any of
+	// them, since this is how we have run our tests historically, and should be
+	// set to true only for tests which explicitly require the other behavior.
+	AllowResourceSubset bool
+
+	// ServerFeaturesIgnoreResourceDeletion, if set, results in a bootstrap config
+	// where the server features list contains `ignore_resource_deletion`. This
+	// results in gRPC ignoring resource deletions from the management server, as
+	// per A53.
+	ServerFeaturesIgnoreResourceDeletion bool
+
+	// The callbacks defined below correspond to the state of the world (sotw)
+	// version of the xDS API on the management server.
+
+	// OnStreamOpen is called when an xDS stream is opened. The callback is
+	// invoked with the assigned stream ID and the type URL from the incoming
+	// request (or "" for ADS).
+	//
+	// Returning an error from this callback will end processing and close the
+	// stream. OnStreamClosed will still be called.
+	OnStreamOpen func(context.Context, int64, string) error
+
+	// OnStreamClosed is called immediately prior to closing an xDS stream. The
+	// callback is invoked with the stream ID of the stream being closed.
+	OnStreamClosed func(int64, *v3corepb.Node)
+
+	// OnStreamRequest is called when a request is received on the stream. The
+	// callback is invoked with the stream ID of the stream on which the request
+	// was received and the received request.
+	//
+	// Returning an error from this callback will end processing and close the
+	// stream. OnStreamClosed will still be called.
+	OnStreamRequest func(int64, *v3discoverypb.DiscoveryRequest) error
+
+	// OnStreamResponse is called immediately prior to sending a response on the
+	// stream. The callback is invoked with the stream ID of the stream on which
+	// the response is being sent along with the incoming request and the outgoing
+	// response.
+	OnStreamResponse func(context.Context, int64, *v3discoverypb.DiscoveryRequest, *v3discoverypb.DiscoveryResponse)
+}
+
+// StartManagementServer initializes a management server which implements the
+// AggregatedDiscoveryService endpoint. The management server is initialized
+// with no resources. Tests should call the Update() method to change the
+// resource snapshot held by the management server, as per by the test logic.
+//
+// Registers a cleanup function on t to stop the management server.
+func StartManagementServer(t *testing.T, opts ManagementServerOptions) *ManagementServer {
+	t.Helper()
+
+	// Create a snapshot cache. The first parameter to NewSnapshotCache()
+	// controls whether the server should wait for all resources to be
+	// explicitly named in the request before responding to any of them.
+	wait := !opts.AllowResourceSubset
+	cache := v3cache.NewSnapshotCache(wait, v3cache.IDHash{}, serverLogger{t})
+	t.Logf("Created new snapshot cache...")
+
+	lis := opts.Listener
+	if lis == nil {
+		var err error
+		lis, err = net.Listen("tcp", "localhost:0")
+		if err != nil {
+			t.Fatalf("Failed to listen on localhost:0: %v", err)
+		}
+	}
+
+	// Cancelling the context passed to the server is the only way of stopping it
+	// at the end of the test.
+	ctx, cancel := context.WithCancel(context.Background())
+	callbacks := v3server.CallbackFuncs{
+		StreamOpenFunc:     opts.OnStreamOpen,
+		StreamClosedFunc:   opts.OnStreamClosed,
+		StreamRequestFunc:  opts.OnStreamRequest,
+		StreamResponseFunc: opts.OnStreamResponse,
+	}
+
+	// Create an xDS management server and register the ADS implementation
+	// provided by it on a gRPC server.
+	xs := v3server.NewServer(ctx, cache, callbacks)
+	gs := grpc.NewServer()
+	v3discoverygrpc.RegisterAggregatedDiscoveryServiceServer(gs, xs)
+	t.Logf("Registered Aggregated Discovery Service (ADS)...")
+
+	mgmtServer := &ManagementServer{
+		Address: lis.Addr().String(),
+		cancel:  cancel,
+		version: 0,
+		gs:      gs,
+		xs:      xs,
+		cache:   cache,
+		logger:  t,
+	}
+	if opts.SupportLoadReportingService {
+		lrs := fakeserver.NewServer(lis.Addr().String())
+		v3lrsgrpc.RegisterLoadReportingServiceServer(gs, lrs)
+		mgmtServer.LRSServer = lrs
+		t.Logf("Registered Load Reporting Service (LRS)...")
+	}
+
+	// Start serving.
+	go gs.Serve(lis)
+	t.Logf("xDS management server serving at: %v...", lis.Addr().String())
+	t.Cleanup(mgmtServer.Stop)
+	return mgmtServer
+}
+
+// UpdateOptions wraps parameters to be passed to the Update() method.
+type UpdateOptions struct {
+	// NodeID is the id of the client to which this update is to be pushed.
+	NodeID string
+	// Endpoints, Clusters, Routes, and Listeners are the updated list of xds
+	// resources for the server.  All must be provided with each Update.
+	Endpoints []*v3endpointpb.ClusterLoadAssignment
+	Clusters  []*v3clusterpb.Cluster
+	Routes    []*v3routepb.RouteConfiguration
+	Listeners []*v3listenerpb.Listener
+	// SkipValidation indicates whether we want to skip validation (by not
+	// calling snapshot.Consistent()). It can be useful for negative tests,
+	// where we send updates that the client will NACK.
+	SkipValidation bool
+}
+
+// Update changes the resource snapshot held by the management server, which
+// updates connected clients as required.
+func (s *ManagementServer) Update(ctx context.Context, opts UpdateOptions) error {
+	s.version++
+
+	// Create a snapshot with the passed in resources.
+	resources := map[v3resource.Type][]types.Resource{
+		v3resource.ListenerType: resourceSlice(opts.Listeners),
+		v3resource.RouteType:    resourceSlice(opts.Routes),
+		v3resource.ClusterType:  resourceSlice(opts.Clusters),
+		v3resource.EndpointType: resourceSlice(opts.Endpoints),
+	}
+	snapshot, err := v3cache.NewSnapshot(strconv.Itoa(s.version), resources)
+	if err != nil {
+		return fmt.Errorf("failed to create new snapshot cache: %v", err)
+
+	}
+	if !opts.SkipValidation {
+		if err := snapshot.Consistent(); err != nil {
+			return fmt.Errorf("failed to create new resource snapshot: %v", err)
+		}
+	}
+	s.logger.Logf("Created new resource snapshot...")
+
+	// Update the cache with the new resource snapshot.
+	if err := s.cache.SetSnapshot(ctx, opts.NodeID, snapshot); err != nil {
+		return fmt.Errorf("failed to update resource snapshot in management server: %v", err)
+	}
+	s.logger.Logf("Updated snapshot cache with resource snapshot...")
+	return nil
+}
+
+// Stop stops the management server.
+func (s *ManagementServer) Stop() {
+	if s.cancel != nil {
+		s.cancel()
+	}
+	s.gs.Stop()
+}
+
+// resourceSlice accepts a slice of any type of proto messages and returns a
+// slice of types.Resource.  Will panic if there is an input type mismatch.
+func resourceSlice(i any) []types.Resource {
+	v := reflect.ValueOf(i)
+	rs := make([]types.Resource, v.Len())
+	for i := 0; i < v.Len(); i++ {
+		rs[i] = v.Index(i).Interface().(types.Resource)
+	}
+	return rs
+}

internal/testutils/xds/e2e/setup/setup.go

@@ -0,0 +1,85 @@
+/*
+ *
+ * Copyright 2024 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 setup implements setup helpers for xDS e2e tests.
+package setup
+
+import (
+	"testing"
+
+	"github.com/google/uuid"
+	"google.golang.org/grpc/internal"
+	"google.golang.org/grpc/internal/testutils/xds/e2e"
+	"google.golang.org/grpc/resolver"
+	_ "google.golang.org/grpc/xds" // Register the xds_resolver.
+)
+
+// ManagementServerAndResolver sets up an xDS management server, creates
+// bootstrap configuration pointing to that server and creates an xDS resolver
+// using that configuration.
+//
+// Registers a cleanup function on t to stop the management server.
+//
+// Returns the following:
+// - the xDS management server
+// - the node ID to use when talking to this management server
+// - bootstrap configuration to use (if creating an xDS-enabled gRPC server)
+// - xDS resolver builder (if creating an xDS-enabled gRPC client)
+func ManagementServerAndResolver(t *testing.T) (*e2e.ManagementServer, string, []byte, resolver.Builder) {
+	// Start an xDS management server.
+	xdsServer := e2e.StartManagementServer(t, e2e.ManagementServerOptions{AllowResourceSubset: true})
+
+	// Create bootstrap configuration pointing to the above management server.
+	nodeID := uuid.New().String()
+	bc := e2e.DefaultBootstrapContents(t, nodeID, xdsServer.Address)
+
+	// Create an xDS resolver with the above bootstrap configuration.
+	if internal.NewXDSResolverWithConfigForTesting == nil {
+		t.Fatalf("internal.NewXDSResolverWithConfigForTesting is nil")
+	}
+	r, err := internal.NewXDSResolverWithConfigForTesting.(func([]byte) (resolver.Builder, error))(bc)
+	if err != nil {
+		t.Fatalf("Failed to create xDS resolver for testing: %v", err)
+	}
+
+	return xdsServer, nodeID, bc, r
+}
+
+// ManagementServerAndResolverWithSPIFFE is exactly the same as
+// ManagementServerAndResolver, except that it uses a bootstrap configuration
+// containing certificate providers utilizing SPIFFE test certificates.
+func ManagementServerAndResolverWithSPIFFE(t *testing.T) (*e2e.ManagementServer,
+	string, []byte, resolver.Builder) {
+	// Start an xDS management server.
+	xdsServer := e2e.StartManagementServer(t, e2e.ManagementServerOptions{AllowResourceSubset: true})
+
+	// Create bootstrap configuration pointing to the above management server.
+	nodeID := uuid.New().String()
+	bc := e2e.SPIFFEBootstrapContents(t, nodeID, xdsServer.Address)
+
+	// Create an xDS resolver with the above bootstrap configuration.
+	if internal.NewXDSResolverWithConfigForTesting == nil {
+		t.Fatalf("internal.NewXDSResolverWithConfigForTesting is nil")
+	}
+	r, err := internal.NewXDSResolverWithConfigForTesting.(func([]byte) (resolver.Builder, error))(bc)
+	if err != nil {
+		t.Fatalf("Failed to create xDS resolver for testing: %v", err)
+	}
+
+	return xdsServer, nodeID, bc, r
+}

internal/testutils/xds/fakeserver/server.go

@@ -0,0 +1,254 @@
+/*
+ *
+ * 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 fakeserver provides a fake implementation of the management server.
+//
+// This package is recommended only for scenarios which cannot be tested using
+// the xDS management server (which uses envoy-go-control-plane) provided by the
+// `internal/testutils/xds/e2e` package.
+package fakeserver
+
+import (
+	"fmt"
+	"io"
+	"net"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/internal/testutils"
+	"google.golang.org/grpc/status"
+	"google.golang.org/protobuf/proto"
+
+	v3discoverygrpc "github.com/envoyproxy/go-control-plane/envoy/service/discovery/v3"
+	v3discoverypb "github.com/envoyproxy/go-control-plane/envoy/service/discovery/v3"
+	v3lrsgrpc "github.com/envoyproxy/go-control-plane/envoy/service/load_stats/v3"
+	v3lrspb "github.com/envoyproxy/go-control-plane/envoy/service/load_stats/v3"
+)
+
+const (
+	// TODO: Make this a var or a field in the server if there is a need to use a
+	// value other than this default.
+	defaultChannelBufferSize = 50
+	defaultDialTimeout       = 5 * time.Second
+)
+
+// Request wraps the request protobuf (xds/LRS) and error received by the
+// Server in a call to stream.Recv().
+type Request struct {
+	Req proto.Message
+	Err error
+}
+
+// Response wraps the response protobuf (xds/LRS) and error that the Server
+// should send out to the client through a call to stream.Send()
+type Response struct {
+	Resp proto.Message
+	Err  error
+}
+
+// Server is a fake implementation of xDS and LRS protocols. It listens on the
+// same port for both services and exposes a bunch of channels to send/receive
+// messages.
+//
+// This server is recommended only for scenarios which cannot be tested using
+// the xDS management server (which uses envoy-go-control-plane) provided by the
+// `internal/testutils/xds/e2e` package.
+type Server struct {
+	// XDSRequestChan is a channel on which received xDS requests are made
+	// available to the users of this Server.
+	XDSRequestChan *testutils.Channel
+	// XDSResponseChan is a channel on which the Server accepts xDS responses
+	// to be sent to the client.
+	XDSResponseChan chan *Response
+	// LRSRequestChan is a channel on which received LRS requests are made
+	// available to the users of this Server.
+	LRSRequestChan *testutils.Channel
+	// LRSResponseChan is a channel on which the Server accepts the LRS
+	// response to be sent to the client.
+	LRSResponseChan chan *Response
+	// LRSStreamOpenChan is a channel on which the Server sends notifications
+	// when a new LRS stream is created.
+	LRSStreamOpenChan *testutils.Channel
+	// LRSStreamCloseChan is a channel on which the Server sends notifications
+	// when an existing LRS stream is closed.
+	LRSStreamCloseChan *testutils.Channel
+	// NewConnChan is a channel on which the fake server notifies receipt of new
+	// connection attempts. Tests can gate on this event before proceeding to
+	// other actions which depend on a connection to the fake server being up.
+	NewConnChan *testutils.Channel
+	// Address is the host:port on which the Server is listening for requests.
+	Address string
+
+	// The underlying fake implementation of xDS and LRS.
+	*xdsServerV3
+	*lrsServerV3
+}
+
+type wrappedListener struct {
+	net.Listener
+	server *Server
+}
+
+func (wl *wrappedListener) Accept() (net.Conn, error) {
+	c, err := wl.Listener.Accept()
+	if err != nil {
+		return nil, err
+	}
+	wl.server.NewConnChan.Send(struct{}{})
+	return c, err
+}
+
+// StartServer makes a new Server and gets it to start listening on the given
+// net.Listener. If the given net.Listener is nil, a new one is created on a
+// local port for gRPC requests. The returned cancel function should be invoked
+// by the caller upon completion of the test.
+func StartServer(lis net.Listener) (*Server, func(), error) {
+	if lis == nil {
+		var err error
+		lis, err = net.Listen("tcp", "localhost:0")
+		if err != nil {
+			return nil, func() {}, fmt.Errorf("net.Listen() failed: %v", err)
+		}
+	}
+
+	s := NewServer(lis.Addr().String())
+	wp := &wrappedListener{
+		Listener: lis,
+		server:   s,
+	}
+
+	server := grpc.NewServer()
+	v3lrsgrpc.RegisterLoadReportingServiceServer(server, s)
+	v3discoverygrpc.RegisterAggregatedDiscoveryServiceServer(server, s)
+	go server.Serve(wp)
+
+	return s, func() { server.Stop() }, nil
+}
+
+// NewServer returns a new instance of Server, set to accept requests on addr.
+// It is the responsibility of the caller to register the exported ADS and LRS
+// services on an appropriate gRPC server. Most usages should prefer
+// StartServer() instead of this.
+func NewServer(addr string) *Server {
+	s := &Server{
+		XDSRequestChan:     testutils.NewChannelWithSize(defaultChannelBufferSize),
+		LRSRequestChan:     testutils.NewChannelWithSize(defaultChannelBufferSize),
+		NewConnChan:        testutils.NewChannelWithSize(defaultChannelBufferSize),
+		XDSResponseChan:    make(chan *Response, defaultChannelBufferSize),
+		LRSResponseChan:    make(chan *Response, 1), // The server only ever sends one response.
+		LRSStreamOpenChan:  testutils.NewChannelWithSize(defaultChannelBufferSize),
+		LRSStreamCloseChan: testutils.NewChannelWithSize(defaultChannelBufferSize),
+		Address:            addr,
+	}
+	s.xdsServerV3 = &xdsServerV3{reqChan: s.XDSRequestChan, respChan: s.XDSResponseChan}
+	s.lrsServerV3 = &lrsServerV3{reqChan: s.LRSRequestChan, respChan: s.LRSResponseChan, streamOpenChan: s.LRSStreamOpenChan, streamCloseChan: s.LRSStreamCloseChan}
+	return s
+}
+
+type xdsServerV3 struct {
+	reqChan  *testutils.Channel
+	respChan chan *Response
+}
+
+func (xdsS *xdsServerV3) StreamAggregatedResources(s v3discoverygrpc.AggregatedDiscoveryService_StreamAggregatedResourcesServer) error {
+	errCh := make(chan error, 2)
+	go func() {
+		for {
+			req, err := s.Recv()
+			if err != nil {
+				errCh <- err
+				return
+			}
+			xdsS.reqChan.Send(&Request{req, err})
+		}
+	}()
+	go func() {
+		var retErr error
+		defer func() {
+			errCh <- retErr
+		}()
+
+		for {
+			select {
+			case r := <-xdsS.respChan:
+				if r.Err != nil {
+					retErr = r.Err
+					return
+				}
+				if err := s.Send(r.Resp.(*v3discoverypb.DiscoveryResponse)); err != nil {
+					retErr = err
+					return
+				}
+			case <-s.Context().Done():
+				retErr = s.Context().Err()
+				return
+			}
+		}
+	}()
+
+	if err := <-errCh; err != nil {
+		return err
+	}
+	return nil
+}
+
+func (xdsS *xdsServerV3) DeltaAggregatedResources(v3discoverygrpc.AggregatedDiscoveryService_DeltaAggregatedResourcesServer) error {
+	return status.Error(codes.Unimplemented, "")
+}
+
+type lrsServerV3 struct {
+	reqChan         *testutils.Channel
+	respChan        chan *Response
+	streamOpenChan  *testutils.Channel
+	streamCloseChan *testutils.Channel
+}
+
+func (lrsS *lrsServerV3) StreamLoadStats(s v3lrsgrpc.LoadReportingService_StreamLoadStatsServer) error {
+	lrsS.streamOpenChan.Send(nil)
+	defer lrsS.streamCloseChan.Send(nil)
+
+	req, err := s.Recv()
+	lrsS.reqChan.Send(&Request{req, err})
+	if err != nil {
+		return err
+	}
+
+	select {
+	case r := <-lrsS.respChan:
+		if r.Err != nil {
+			return r.Err
+		}
+		if err := s.Send(r.Resp.(*v3lrspb.LoadStatsResponse)); err != nil {
+			return err
+		}
+	case <-s.Context().Done():
+		return s.Context().Err()
+	}
+
+	for {
+		req, err := s.Recv()
+		lrsS.reqChan.Send(&Request{req, err})
+		if err != nil {
+			if err == io.EOF {
+				return nil
+			}
+			return err
+		}
+	}
+}

internal/testutils/xds_bootstrap.go

@@ -0,0 +1,52 @@
+/*
+ *
+ * Copyright 2024 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 testutils
+
+import (
+	"os"
+	"testing"
+
+	"google.golang.org/grpc/internal/envconfig"
+)
+
+// CreateBootstrapFileForTesting creates a temporary file with the provided
+// bootstrap contents, and updates the bootstrap environment variable to point
+// to this file.
+//
+// Registers a cleanup function on the provided testing.T, that deletes the
+// temporary file and resets the bootstrap environment variable.
+func CreateBootstrapFileForTesting(t *testing.T, bootstrapContents []byte) {
+	t.Helper()
+
+	f, err := os.CreateTemp("", "test_xds_bootstrap_*")
+	if err != nil {
+		t.Fatalf("Failed to created bootstrap file: %v", err)
+	}
+
+	if err := os.WriteFile(f.Name(), bootstrapContents, 0644); err != nil {
+		t.Fatalf("Failed to created bootstrap file: %v", err)
+	}
+
+	origBootstrapFileName := envconfig.XDSBootstrapFileName
+	envconfig.XDSBootstrapFileName = f.Name()
+	t.Cleanup(func() {
+		os.Remove(f.Name())
+		envconfig.XDSBootstrapFileName = origBootstrapFileName
+	})
+}