2 files changed, 21 insertions, 27 deletions

diff --git a/test/cpp/end2end/client_lb_end2end_test.cc b/test/cpp/end2end/client_lb_end2end_test.cc
index ee5adbc6fa..386d18e95d 100644
--- a/test/cpp/end2end/client_lb_end2end_test.cc
+++ b/test/cpp/end2end/client_lb_end2end_test.cc
@@ -40,6 +40,7 @@
 #include "src/core/lib/backoff/backoff.h"
 #include "src/core/lib/gpr/env.h"
 #include "src/core/lib/gprpp/debug_location.h"
+#include "src/core/lib/gprpp/ref_counted_ptr.h"
 
 #include "src/proto/grpc/testing/echo.grpc.pb.h"
 #include "test/core/util/port.h"
@@ -119,11 +120,11 @@ class ClientLbEnd2endTest : public ::testing::Test {
   }
 
   void SetUp() override {
-    response_generator_ = grpc_fake_resolver_response_generator_create();
+    response_generator_ =
+        grpc_core::MakeRefCounted<grpc_core::FakeResolverResponseGenerator>();
   }
 
   void TearDown() override {
-    grpc_fake_resolver_response_generator_unref(response_generator_);
     for (size_t i = 0; i < servers_.size(); ++i) {
       servers_[i]->Shutdown();
     }
@@ -153,13 +154,9 @@ class ClientLbEnd2endTest : public ::testing::Test {
       grpc_uri_destroy(lb_uri);
       gpr_free(lb_uri_str);
     }
-    const grpc_arg fake_addresses =
-        grpc_lb_addresses_create_channel_arg(addresses);
-    grpc_channel_args* fake_result =
-        grpc_channel_args_copy_and_add(nullptr, &fake_addresses, 1);
-    grpc_fake_resolver_response_generator_set_response(response_generator_,
-                                                       fake_result);
-    grpc_channel_args_destroy(fake_result);
+    grpc_arg fake_addresses = grpc_lb_addresses_create_channel_arg(addresses);
+    grpc_channel_args fake_result = {1, &fake_addresses};
+    response_generator_->SetResponse(&fake_result);
     grpc_lb_addresses_destroy(addresses);
   }
 
@@ -178,13 +175,9 @@ class ClientLbEnd2endTest : public ::testing::Test {
       grpc_uri_destroy(lb_uri);
       gpr_free(lb_uri_str);
     }
-    const grpc_arg fake_addresses =
-        grpc_lb_addresses_create_channel_arg(addresses);
-    grpc_channel_args* fake_result =
-        grpc_channel_args_copy_and_add(nullptr, &fake_addresses, 1);
-    grpc_fake_resolver_response_generator_set_response_upon_error(
-        response_generator_, fake_result);
-    grpc_channel_args_destroy(fake_result);
+    grpc_arg fake_addresses = grpc_lb_addresses_create_channel_arg(addresses);
+    grpc_channel_args fake_result = {1, &fake_addresses};
+    response_generator_->SetReresolutionResponse(&fake_result);
     grpc_lb_addresses_destroy(addresses);
   }
 
@@ -206,7 +199,7 @@ class ClientLbEnd2endTest : public ::testing::Test {
       args.SetLoadBalancingPolicyName(lb_policy_name);
     }  // else, default to pick first
     args.SetPointer(GRPC_ARG_FAKE_RESOLVER_RESPONSE_GENERATOR,
-                    response_generator_);
+                    response_generator_.get());
     return CreateCustomChannel("fake:///", InsecureChannelCredentials(), args);
   }
 
@@ -323,7 +316,8 @@ class ClientLbEnd2endTest : public ::testing::Test {
   const grpc::string server_host_;
   std::unique_ptr<grpc::testing::EchoTestService::Stub> stub_;
   std::vector<std::unique_ptr<ServerData>> servers_;
-  grpc_fake_resolver_response_generator* response_generator_;
+  grpc_core::RefCountedPtr<grpc_core::FakeResolverResponseGenerator>
+      response_generator_;
   const grpc::string kRequestMessage_;
 };
 
