start fixing LB policy unit tests

markdroth · Sep 8, 2023 · 3afe12f · 3afe12f
1 parent 6261217
commit 3afe12f
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 3 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
@@ -153,6 +153,8 @@ class LoadBalancingPolicyTest : public ::testing::Test {
 
         void OnConnectivityStateChange(grpc_connectivity_state new_state,
                                        const absl::Status& status) override {
+          gpr_log(GPR_INFO, "notifying watcher: state=%s status=%s",
+                  ConnectivityStateName(new_state), status.ToString().c_str());
           watcher()->OnConnectivityStateChange(new_state, status);
         }
 
@@ -330,16 +332,19 @@ class LoadBalancingPolicyTest : public ::testing::Test {
             << "bug in test: " << ConnectivityStateName(state)
             << " must have OK status: " << status;
       }
+      absl::Notification notification;
       work_serializer_->Run(
-          [this, state, status, validate_state_transition, location]()
-              ABSL_EXCLUSIVE_LOCKS_REQUIRED(*work_serializer_) {
+          [&]() ABSL_EXCLUSIVE_LOCKS_REQUIRED(*work_serializer_) {
                 if (validate_state_transition) {
                   AssertValidConnectivityStateTransition(state_tracker_.state(),
                                                          state, location);
                 }
                 state_tracker_.SetState(state, status, "set from test");
+                work_serializer_->Run([&]() { notification.Notify(); },
+                                      DEBUG_LOCATION);
               },
           DEBUG_LOCATION);
+      notification.WaitForNotification();
     }
 
     // Indicates if any of the associated SubchannelInterface objects
@@ -704,7 +709,8 @@ class LoadBalancingPolicyTest : public ::testing::Test {
     work_serializer_->Run(
         [&]() {
           status = lb_policy->UpdateLocked(std::move(update_args));
-          notification.Notify();
+          work_serializer_->Run([&]() { notification.Notify(); },
+                                DEBUG_LOCATION);
         },
         DEBUG_LOCATION);
     notification.WaitForNotification();
@@ -1097,6 +1103,14 @@ class LoadBalancingPolicyTest : public ::testing::Test {
     return &it->second;
   }
 
+  void WaitForWorkSerializerToFlush() {
+    gpr_log(GPR_INFO, "waiting for WorkSerializer to flush...");
+    absl::Notification notification;
+    work_serializer_->Run([&]() { notification.Notify(); }, DEBUG_LOCATION);
+    notification.WaitForNotification();
+    gpr_log(GPR_INFO, "WorkSerializer flush complete");
+  }
+
   std::shared_ptr<WorkSerializer> work_serializer_;
   std::shared_ptr<grpc_event_engine::experimental::EventEngine> event_engine_ =
       grpc_event_engine::experimental::GetDefaultEventEngine();

diff --git a/test/core/client_channel/lb_policy/pick_first_test.cc b/test/core/client_channel/lb_policy/pick_first_test.cc
@@ -67,6 +67,10 @@ class PickFirstTest : public LoadBalancingPolicyTest {
       std::vector<absl::string_view>* out_address_order) {
     work_serializer_->Run([&]() { lb_policy_->ExitIdleLocked(); },
                           DEBUG_LOCATION);
+    // First flush is for ExitIdle(), second flush is for the resulting
+    // subchannel connectivity state notifications.
+    WaitForWorkSerializerToFlush();
+    WaitForWorkSerializerToFlush();
     out_address_order->clear();
     // Construct a map of subchannel to address.
     // We will remove entries as each subchannel starts to connect.
@@ -445,6 +449,13 @@ TEST_F(PickFirstTest, GoesIdleWhenConnectionFailsThenCanReconnect) {
   // By checking the picker, we told the LB policy to trigger a new
   // connection attempt, so it should start over with the first
   // subchannel.
+  // Note that the picker will have enqueued the ExitIdle() call in the
+  // WorkSerializer, so the first flush will execute that call.  But
+  // executing that call will result in enqueueing subchannel
+  // connectivity state notifications, so we need to flush again to make
+  // sure all of that work is done before we continue.
+  WaitForWorkSerializerToFlush();
+  WaitForWorkSerializerToFlush();
   EXPECT_TRUE(subchannel->ConnectionRequested());
   // The subchannel starts connecting.
   subchannel->SetConnectivityState(GRPC_CHANNEL_CONNECTING);