From a2e2f565651f80ffefc58f4b1527203e549c751c Mon Sep 17 00:00:00 2001
From: yifeizhuang <zivy@google.com>
Date: Thu, 7 Oct 2021 16:17:08 -0700
Subject: [PATCH] xds: override bootstrap for xds server (#8575)

added xdsServerBuilder method `overrideBootstrapForTest()`. Fix issue https://github.com/grpc/grpc-java/issues/7819
---
 xds/src/main/java/io/grpc/xds/XdsServerBuilder.java | 13 +++++++++++++
 .../test/java/io/grpc/xds/XdsServerBuilderTest.java | 13 ++++++++++++-
 .../test/java/io/grpc/xds/XdsServerTestHelper.java  |  3 ++-
 3 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/xds/src/main/java/io/grpc/xds/XdsServerBuilder.java b/xds/src/main/java/io/grpc/xds/XdsServerBuilder.java
index c95c1e6d48..d4df317a7e 100644
--- a/xds/src/main/java/io/grpc/xds/XdsServerBuilder.java
+++ b/xds/src/main/java/io/grpc/xds/XdsServerBuilder.java
@@ -37,6 +37,7 @@ import io.grpc.netty.InternalProtocolNegotiator;
 import io.grpc.netty.NettyServerBuilder;
 import io.grpc.xds.FilterChainMatchingProtocolNegotiators.FilterChainMatchingNegotiatorServerFactory;
 import io.grpc.xds.XdsNameResolverProvider.XdsClientPoolFactory;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.logging.Logger;
@@ -136,6 +137,18 @@ public final class XdsServerBuilder extends ForwardingServerBuilder<XdsServerBui
     return this;
   }
 
+  /**
+   * Allows providing bootstrap override, useful for testing.
+   */
+  public XdsServerBuilder overrideBootstrapForTest(Map<String, ?> bootstrapOverride) {
+    checkNotNull(bootstrapOverride, "bootstrapOverride");
+    if (this.xdsClientPoolFactory == SharedXdsClientPoolProvider.getDefaultProvider()) {
+      this.xdsClientPoolFactory = new SharedXdsClientPoolProvider();
+    }
+    this.xdsClientPoolFactory.setBootstrapOverride(bootstrapOverride);
+    return this;
+  }
+
   /**
    * Returns the delegate {@link NettyServerBuilder} to allow experimental level
    * transport-specific configuration. Note this API will always be experimental.
diff --git a/xds/src/test/java/io/grpc/xds/XdsServerBuilderTest.java b/xds/src/test/java/io/grpc/xds/XdsServerBuilderTest.java
index 0d15c1f660..d67ed9d09f 100644
--- a/xds/src/test/java/io/grpc/xds/XdsServerBuilderTest.java
+++ b/xds/src/test/java/io/grpc/xds/XdsServerBuilderTest.java
@@ -40,7 +40,9 @@ import java.io.IOException;
 import java.net.InetSocketAddress;
 import java.net.ServerSocket;
 import java.net.SocketAddress;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -65,6 +67,7 @@ public class XdsServerBuilderTest {
   private int port;
   private TlsContextManager tlsContextManager;
   private FakeXdsClient xdsClient = new FakeXdsClient();
+  private FakeXdsClientPoolFactory xdsClientPoolFactory = new FakeXdsClientPoolFactory(xdsClient);
 
   private void buildServer(XdsServerBuilder.XdsServingStatusListener xdsServingStatusListener)
       throws IOException {
@@ -77,7 +80,7 @@ public class XdsServerBuilderTest {
     builder =
         XdsServerBuilder.forPort(
             port, XdsServerCredentials.create(InsecureServerCredentials.create()));
-    builder.xdsClientPoolFactory(new FakeXdsClientPoolFactory(xdsClient));
+    builder.xdsClientPoolFactory(xdsClientPoolFactory);
     if (xdsServingStatusListener != null) {
       builder.xdsServingStatusListener(xdsServingStatusListener);
     }
@@ -292,4 +295,12 @@ public class XdsServerBuilderTest {
       assertThat(expected).hasMessageThat().contains("drain grace time");
     }
   }
+
+  @Test
+  public void testOverrideBootstrap() throws Exception {
+    Map<String, Object> b = new HashMap<>();
+    buildBuilder(null);
+    builder.overrideBootstrapForTest(b);
+    assertThat(xdsClientPoolFactory.savedBootstrap).isEqualTo(b);
+  }
 }
diff --git a/xds/src/test/java/io/grpc/xds/XdsServerTestHelper.java b/xds/src/test/java/io/grpc/xds/XdsServerTestHelper.java
index f289c4726f..ffe4a72f52 100644
--- a/xds/src/test/java/io/grpc/xds/XdsServerTestHelper.java
+++ b/xds/src/test/java/io/grpc/xds/XdsServerTestHelper.java
@@ -113,6 +113,7 @@ public class XdsServerTestHelper {
         implements XdsNameResolverProvider.XdsClientPoolFactory {
 
     private XdsClient xdsClient;
+    Map<String, ?> savedBootstrap;
 
     FakeXdsClientPoolFactory(XdsClient xdsClient) {
       this.xdsClient = xdsClient;
@@ -120,7 +121,7 @@ public class XdsServerTestHelper {
 
     @Override
     public void setBootstrapOverride(Map<String, ?> bootstrap) {
-      throw new UnsupportedOperationException("Should not be called");
+      this.savedBootstrap = bootstrap;
     }
 
     @Override