Rework how tests deal with future timeout

firestore/integration_test_internal/src/android/firestore_integration_test_android.cc

@@ -87,17 +87,9 @@ Local<Throwable> FirestoreAndroidIntegrationTest::CreateException(
 }
 
 void FirestoreAndroidIntegrationTest::Await(Env& env, const Task& task) {
-  int cycles = kTimeOutMillis / kCheckIntervalMillis;
-  while (env.ok() && !task.IsComplete(env)) {
-    if (ProcessEvents(kCheckIntervalMillis)) {
-      std::cout << "WARNING: app receives an event requesting exit."
-                << std::endl;
-      break;
-    }
-    --cycles;
-  }
+  bool success = WaitUntil([&] { return env.ok() && task.IsComplete(env); });
   if (env.ok()) {
-    EXPECT_GT(cycles, 0) << "Waiting for Task timed out.";
+    EXPECT_TRUE(success) << "Waiting for Task timed out.";
   }
 }
 

firestore/integration_test_internal/src/android/promise_android_test.cc

@@ -136,24 +136,10 @@ class TestCompletionBase : public Promise<PublicType,
   // Waits for `CompleteWith()` to be invoked. Returns `true` if an invocation
   // occurred prior to timing out or `false` otherwise.
   bool AwaitCompletion() {
-    int cycles_remaining = kTimeOutMillis / kCheckIntervalMillis;
-    while (true) {
-      {
-        MutexLock lock(mutex_);
-        if (invocation_count_ > 0) {
-          return true;
-        }
-      }
-
-      if (ProcessEvents(kCheckIntervalMillis)) {
-        return false;
-      }
-
-      cycles_remaining--;
-      if (cycles_remaining == 0) {
-        return false;
-      }
-    }
+    return WaitUntil([&] {
+      MutexLock lock(mutex_);
+      return invocation_count_ > 0;
+    });
   }
 
   // Returns the number of times that `CompleteWith()` has been invoked.
