Fix flaky bundle tests. (#547)

* Ensure progress updates are all received. * Fix flaky bundle tests.

Author: wu-hui

firestore/integration_test_internal/src/bundle_test.cc

@@ -65,6 +65,15 @@ std::string CreateTestBundle(Firestore* db) {
  return CreateBundle(db->app()->options().project_id());
 }
 
+void SetPromiseValueWhenUpdateIsFinal(
+ LoadBundleTaskProgress progress,
+ std::promise<void>& final_update_received) {
+ if (progress.state() == LoadBundleTaskProgress::State::kError ||
+ progress.state() == LoadBundleTaskProgress::State::kSuccess) {
+ final_update_received.set_value();
+ }
+}
+
 class BundleTest : public FirestoreIntegrationTest {
  protected:
  void SetUp() override {
@@ -133,15 +142,19 @@ TEST_F(BundleTest, CanLoadBundlesWithProgressUpdates) {
  auto bundle = CreateTestBundle(db);
 
  std::vector<LoadBundleTaskProgress> progresses;
+ std::promise<void> final_update;
  Future<LoadBundleTaskProgress> result = db->LoadBundle(
- bundle, [&progresses](const LoadBundleTaskProgress& progress) {
+ bundle,
+ [&progresses, &final_update](const LoadBundleTaskProgress& progress) {
  progresses.push_back(progress);
+ SetPromiseValueWhenUpdateIsFinal(progress, final_update);
  });
 
  auto final_progress = AwaitResult(result);
 
  // 4 progresses will be reported: initial, document 1, document 2, final
  // success.
+ final_update.get_future().wait();
  ASSERT_EQ(progresses.size(), 4);
  EXPECT_THAT(progresses[0], InProgressWithLoadedDocuments(0));
  EXPECT_THAT(progresses[1], InProgressWithLoadedDocuments(1));
@@ -195,12 +208,16 @@ TEST_F(BundleTest, LoadBundlesForASecondTimeSkips) {
  VerifySuccessProgress(first_load);
 
  std::vector<LoadBundleTaskProgress> progresses;
+ std::promise<void> final_update;
  LoadBundleTaskProgress second_load = AwaitResult(db->LoadBundle(
- bundle, [&progresses](const LoadBundleTaskProgress& progress) {
+ bundle,
+ [&progresses, &final_update](const LoadBundleTaskProgress& progress) {
  progresses.push_back(progress);
+ SetPromiseValueWhenUpdateIsFinal(progress, final_update);
  }));
 
  // There will be 4 progress updates if it does not skip loading.
+ final_update.get_future().wait();
  ASSERT_EQ(progresses.size(), 1);
  VerifySuccessProgress(progresses[0]);
  EXPECT_EQ(progresses[0], second_load);
@@ -220,13 +237,18 @@ TEST_F(BundleTest, LoadInvalidBundlesShouldFail) {
  };
  for (const auto& bundle : invalid_bundles) {
  std::vector<LoadBundleTaskProgress> progresses;
+ std::promise<void> final_update;
  Future<LoadBundleTaskProgress> result = db->LoadBundle(
- bundle, [&progresses](const LoadBundleTaskProgress& progress) {
+ bundle,
+ [&progresses, &final_update](const LoadBundleTaskProgress& progress) {
  progresses.push_back(progress);
+ SetPromiseValueWhenUpdateIsFinal(progress, final_update);
  });
- Await(result);
 
+ Await(result);
  EXPECT_NE(result.error(), Error::kErrorOk);
+
+ final_update.get_future().wait();
  ASSERT_EQ(progresses.size(), 1);
  VerifyErrorProgress(progresses[0]);
  }
@@ -306,13 +328,17 @@ TEST_F(BundleTest, LoadDocumentsFromOtherProjectsShouldFail) {
  Firestore* db = TestFirestore();
  auto bundle = CreateBundle("other-project");
  std::vector<LoadBundleTaskProgress> progresses;
+ std::promise<void> final_update;
  Future<LoadBundleTaskProgress> result = db->LoadBundle(
- bundle, [&progresses](const LoadBundleTaskProgress& progress) {
+ bundle,
+ [&progresses, &final_update](const LoadBundleTaskProgress& progress) {
  progresses.push_back(progress);
+ SetPromiseValueWhenUpdateIsFinal(progress, final_update);
  });
  Await(result);
-
  EXPECT_NE(result.error(), Error::kErrorOk);
+
+ final_update.get_future().wait();
  ASSERT_EQ(progresses.size(), 2);
  EXPECT_THAT(progresses[0], InProgressWithLoadedDocuments(0));
  VerifyErrorProgress(progresses[1]);