From 9a4ed686b1a0a0f3b655a34e61151b1dcb0db744 Mon Sep 17 00:00:00 2001
From: David Garcia Quintas <dgq@google.com>
Date: Thu, 13 Jul 2017 16:12:30 -0700
Subject: [PATCH] Improved testing

---
 .../ext/filters/client_channel/subchannel_index.c   |  7 +++++++
 .../ext/filters/client_channel/subchannel_index.h   | 12 ++++++++++++
 test/cpp/end2end/client_lb_end2end_test.cc          | 13 +++++++++----
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/src/core/ext/filters/client_channel/subchannel_index.c b/src/core/ext/filters/client_channel/subchannel_index.c
index e291ca9db99..a33ab950bff 100644
--- a/src/core/ext/filters/client_channel/subchannel_index.c
+++ b/src/core/ext/filters/client_channel/subchannel_index.c
@@ -40,6 +40,8 @@ struct grpc_subchannel_key {
 
 GPR_TLS_DECL(subchannel_index_exec_ctx);
 
+static bool g_force_creation = false;
+
 static void enter_ctx(grpc_exec_ctx *exec_ctx) {
   GPR_ASSERT(gpr_tls_get(&subchannel_index_exec_ctx) == 0);
   gpr_tls_set(&subchannel_index_exec_ctx, (intptr_t)exec_ctx);
@@ -84,6 +86,7 @@ static grpc_subchannel_key *subchannel_key_copy(grpc_subchannel_key *k) {
 
 int grpc_subchannel_key_compare(const grpc_subchannel_key *a,
                                 const grpc_subchannel_key *b) {
+  if (g_force_creation) return false;
   int c = GPR_ICMP(a->args.filter_count, b->args.filter_count);
   if (c != 0) return c;
   if (a->args.filter_count > 0) {
@@ -250,3 +253,7 @@ void grpc_subchannel_index_unregister(grpc_exec_ctx *exec_ctx,
 
   leave_ctx(exec_ctx);
 }
+
+void grpc_subchannel_index_test_only_set_force_creation(bool force_creation) {
+  g_force_creation = force_creation;
+}
diff --git a/src/core/ext/filters/client_channel/subchannel_index.h b/src/core/ext/filters/client_channel/subchannel_index.h
index e303bfaa05f..98d882a453f 100644
--- a/src/core/ext/filters/client_channel/subchannel_index.h
+++ b/src/core/ext/filters/client_channel/subchannel_index.h
@@ -59,4 +59,16 @@ void grpc_subchannel_index_init(void);
 /** Shutdown the subchannel index (global) */
 void grpc_subchannel_index_shutdown(void);
 
+/** \em TEST ONLY.
+ * If \a force_creation is true, all key comparisons will be false, resulting in
+ * new subchannels always being created. Otherwise, the keys will be compared as
+ * usual.
+ *
+ * This function is *not* threadsafe on purpose: it should *only* be used in
+ * test code.
+ *
+ * Tests using this function \em MUST run tests with and without \a
+ * force_creation set. */
+void grpc_subchannel_index_test_only_set_force_creation(bool force_creation);
+
 #endif /* GRPC_CORE_EXT_FILTERS_CLIENT_CHANNEL_SUBCHANNEL_INDEX_H */
diff --git a/test/cpp/end2end/client_lb_end2end_test.cc b/test/cpp/end2end/client_lb_end2end_test.cc
index 6d3f5a9d46a..e1160ecdc6d 100644
--- a/test/cpp/end2end/client_lb_end2end_test.cc
+++ b/test/cpp/end2end/client_lb_end2end_test.cc
@@ -35,6 +35,7 @@
 
 extern "C" {
 #include "src/core/ext/filters/client_channel/resolver/fake/fake_resolver.h"
+#include "src/core/ext/filters/client_channel/subchannel_index.h"
 }
 
 #include "src/proto/grpc/testing/echo.grpc.pb.h"
@@ -331,10 +332,14 @@ TEST_F(ClientLbEnd2endTest, PickFirstManyUpdates) {
   for (size_t i = 0; i < servers_.size(); ++i) {
     ports.emplace_back(servers_[i]->port_);
   }
-  for (size_t i = 0; i < 1000; ++i) {
-    std::random_shuffle(ports.begin(), ports.end());
-    SetNextResolution(ports);
-    if (i % 10 == 0) SendRpc();
+  for (const bool force_creation : {true, false}) {
+    grpc_subchannel_index_test_only_set_force_creation(force_creation);
+    gpr_log(GPR_INFO, "Force subchannel creation: %d", force_creation);
+    for (size_t i = 0; i < 1000; ++i) {
+      std::random_shuffle(ports.begin(), ports.end());
+      SetNextResolution(ports);
+      if (i % 10 == 0) SendRpc();
+    }
   }
   // Check LB policy name for the channel.
   EXPECT_EQ("pick_first", channel_->GetLoadBalancingPolicyName());