@@ -243,7 +229,7 @@ TEST_F(PromiseTest, FutureVoidShouldSucceedWhenTaskSucceeds) {
 
   SetTaskResult(0);
 
-  EXPECT_GT(WaitFor(future), 0);
+  EXPECT_TRUE(WaitUntilFutureCompletes(future));
   EXPECT_EQ(future.status(), FutureStatus::kFutureStatusComplete);
   EXPECT_EQ(future.error(), 0);
   EXPECT_EQ(future.result(), nullptr);
@@ -257,7 +243,7 @@ TEST_F(PromiseTest, FutureNonVoidShouldSucceedWhenTaskSucceeds) {
 
   SetTaskResult(42);
 
-  EXPECT_GT(WaitFor(future), 0);
+  EXPECT_TRUE(WaitUntilFutureCompletes(future));
   EXPECT_EQ(future.status(), FutureStatus::kFutureStatusComplete);
   EXPECT_EQ(future.error(), 0);
   EXPECT_EQ(*future.result(), "42");
@@ -270,7 +256,7 @@ TEST_F(PromiseTest, FutureVoidShouldFailWhenTaskFails) {
 
   SetTaskException(Error::kErrorFailedPrecondition, "Simulated failure");
 
-  EXPECT_GT(WaitFor(future), 0);
+  EXPECT_TRUE(WaitUntilFutureCompletes(future));
   EXPECT_EQ(future.status(), FutureStatus::kFutureStatusComplete);
   EXPECT_EQ(future.error(), Error::kErrorFailedPrecondition);
   EXPECT_EQ(future.error_message(), std::string("Simulated failure"));
@@ -285,7 +271,7 @@ TEST_F(PromiseTest, FutureNonVoidShouldFailWhenTaskFails) {
 
   SetTaskException(Error::kErrorFailedPrecondition, "Simulated failure");
 
-  EXPECT_GT(WaitFor(future), 0);
+  EXPECT_TRUE(WaitUntilFutureCompletes(future));
   EXPECT_EQ(future.status(), FutureStatus::kFutureStatusComplete);
   EXPECT_EQ(future.error(), Error::kErrorFailedPrecondition);
   EXPECT_EQ(future.error_message(), std::string("Simulated failure"));
@@ -299,7 +285,7 @@ TEST_F(PromiseTest, FutureVoidShouldCancelWhenTaskCancels) {
 
   CancelTask();
 
-  EXPECT_GT(WaitFor(future), 0);
+  EXPECT_TRUE(WaitUntilFutureCompletes(future));
   EXPECT_EQ(future.status(), FutureStatus::kFutureStatusComplete);
   EXPECT_EQ(future.error(), Error::kErrorCancelled);
   EXPECT_EQ(future.error_message(), std::string("cancelled"));
@@ -314,7 +300,7 @@ TEST_F(PromiseTest, FutureNonVoidShouldCancelWhenTaskCancels) {
 
   CancelTask();
 
-  EXPECT_GT(WaitFor(future), 0);
+  EXPECT_TRUE(WaitUntilFutureCompletes(future));
   EXPECT_EQ(future.status(), FutureStatus::kFutureStatusComplete);
   EXPECT_EQ(future.error(), Error::kErrorCancelled);
   EXPECT_EQ(future.error_message(), std::string("cancelled"));

firestore/integration_test_internal/src/bundle_test.cc

@@ -205,7 +205,8 @@ TEST_F(BundleTest, CanDeleteFirestoreFromProgressUpdate) {
 
   // This future is not completed, and returns back a nullptr for result when it
   // times out.
-  EXPECT_EQ(Await(result), nullptr);
+  EXPECT_TRUE(WaitUntilFutureCompletes(result));
+  EXPECT_EQ(result.result(), nullptr);
 
   // 3 progresses will be reported: initial, document 1, document 2.
   // Final progress update is missing because Firestore is deleted before that.

firestore/integration_test_internal/src/firestore_integration_test.cc

@@ -111,18 +111,10 @@ std::string ToFirestoreErrorCodeName(int error_code) {
   }
 }
 
-int WaitFor(const FutureBase& future) {
-  // Instead of getting a clock, we count the cycles instead.
-  int cycles = kTimeOutMillis / kCheckIntervalMillis;
-  while (future.status() == FutureStatus::kFutureStatusPending && cycles > 0) {
-    if (ProcessEvents(kCheckIntervalMillis)) {
-      std::cout << "WARNING: app receives an event requesting exit."
-                << std::endl;
-      break;
-    }
-    --cycles;
-  }
-  return cycles;
+bool WaitUntilFutureCompletes(const FutureBase& future, int timeout_ms) {
+  return WaitUntil(
+      [&] { return future.status() != FutureStatus::kFutureStatusPending; },
+      timeout_ms);
 }
 
 FirestoreIntegrationTest::FirestoreIntegrationTest() {
@@ -299,18 +291,12 @@ FirestoreIntegrationTest::QuerySnapshotToMap(
   return result;
 }
 
-/* static */
-void FirestoreIntegrationTest::Await(const Future<void>& future) {
-  while (future.status() == FutureStatus::kFutureStatusPending) {
-    if (ProcessEvents(kCheckIntervalMillis)) {
-      std::cout << "WARNING: app received an event requesting exit."
-                << std::endl;
-      break;
-    }
-  }
+void FirestoreIntegrationTest::Await(const Future<void>& future,
+                                     int timeout_ms) {
+  EXPECT_TRUE(WaitUntilFutureCompletes(future, timeout_ms))
+      << "Future<void> timed out.";
 }
 
-/* static */
 bool FirestoreIntegrationTest::FailIfUnsuccessful(const char* operation,
                                                   const FutureBase& future) {
   if (future.status() != FutureStatus::kFutureStatusComplete) {
@@ -326,11 +312,13 @@ bool FirestoreIntegrationTest::FailIfUnsuccessful(const char* operation,
   }
 }
 
-/* static */
 std::string FirestoreIntegrationTest::DescribeFailedFuture(
     const FutureBase& future) {
-  return "Future failed: " + ToFirestoreErrorCodeName(future.error()) + " (" +
-         std::to_string(future.error()) + "): " + future.error_message();
+  std::string error_message =
+      future.error_message() ? future.error_message() : "[no additional info]";
+  return std::string("Future failed: ") +
+         ToFirestoreErrorCodeName(future.error()) + " (" +
+         std::to_string(future.error()) + "): " + error_message;
 }
 
 bool ProcessEvents(int msec) { return app_framework::ProcessEvents(msec); }

firestore/integration_test_internal/src/firestore_integration_test.h

@@ -17,6 +17,7 @@
 #ifndef FIREBASE_FIRESTORE_INTEGRATION_TEST_INTERNAL_SRC_FIRESTORE_INTEGRATION_TEST_H_
 #define FIREBASE_FIRESTORE_INTEGRATION_TEST_INTERNAL_SRC_FIRESTORE_INTEGRATION_TEST_H_
 
+#include <chrono>  // NOLINT(build/c++11)
 #include <cstdio>
 #include <iostream>
 #include <map>
@@ -59,11 +60,16 @@ bool ProcessEvents(int msec);
 // enum, but this function will gracefully handle the case where it is not.
 std::string ToFirestoreErrorCodeName(int error_code);
 
-// Waits for a Future to complete. If a timeout is reached then this method
-// returns as if successful; therefore, the caller should verify the status of
-// the given Future after this function returns. Returns the number of cycles
-// that were left before a timeout would have occurred.
-int WaitFor(const FutureBase& future);
+// Blocks until either the given predicate `pred` returns `true` or timeout
+// occurs. Returns `true` on success, `false` on timeout or if exit signal was
+// received.
+template <typename PredT>
+bool WaitUntil(const PredT& pred, int timeout_ms = kTimeOutMillis);
+
+// Blocks until either the future completes or timeout occurs. Returns `true`
+// on success, `false` on timeout or if exit signal was received.
+bool WaitUntilFutureCompletes(const FutureBase& future,
+                              int timeout_ms = kTimeOutMillis);
 
 template <typename T>
 class EventAccumulator;
@@ -295,37 +301,30 @@ class FirestoreIntegrationTest : public testing::Test {
 
   // TODO(zxu): add a helper function to block on signal.
 
-  // A helper function to block until the future completes.
+  // Blocks until the future is completed and returns the future result.
   template <typename T>
-  static const T* Await(const Future<T>& future) {
-    int cycles = WaitFor(future);
-    EXPECT_GT(cycles, 0) << "Waiting future timed out.";
-    if (future.status() == FutureStatus::kFutureStatusComplete) {
-      if (future.result() == nullptr) {
-        std::cout << "WARNING: " << DescribeFailedFuture(future) << std::endl;
-      }
-    } else {
-      std::cout << "WARNING: Future is not completed." << std::endl;
-    }
+  static const T* Await(const Future<T>& future,
+                        int timeout_ms = kTimeOutMillis) {
+    EXPECT_TRUE(WaitUntilFutureCompletes(future, timeout_ms))
+        << "Future<T> timed out.";
+
+    EXPECT_EQ(future.status(), FutureStatus::kFutureStatusComplete)
+        << DescribeFailedFuture(future);
+    EXPECT_NE(future.result(), nullptr) << DescribeFailedFuture(future);
+
     return future.result();
   }
 
-  static void Await(const Future<void>& future);
+  // Blocks until the future completes.
+  static void Await(const Future<void>& future,
+                    int timeout_ms = kTimeOutMillis);
 
-  // A helper function to block until there is at least n event.
+  // Blocks until there is at least `event_count` events.
   template <typename T>
-  static void Await(const TestEventListener<T>& listener, int n = 1) {
-    // Instead of getting a clock, we count the cycles instead.
-    int cycles = kTimeOutMillis / kCheckIntervalMillis;
-    while (listener.event_count() < n && cycles > 0) {
-      if (ProcessEvents(kCheckIntervalMillis)) {
-        std::cout << "WARNING: app receives an event requesting exit."
-                  << std::endl;
-        return;
-      }
-      --cycles;
-    }
-    EXPECT_GT(cycles, 0) << "Waiting listener timed out.";
+  static void Await(const TestEventListener<T>& listener, int event_count = 1) {
+    bool success =
+        WaitUntil([&] { return listener.event_count() >= event_count; });
+    EXPECT_TRUE(success) << "Waiting for a listener timed out.";
   }
 
   // Fails the current test if the given future did not complete or contained an
@@ -370,6 +369,24 @@ class FirestoreIntegrationTest : public testing::Test {
   mutable std::unordered_map<Firestore*, FirestoreInfo> firestores_;
 };
 
+template <typename PredT>
+bool WaitUntil(const PredT& pred, int timeout_ms) {
+  auto now = std::chrono::steady_clock::now();
+  auto timeout_time = now + std::chrono::milliseconds(timeout_ms);
+
+  while (!pred() && now < timeout_time) {
+    if (ProcessEvents(kCheckIntervalMillis)) {
+      std::cout << "WARNING: app received an event requesting exit."
+                << std::endl;
+      return false;
+    }
+
+    now = std::chrono::steady_clock::now();
+  }
+
+  return now < timeout_time;
+}
+
 }  // namespace firestore
 }  // namespace firebase
 

firestore/integration_test_internal/src/transaction_extra_test.cc

@@ -103,7 +103,11 @@ TEST_F(TransactionExtraTest, TestReadingADocTwiceWithDifferentVersions) {
     transaction.Set(doc, MapFieldValue{{"count", FieldValue::Double(16.0)}});
     return error;
   });
-  Await(future);
+
+  // Because the transaction retries a few times with exponential backoff, it
+  // might time out with the default timeout time (the default timeout value is
+  // 15 seconds and the transaction can take 11-16 seconds).
+  Await(future, kTimeOutMillis * 2);
   EXPECT_EQ(Error::kErrorAborted, future.error());
   EXPECT_STREQ("Document version changed between two reads.",
                future.error_message());

firestore/integration_test_internal/src/transaction_test.cc

@@ -53,12 +53,13 @@ class TransactionTest : public FirestoreIntegrationTest {
   void RunTransactionAndExpect(
       Error error,
       const char* message,
-      std::function<Error(Transaction&, std::string&)> update) {
+      std::function<Error(Transaction&, std::string&)> update,
+      int timeout_ms = kTimeOutMillis) {
     Future<void> future;
     // Re-try 5 times in case server is unavailable.
     for (int i = 0; i < 5; ++i) {
       future = TestFirestore()->RunTransaction(update);
-      Await(future);
+      Await(future, timeout_ms);
       if (future.error() == Error::kErrorUnavailable) {
         std::cout << "Could not reach backend. Retrying transaction test."
                   << std::endl;
@@ -71,10 +72,13 @@ class TransactionTest : public FirestoreIntegrationTest {
   }
 
   void RunTransactionAndExpect(
-      Error error, std::function<Error(Transaction&, std::string&)> update) {
+      Error error,
+      std::function<Error(Transaction&, std::string&)> update,
+      int timeout_ms = kTimeOutMillis) {
     switch (error) {
       case Error::kErrorOk:
-        RunTransactionAndExpect(Error::kErrorOk, "", std::move(update));
+        RunTransactionAndExpect(Error::kErrorOk, "", std::move(update),
+                                timeout_ms);
         break;
       case Error::kErrorAborted:
         RunTransactionAndExpect(
@@ -83,15 +87,15 @@ class TransactionTest : public FirestoreIntegrationTest {
 #else
             Error::kErrorAborted,
 #endif
-            "Transaction failed all retries.", std::move(update));
+            "Transaction failed all retries.", std::move(update), timeout_ms);
         break;
       case Error::kErrorFailedPrecondition:
         // Here specifies error message of the most common cause. There are
         // other causes for FailedPrecondition as well. Use the one with message
         // parameter if the expected error message is different.
         RunTransactionAndExpect(Error::kErrorFailedPrecondition,
                                 "Can't update a document that doesn't exist.",
-                                std::move(update));
+                                std::move(update), timeout_ms);
         break;
       default:
         FAIL() << "Unexpected error code: " << error;
@@ -708,6 +712,10 @@ TEST_F(TransactionTest, TestCancellationOnError) {
   int count = 0;
 
   SCOPED_TRACE("TestCancellationOnError");
+
+  // Because the transaction retries a few times with exponential backoff, it
+  // might time out with the default timeout time (the default timeout value is
+  // 15 seconds and the transaction can take up to 20 seconds).
   RunTransactionAndExpect(
       Error::kErrorDeadlineExceeded, "no",
       [doc, &count_locker, &count](Transaction& transaction,
@@ -719,7 +727,8 @@ TEST_F(TransactionTest, TestCancellationOnError) {
         transaction.Set(doc, MapFieldValue{{"foo", FieldValue::String("bar")}});
         error_message = "no";
         return Error::kErrorDeadlineExceeded;
-      });
+      },
+      kTimeOutMillis * 3);
 
   // TODO(varconst): uncomment. Currently, there is no way in C++ to distinguish
   // user error, so the transaction gets retried, and the counter goes up to 6.