From c9d1b69d415e34fdca8d19c55e972355ebf81939 Mon Sep 17 00:00:00 2001
From: Janusz Sobczak <jsobczak@google.com>
Date: Tue, 25 Jul 2023 18:00:39 -0700
Subject: [PATCH] Set expected latch before sending messages

This closes a time window where the sent message was received before the latch was
configured.

PiperOrigin-RevId: 551050923
---
 .../implementation/offline_service_controller_test.cc       | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/connections/implementation/offline_service_controller_test.cc b/connections/implementation/offline_service_controller_test.cc
index 3cedc81c..43daafea 100644
--- a/connections/implementation/offline_service_controller_test.cc
+++ b/connections/implementation/offline_service_controller_test.cc
@@ -284,10 +284,10 @@ TEST_P(OfflineServiceControllerTest, CanSendBytePayload) {
   env_.Start();
   OfflineSimulationUser user_a(kDeviceA, GetParam());
   OfflineSimulationUser user_b(kDeviceB, GetParam());
+  user_b.ExpectPayload(payload_latch_);
   ASSERT_TRUE(SetupConnection(user_a, user_b));
   ByteArray message(std::string{kMessage});
   user_a.SendPayload(Payload(message));
-  user_b.ExpectPayload(payload_latch_);
   EXPECT_TRUE(payload_latch_.Await(kLongTimeout));
   EXPECT_EQ(user_b.GetPayload().AsBytes(), message);
   user_a.Stop();
@@ -299,6 +299,7 @@ TEST_P(OfflineServiceControllerTest, CanSendStreamPayload) {
   env_.Start();
   OfflineSimulationUser user_a(kDeviceA, GetParam());
   OfflineSimulationUser user_b(kDeviceB, GetParam());
+  user_b.ExpectPayload(payload_latch_);
   ASSERT_TRUE(SetupConnection(user_a, user_b));
   ByteArray message(std::string{kMessage});
   auto pipe = std::make_shared<Pipe>();
@@ -306,7 +307,6 @@ TEST_P(OfflineServiceControllerTest, CanSendStreamPayload) {
   user_a.SendPayload(Payload([pipe]() -> InputStream& {
     return pipe->GetInputStream();  // NOLINT
   }));
-  user_b.ExpectPayload(payload_latch_);
   tx.Write(message);
   EXPECT_TRUE(payload_latch_.Await(kLongTimeout));
   ASSERT_NE(user_b.GetPayload().AsStream(), nullptr);
@@ -326,6 +326,7 @@ TEST_P(OfflineServiceControllerTest, CanCancelStreamPayload) {
   env_.Start();
   OfflineSimulationUser user_a(kDeviceA, GetParam());
   OfflineSimulationUser user_b(kDeviceB, GetParam());
+  user_b.ExpectPayload(payload_latch_);
   ASSERT_TRUE(SetupConnection(user_a, user_b));
   ByteArray message(std::string{kMessage});
   auto pipe = std::make_shared<Pipe>();
@@ -333,7 +334,6 @@ TEST_P(OfflineServiceControllerTest, CanCancelStreamPayload) {
   user_a.SendPayload(Payload([pipe]() -> InputStream& {
     return pipe->GetInputStream();  // NOLINT
   }));
-  user_b.ExpectPayload(payload_latch_);
   tx.Write(message);
   EXPECT_TRUE(payload_latch_.Await(kLongTimeout));
   ASSERT_NE(user_b.GetPayload().AsStream(), nullptr);