diff --git a/test/cpp/end2end/grpclb_end2end_test.cc b/test/cpp/end2end/grpclb_end2end_test.cc
index 5591acf880..89c9c2bf76 100644
--- a/test/cpp/end2end/grpclb_end2end_test.cc
+++ b/test/cpp/end2end/grpclb_end2end_test.cc
@@ -35,6 +35,7 @@
 
 #include "src/core/ext/filters/client_channel/resolver/fake/fake_resolver.h"
 #include "src/core/lib/gpr/env.h"
+#include "src/core/lib/gprpp/ref_counted_ptr.h"
 #include "src/core/lib/iomgr/sockaddr.h"
 
 #include "test/core/util/port.h"
@@ -339,7 +340,8 @@ class GrpclbEnd2endTest : public ::testing::Test {
   }
 
   void SetUp() override {
-    response_generator_ = grpc_fake_resolver_response_generator_create();
+    response_generator_ =
+        grpc_core::MakeRefCounted<grpc_core::FakeResolverResponseGenerator>();
     // Start the backends.
     for (size_t i = 0; i < num_backends_; ++i) {
       backends_.emplace_back(new BackendServiceImpl());
@@ -363,7 +365,6 @@ class GrpclbEnd2endTest : public ::testing::Test {
     for (size_t i = 0; i < balancers_.size(); ++i) {
       if (balancers_[i]->Shutdown()) balancer_servers_[i].Shutdown();
     }
-    grpc_fake_resolver_response_generator_unref(response_generator_);
   }
 
   void SetNextResolutionAllBalancers() {
@@ -378,7 +379,7 @@ class GrpclbEnd2endTest : public ::testing::Test {
     ChannelArguments args;
     args.SetGrpclbFallbackTimeout(fallback_timeout);
     args.SetPointer(GRPC_ARG_FAKE_RESOLVER_RESPONSE_GENERATOR,
-                    response_generator_);
+                    response_generator_.get());
     std::ostringstream uri;
     uri << "fake:///servername_not_used";
     channel_ =
@@ -478,20 +479,18 @@ class GrpclbEnd2endTest : public ::testing::Test {
         CreateLbAddressesFromAddressDataList(address_data);
     grpc_arg fake_addresses = grpc_lb_addresses_create_channel_arg(addresses);
     grpc_channel_args fake_result = {1, &fake_addresses};
-    grpc_fake_resolver_response_generator_set_response(response_generator_,
-                                                       &fake_result);
+    response_generator_->SetResponse(&fake_result);
     grpc_lb_addresses_destroy(addresses);
   }
 
-  void SetNextResolutionUponError(
+  void SetNextReresolutionResponse(
       const std::vector<AddressData>& address_data) {
     grpc_core::ExecCtx exec_ctx;
     grpc_lb_addresses* addresses =
         CreateLbAddressesFromAddressDataList(address_data);
     grpc_arg fake_addresses = grpc_lb_addresses_create_channel_arg(addresses);
     grpc_channel_args fake_result = {1, &fake_addresses};
-    grpc_fake_resolver_response_generator_set_response_upon_error(
-        response_generator_, &fake_result);
+    response_generator_->SetReresolutionResponse(&fake_result);
     grpc_lb_addresses_destroy(addresses);
   }
 
@@ -593,7 +592,8 @@ class GrpclbEnd2endTest : public ::testing::Test {
   std::vector<std::unique_ptr<BalancerServiceImpl>> balancers_;
   std::vector<ServerThread<BackendService>> backend_servers_;
   std::vector<ServerThread<BalancerService>> balancer_servers_;
-  grpc_fake_resolver_response_generator* response_generator_;
+  grpc_core::RefCountedPtr<grpc_core::FakeResolverResponseGenerator>
+      response_generator_;
   const grpc::string kRequestMessage_ = "Live long and prosper.";
 };