From d82e137c7bf94301c95b38018d637b0166e199d7 Mon Sep 17 00:00:00 2001
From: Juanli Shen <juanlishen@google.com>
Date: Thu, 1 Feb 2018 14:47:38 -0800
Subject: Add a results_upon_error setter to fake resolver

---
 test/cpp/end2end/grpclb_end2end_test.cc | 23 +++++++++++++++++++++--
 1 file changed, 21 insertions(+), 2 deletions(-)

(limited to 'test/cpp/end2end')

diff --git a/test/cpp/end2end/grpclb_end2end_test.cc b/test/cpp/end2end/grpclb_end2end_test.cc
index 78527587cf..5591acf880 100644
--- a/test/cpp/end2end/grpclb_end2end_test.cc
+++ b/test/cpp/end2end/grpclb_end2end_test.cc
@@ -454,8 +454,8 @@ class GrpclbEnd2endTest : public ::testing::Test {
     grpc::string balancer_name;
   };
 
-  void SetNextResolution(const std::vector<AddressData>& address_data) {
-    grpc_core::ExecCtx exec_ctx;
+  grpc_lb_addresses* CreateLbAddressesFromAddressDataList(
+      const std::vector<AddressData>& address_data) {
     grpc_lb_addresses* addresses =
         grpc_lb_addresses_create(address_data.size(), nullptr);
     for (size_t i = 0; i < address_data.size(); ++i) {
@@ -469,6 +469,13 @@ class GrpclbEnd2endTest : public ::testing::Test {
       grpc_uri_destroy(lb_uri);
       gpr_free(lb_uri_str);
     }
+    return addresses;
+  }
+
+  void SetNextResolution(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(response_generator_,
@@ -476,6 +483,18 @@ class GrpclbEnd2endTest : public ::testing::Test {
     grpc_lb_addresses_destroy(addresses);
   }
 
+  void SetNextResolutionUponError(
+      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);
+    grpc_lb_addresses_destroy(addresses);
+  }
+
   const std::vector<int> GetBackendPorts(const size_t start_index = 0) const {
     std::vector<int> backend_ports;
     for (size_t i = start_index; i < backend_servers_.size(); ++i) {
-- 
cgit v1.2.3