From f8ea440d76b8ddd11e858af65c7fd1c7d5dbe8e8 Mon Sep 17 00:00:00 2001
From: Menghan Li <menghanl@google.com>
Date: Wed, 3 Nov 2021 13:29:17 -0700
Subject: [PATCH] [federation_target_parsing] add test for client close

---
 xds/internal/resolver/xds_resolver_test.go | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/xds/internal/resolver/xds_resolver_test.go b/xds/internal/resolver/xds_resolver_test.go
index ee0e2d91000b..adeb37159bdf 100644
--- a/xds/internal/resolver/xds_resolver_test.go
+++ b/xds/internal/resolver/xds_resolver_test.go
@@ -193,13 +193,20 @@ func (s) TestResolverBuilder(t *testing.T) {
 func (s) TestResolverBuilder_xdsCredsBootstrapMismatch(t *testing.T) {
 	// Fake out the xdsClient creation process by providing a fake, which does
 	// not have any certificate provider configuration.
+	fc := fakeclient.NewClient()
+	fc.SetBootstrapConfig(&bootstrap.Config{})
 	oldClientMaker := newXDSClient
 	newXDSClient = func() (xdsclient.XDSClient, error) {
-		fc := fakeclient.NewClient()
-		fc.SetBootstrapConfig(&bootstrap.Config{})
 		return fc, nil
 	}
 	defer func() { newXDSClient = oldClientMaker }()
+	defer func() {
+		select {
+		case <-time.After(defaultTestTimeout):
+			t.Fatalf("timeout waiting for close")
+		case <-fc.Closed.Done():
+		}
+	}()
 
 	builder := resolver.Get(xdsScheme)
 	if builder == nil {