From d30e2c91a0545bd393774c3775cd9f9c5f5a5673 Mon Sep 17 00:00:00 2001
From: Doug Fawley <dfawley@google.com>
Date: Mon, 7 Jun 2021 17:13:48 -0700
Subject: [PATCH] xds/resolver: test xds client closed by resolver Close
 (#4509)

---
 xds/internal/resolver/xds_resolver_test.go  | 14 ++++++++++++++
 xds/internal/testutils/fakeclient/client.go | 17 ++++++-----------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/xds/internal/resolver/xds_resolver_test.go b/xds/internal/resolver/xds_resolver_test.go
index d3a66595a35..d588ff157cd 100644
--- a/xds/internal/resolver/xds_resolver_test.go
+++ b/xds/internal/resolver/xds_resolver_test.go
@@ -281,6 +281,20 @@ func (s) TestXDSResolverWatchCallbackAfterClose(t *testing.T) {
 	}
 }
 
+// TestXDSResolverCloseClosesXDSClient tests that the XDS resolver's Close
+// method closes the XDS client.
+func (s) TestXDSResolverCloseClosesXDSClient(t *testing.T) {
+	xdsC := fakeclient.NewClient()
+	xdsR, _, cancel := testSetup(t, setupOpts{
+		xdsClientFunc: func() (xdsClient, error) { return xdsC, nil },
+	})
+	defer cancel()
+	xdsR.Close()
+	if !xdsC.Closed.HasFired() {
+		t.Fatalf("xds client not closed by xds resolver Close method")
+	}
+}
+
 // TestXDSResolverBadServiceUpdate tests the case the xdsClient returns a bad
 // service update.
 func (s) TestXDSResolverBadServiceUpdate(t *testing.T) {
diff --git a/xds/internal/testutils/fakeclient/client.go b/xds/internal/testutils/fakeclient/client.go
index ce0996cb8e9..37e84f998b9 100644
--- a/xds/internal/testutils/fakeclient/client.go
+++ b/xds/internal/testutils/fakeclient/client.go
@@ -22,6 +22,7 @@ package fakeclient
 import (
 	"context"
 
+	"google.golang.org/grpc/internal/grpcsync"
 	"google.golang.org/grpc/internal/testutils"
 	"google.golang.org/grpc/xds/internal/xdsclient"
 	"google.golang.org/grpc/xds/internal/xdsclient/bootstrap"
@@ -41,7 +42,6 @@ type Client struct {
 	cdsCancelCh  *testutils.Channel
 	edsCancelCh  *testutils.Channel
 	loadReportCh *testutils.Channel
-	closeCh      *testutils.Channel
 	loadStore    *load.Store
 	bootstrapCfg *bootstrap.Config
 
@@ -49,6 +49,8 @@ type Client struct {
 	rdsCb  func(xdsclient.RouteConfigUpdate, error)
 	cdsCbs map[string]func(xdsclient.ClusterUpdate, error)
 	edsCb  func(xdsclient.EndpointsUpdate, error)
+
+	Closed *grpcsync.Event // fired when Close is called.
 }
 
 // WatchListener registers a LDS watch.
@@ -228,16 +230,9 @@ func (xdsC *Client) WaitForReportLoad(ctx context.Context) (ReportLoadArgs, erro
 	return val.(ReportLoadArgs), err
 }
 
-// Close closes the xds client.
+// Close fires xdsC.Closed, indicating it was called.
 func (xdsC *Client) Close() {
-	xdsC.closeCh.Send(nil)
-}
-
-// WaitForClose waits for Close to be invoked on this client and returns
-// context.DeadlineExceeded otherwise.
-func (xdsC *Client) WaitForClose(ctx context.Context) error {
-	_, err := xdsC.closeCh.Receive(ctx)
-	return err
+	xdsC.Closed.Fire()
 }
 
 // BootstrapConfig returns the bootstrap config.
@@ -275,8 +270,8 @@ func NewClientWithName(name string) *Client {
 		cdsCancelCh:  testutils.NewChannelWithSize(10),
 		edsCancelCh:  testutils.NewChannel(),
 		loadReportCh: testutils.NewChannel(),
-		closeCh:      testutils.NewChannel(),
 		loadStore:    load.NewStore(),
 		cdsCbs:       make(map[string]func(xdsclient.ClusterUpdate, error)),
+		Closed:       grpcsync.NewEvent(),
 	}
 }