From c4c8b113050ecd856b118b5a58bda24e17763d59 Mon Sep 17 00:00:00 2001
From: Easwar Swaminathan <easwars@google.com>
Date: Thu, 24 Oct 2024 10:02:09 -0700
Subject: [PATCH] xds/resolver: add a way to specify the xDS client to use for
 testing purposes (#7771)

---
 internal/internal.go                  | 14 ++++++++++++++
 xds/internal/resolver/xds_resolver.go | 24 +++++++++++++++++++-----
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/internal/internal.go b/internal/internal.go
index 20b4dc3d3536..88900fa9bbc6 100644
--- a/internal/internal.go
+++ b/internal/internal.go
@@ -149,6 +149,20 @@ var (
 	// other features, including the CSDS service.
 	NewXDSResolverWithConfigForTesting any // func([]byte) (resolver.Builder, error)
 
+	// NewXDSResolverWithClientForTesting creates a new xDS resolver builder
+	// using the provided xDS client instead of creating a new one using the
+	// bootstrap configuration specified by the supported environment variables.
+	// The resolver.Builder is meant to be used in conjunction with the
+	// grpc.WithResolvers DialOption. The resolver.Builder does not take
+	// ownership of the provided xDS client and it is the responsibility of the
+	// caller to close the client when no longer required.
+	//
+	// Testing Only
+	//
+	// This function should ONLY be used for testing and may not work with some
+	// other features, including the CSDS service.
+	NewXDSResolverWithClientForTesting any // func(xdsclient.XDSClient) (resolver.Builder, error)
+
 	// RegisterRLSClusterSpecifierPluginForTesting registers the RLS Cluster
 	// Specifier Plugin for testing purposes, regardless of the XDSRLS environment
 	// variable.
diff --git a/xds/internal/resolver/xds_resolver.go b/xds/internal/resolver/xds_resolver.go
index b5d24e4bf214..de339a7c9b69 100644
--- a/xds/internal/resolver/xds_resolver.go
+++ b/xds/internal/resolver/xds_resolver.go
@@ -44,10 +44,10 @@ import (
 // xdsresolver.Scheme
 const Scheme = "xds"
 
-// newBuilderForTesting creates a new xds resolver builder using a specific xds
-// bootstrap config, so tests can use multiple xds clients in different
-// ClientConns at the same time.
-func newBuilderForTesting(config []byte) (resolver.Builder, error) {
+// newBuilderWithConfigForTesting creates a new xds resolver builder using a
+// specific xds bootstrap config, so tests can use multiple xds clients in
+// different ClientConns at the same time.
+func newBuilderWithConfigForTesting(config []byte) (resolver.Builder, error) {
 	return &xdsResolverBuilder{
 		newXDSClient: func(name string) (xdsclient.XDSClient, func(), error) {
 			return xdsclient.NewForTesting(xdsclient.OptionsForTesting{Name: name, Contents: config})
@@ -55,9 +55,23 @@ func newBuilderForTesting(config []byte) (resolver.Builder, error) {
 	}, nil
 }
 
+// newBuilderWithClientForTesting creates a new xds resolver builder using the
+// specific xDS client, so that tests have complete control over the exact
+// specific xDS client being used.
+func newBuilderWithClientForTesting(client xdsclient.XDSClient) (resolver.Builder, error) {
+	return &xdsResolverBuilder{
+		newXDSClient: func(string) (xdsclient.XDSClient, func(), error) {
+			// Returning an empty close func here means that the responsibility
+			// of closing the client lies with the caller.
+			return client, func() {}, nil
+		},
+	}, nil
+}
+
 func init() {
 	resolver.Register(&xdsResolverBuilder{})
-	internal.NewXDSResolverWithConfigForTesting = newBuilderForTesting
+	internal.NewXDSResolverWithConfigForTesting = newBuilderWithConfigForTesting
+	internal.NewXDSResolverWithClientForTesting = newBuilderWithClientForTesting
 
 	rinternal.NewWRR = wrr.NewRandom
 	rinternal.NewXDSClient = xdsclient.New