From 247c408200f3db8994f1e4d0452bb082776f0d15 Mon Sep 17 00:00:00 2001
From: Alexander Polcyn <apolcyn@google.com>
Date: Thu, 8 Feb 2024 19:57:54 +0000
Subject: [PATCH] Add ExitIdle and CompleteCall helpers to LB policy test lib

---
 .../lb_policy/lb_policy_test_lib.h            | 39 +++++++++++++++----
 .../lb_policy/pick_first_test.cc              | 14 +------
 2 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/test/core/client_channel/lb_policy/lb_policy_test_lib.h b/test/core/client_channel/lb_policy/lb_policy_test_lib.h
index 83ed95e9e09..a1af2dc4a8e 100644
--- a/test/core/client_channel/lb_policy/lb_policy_test_lib.h
+++ b/test/core/client_channel/lb_policy/lb_policy_test_lib.h
@@ -861,6 +861,23 @@ class LoadBalancingPolicyTest : public ::testing::Test {
     return status;
   }
 
+  // Invoke ExitIdle on the LB policy
+  void ExitIdle() {
+    ExecCtx exec_ctx;
+    absl::Notification notification;
+    // Note: ExitIdle() will enqueue a bunch of connectivity state
+    // notifications on the WorkSerializer, and we want to wait until
+    // those are delivered to the LB policy.
+    work_serializer_->Run(
+        [&]() {
+          lb_policy_->ExitIdleLocked();
+          work_serializer_->Run([&]() { notification.Notify(); },
+                                DEBUG_LOCATION);
+        },
+        DEBUG_LOCATION);
+    notification.WaitForNotification();
+  }
+
   void ExpectQueueEmpty(SourceLocation location = SourceLocation()) {
     helper_->ExpectQueueEmpty(location);
   }
@@ -1087,6 +1104,7 @@ class LoadBalancingPolicyTest : public ::testing::Test {
       const CallAttributes& call_attributes = {},
       std::unique_ptr<LoadBalancingPolicy::SubchannelCallTrackerInterface>*
           subchannel_call_tracker = nullptr,
+      SubchannelState::FakeSubchannel** picked_subchannel = nullptr,
       SourceLocation location = SourceLocation()) {
     EXPECT_NE(picker, nullptr);
     if (picker == nullptr) {
@@ -1100,21 +1118,28 @@ class LoadBalancingPolicyTest : public ::testing::Test {
     if (complete == nullptr) return absl::nullopt;
     auto* subchannel = static_cast<SubchannelState::FakeSubchannel*>(
         complete->subchannel.get());
+    if (picked_subchannel != nullptr) *picked_subchannel = subchannel;
     std::string address = subchannel->state()->address();
     if (complete->subchannel_call_tracker != nullptr) {
       if (subchannel_call_tracker != nullptr) {
         *subchannel_call_tracker = std::move(complete->subchannel_call_tracker);
       } else {
-        complete->subchannel_call_tracker->Start();
-        FakeMetadata metadata({});
-        FakeBackendMetricAccessor backend_metric_accessor({});
-        LoadBalancingPolicy::SubchannelCallTrackerInterface::FinishArgs args = {
-            address, absl::OkStatus(), &metadata, &backend_metric_accessor};
-        complete->subchannel_call_tracker->Finish(args);
+        CompleteCall(std::move(complete->subchannel_call_tracker), address);
       }
     }
     return address;
   }
+  void CompleteCall(
+      std::unique_ptr<LoadBalancingPolicy::SubchannelCallTrackerInterface>
+          subchannel_call_tracker,
+      absl::string_view address, absl::Status status = absl::OkStatus()) {
+    subchannel_call_tracker->Start();
+    FakeMetadata metadata({});
+    FakeBackendMetricAccessor backend_metric_accessor({});
+    LoadBalancingPolicy::SubchannelCallTrackerInterface::FinishArgs args = {
+        address, status, &metadata, &backend_metric_accessor};
+    subchannel_call_tracker->Finish(args);
+  }
 
   // Gets num_picks complete picks from picker and returns the resulting
   // list of addresses, or nullopt if a non-complete pick was returned.
@@ -1137,7 +1162,7 @@ class LoadBalancingPolicyTest : public ::testing::Test {
                                         subchannel_call_trackers == nullptr
                                             ? nullptr
                                             : &subchannel_call_tracker,
-                                        location);
+                                        nullptr, location);
       if (!address.has_value()) return absl::nullopt;
       results.emplace_back(std::move(*address));
       if (subchannel_call_trackers != nullptr) {
diff --git a/test/core/client_channel/lb_policy/pick_first_test.cc b/test/core/client_channel/lb_policy/pick_first_test.cc
index b5cdf28801e..fccc8af16ee 100644
--- a/test/core/client_channel/lb_policy/pick_first_test.cc
+++ b/test/core/client_channel/lb_policy/pick_first_test.cc
@@ -76,20 +76,8 @@ class PickFirstTest : public LoadBalancingPolicyTest {
   void GetOrderAddressesArePicked(
       absl::Span<const absl::string_view> addresses,
       std::vector<absl::string_view>* out_address_order) {
-    ExecCtx exec_ctx;
     out_address_order->clear();
-    // Note: ExitIdle() will enqueue a bunch of connectivity state
-    // notifications on the WorkSerializer, and we want to wait until
-    // those are delivered to the LB policy.
-    absl::Notification notification;
-    work_serializer_->Run(
-        [&]() {
-          lb_policy()->ExitIdleLocked();
-          work_serializer_->Run([&]() { notification.Notify(); },
-                                DEBUG_LOCATION);
-        },
-        DEBUG_LOCATION);
-    notification.WaitForNotification();
+    ExitIdle();
     // Construct a map of subchannel to address.
     // We will remove entries as each subchannel starts to connect.
     std::map<SubchannelState*, absl::string_view> subchannels;