[testing] Add ExitIdle and CompleteCall helpers to LB policy test lib (#35860)

As title. Pulling these additions out from a larger change.

Related: cl/563857636

Closes #35860

COPYBARA_INTEGRATE_REVIEW=https://github.com/grpc/grpc/pull/35860 from apolcyn:test_lib_changes 40b6455638
PiperOrigin-RevId: 606708025

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,22 +1118,31 @@ 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();
+        ReportCompletionToCallTracker(
-        FakeMetadata metadata({});
+            std::move(complete->subchannel_call_tracker), address);
-        FakeBackendMetricAccessor backend_metric_accessor({});
-        LoadBalancingPolicy::SubchannelCallTrackerInterface::FinishArgs args = {
-            address, absl::OkStatus(), &metadata, &backend_metric_accessor};
-        complete->subchannel_call_tracker->Finish(args);
       }
     }
     return address;
   }
 
+  void ReportCompletionToCallTracker(
+      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.
   absl::optional<std::vector<std::string>> GetCompletePicks(
@@ -1137,7 +1164,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) {

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
+    ExitIdle();
-    // 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();
     // Construct a map of subchannel to address.
     // We will remove entries as each subchannel starts to connect.
     std::map<SubchannelState*, absl::string_view> subchannels;

test/core/client_channel/lb_policy/xds_override_host_test.cc

@@ -187,10 +187,10 @@ class XdsOverrideHostTest : public LoadBalancingPolicyTest {
     }
     std::string expected_addresses_str = absl::StrJoin(expected_addresses, ",");
     for (size_t i = 0; i < 3; ++i) {
-      EXPECT_EQ(
+      EXPECT_EQ(ExpectPickComplete(picker, {attribute},
-          ExpectPickComplete(picker, {attribute},
+                                   /*subchannel_call_tracker=*/nullptr,
-                             /*subchannel_call_tracker=*/nullptr, location),
+                                   /*picked_subchannel=*/nullptr, location),
-          expected)
+                expected)
           << location.file() << ":" << location.line();
       EXPECT_EQ(attribute->actual_address_list(), expected_addresses_str)
           << "  Actual: " << attribute->actual_address_list() << "\n"
@@ -207,7 +207,8 @@ class XdsOverrideHostTest : public LoadBalancingPolicyTest {
     std::vector<std::string> actual_picks;
     for (size_t i = 0; i < expected.size(); ++i) {
       auto address = ExpectPickComplete(
-          picker, {attribute}, /*subchannel_call_tracker=*/nullptr, location);
+          picker, {attribute}, /*subchannel_call_tracker=*/nullptr,
+          /*picked_subchannel=*/nullptr, location);
       ASSERT_TRUE(address.has_value())
           << location.file() << ":" << location.line();
       EXPECT_THAT(*address, ::testing::AnyOfArray(expected))