Test cleanup: Adds a helper for waiting for FSTVoidErrorBlock callbacks.

* Add helper for waiting for FSTVoidErrorBlock callbacks.
* Remove errorEventHandler from FSTEventAccumulator.h too.
* Add synchronous enableNetwork / disableNetwork helpers.
* Workaround for batch writes test flakiness.

10 files changed, 47 insertions, 60 deletions

diff --git a/Firestore/Example/Tests/Integration/API/FIRQueryTests.m b/Firestore/Example/Tests/Integration/API/FIRQueryTests.m
index 14351a8..e92bbcf 100644
--- a/Firestore/Example/Tests/Integration/API/FIRQueryTests.m
+++ b/Firestore/Example/Tests/Integration/API/FIRQueryTests.m
@@ -227,12 +227,8 @@
   FIRQuerySnapshot *querySnap = [self.eventAccumulator awaitEventWithName:@"initial event"];
   XCTAssertEqualObjects(FIRQuerySnapshotGetData(querySnap), @[ @{ @"foo" : @1 } ]);
   XCTAssertEqual(querySnap.metadata.isFromCache, NO);
-  XCTestExpectation *networkDisabled = [self expectationWithDescription:@"disable network"];
-  [collection.firestore.client disableNetworkWithCompletion:^(NSError *error) {
-    [networkDisabled fulfill];
-  }];
-  [self awaitExpectations];
 
+  [self disableNetwork];
   querySnap = [self.eventAccumulator awaitEventWithName:@"offline event with isFromCache=YES"];
   XCTAssertEqual(querySnap.metadata.isFromCache, YES);
 
@@ -241,12 +237,7 @@
   // sufficient.
   [NSThread sleepForTimeInterval:0.2f];
 
-  XCTestExpectation *networkEnabled = [self expectationWithDescription:@"enable network"];
-  [collection.firestore.client enableNetworkWithCompletion:^(NSError *error) {
-    [networkEnabled fulfill];
-  }];
-  [self awaitExpectations];
-
+  [self enableNetwork];
   querySnap = [self.eventAccumulator awaitEventWithName:@"back online event with isFromCache=NO"];
   XCTAssertEqual(querySnap.metadata.isFromCache, NO);
 }
diff --git a/Firestore/Example/Tests/Integration/API/FIRServerTimestampTests.m b/Firestore/Example/Tests/Integration/API/FIRServerTimestampTests.m
index 2a21dc0..5cda053 100644
--- a/Firestore/Example/Tests/Integration/API/FIRServerTimestampTests.m
+++ b/Firestore/Example/Tests/Integration/API/FIRServerTimestampTests.m
@@ -167,18 +167,6 @@
   [self awaitExpectations];
 }
 
-/** Disables the network synchronously. */
-- (void)disableNetwork {
-  [_docRef.firestore.client disableNetworkWithCompletion:_accumulator.errorEventHandler];
-  [_accumulator awaitEventWithName:@"Disconnect event."];
-}
-
-/** Enables the network synchronously. */
-- (void)enableNetwork {
-  [_docRef.firestore.client enableNetworkWithCompletion:_accumulator.errorEventHandler];
-  [_accumulator awaitEventWithName:@"Reconnect event."];
-}
-
 #pragma mark - Test Cases
 
 - (void)testServerTimestampsWorkViaSet {
diff --git a/Firestore/Example/Tests/Integration/API/FIRWriteBatchTests.m b/Firestore/Example/Tests/Integration/API/FIRWriteBatchTests.m
index 8ac5491..5e7f6d7 100644
--- a/Firestore/Example/Tests/Integration/API/FIRWriteBatchTests.m
+++ b/Firestore/Example/Tests/Integration/API/FIRWriteBatchTests.m
@@ -47,6 +47,12 @@
   FIRWriteBatch *batch2 = [doc.firestore batch];
   [batch2 setData:@{@"cc" : @"dd"} forDocument:doc];
   [batch2 commit];
+
+  // TODO(b/70631617): There's currently a backend bug that prevents us from using a resume token
+  // right away (against hexa at least). So we sleep. :-( :-( Anything over ~10ms seems to be
+  // sufficient.
+  [NSThread sleepForTimeInterval:0.2f];
+
   FIRDocumentSnapshot *snapshot2 = [self readDocumentForRef:doc];
   XCTAssertTrue(snapshot2.exists);
   XCTAssertEqualObjects(snapshot2.data, @{@"cc" : @"dd"});
diff --git a/Firestore/Example/Tests/Util/FSTEventAccumulator.h b/Firestore/Example/Tests/Util/FSTEventAccumulator.h
index 424f4c8..baa501b 100644
--- a/Firestore/Example/Tests/Util/FSTEventAccumulator.h
+++ b/Firestore/Example/Tests/Util/FSTEventAccumulator.h
@@ -24,7 +24,6 @@
 NS_ASSUME_NONNULL_BEGIN
 
 typedef void (^FSTValueEventHandler)(id _Nullable, NSError *_Nullable error);
-typedef void (^FSTErrorEventHandler)(NSError *_Nullable error);
 
 @interface FSTEventAccumulator : NSObject
 
@@ -37,7 +36,6 @@ typedef void (^FSTErrorEventHandler)(NSError *_Nullable error);
 - (NSArray<id> *)awaitEvents:(NSUInteger)events name:(NSString *)name;
 
 @property(nonatomic, strong, readonly) FSTValueEventHandler valueEventHandler;
-@property(nonatomic, strong, readonly) FSTErrorEventHandler errorEventHandler;
 
 @end
 
diff --git a/Firestore/Example/Tests/Util/FSTEventAccumulator.m b/Firestore/Example/Tests/Util/FSTEventAccumulator.m
index 1b1e276..c4c1602 100644
--- a/Firestore/Example/Tests/Util/FSTEventAccumulator.m
+++ b/Firestore/Example/Tests/Util/FSTEventAccumulator.m
@@ -68,7 +68,6 @@ NS_ASSUME_NONNULL_BEGIN
   return events[0];
 }
 
-// Override the valueEventHandler property
 - (void (^)(id _Nullable, NSError *_Nullable))valueEventHandler {
   return ^void(id _Nullable value, NSError *_Nullable error) {
     // We can't store nil in the _events array, but these are still interesting to tests so store
@@ -82,16 +81,6 @@ NS_ASSUME_NONNULL_BEGIN
   };
 }
 
-// Override the errorEventHandler property
-- (void (^)(NSError *_Nullable))errorEventHandler {
-  return ^void(NSError *_Nullable error) {
-    @synchronized(self) {
-      [_events addObject:[NSNull null]];
-      [self checkFulfilled];
-    }
-  };
-}
-
 - (void)checkFulfilled {
   if (_events.count >= self.maxEvents) {
     [self.expectation fulfill];
diff --git a/Firestore/Example/Tests/Util/FSTIntegrationTestCase.h b/Firestore/Example/Tests/Util/FSTIntegrationTestCase.h
index 88f9346..ac54253 100644
--- a/Firestore/Example/Tests/Util/FSTIntegrationTestCase.h
+++ b/Firestore/Example/Tests/Util/FSTIntegrationTestCase.h
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#import <Firestore/Source/Core/FSTTypes.h>
 #import <Foundation/Foundation.h>
 #import <XCTest/XCTest.h>
 
@@ -82,6 +83,10 @@ extern "C" {
 
 - (void)deleteDocumentRef:(FIRDocumentReference *)ref;
 
+- (void)disableNetwork;
+
+- (void)enableNetwork;
+
 /**
  * "Blocks" the current thread/run loop until the block returns YES.
  * Should only be called on the main thread.
diff --git a/Firestore/Example/Tests/Util/FSTIntegrationTestCase.mm b/Firestore/Example/Tests/Util/FSTIntegrationTestCase.mm
index 3d30a77..839e4a5 100644
--- a/Firestore/Example/Tests/Util/FSTIntegrationTestCase.mm
+++ b/Firestore/Example/Tests/Util/FSTIntegrationTestCase.mm
@@ -18,6 +18,7 @@
 
 #import <FirebaseCore/FIRLogger.h>
 #import <FirebaseFirestore/FirebaseFirestore-umbrella.h>
+#import <Firestore/Source/Core/FSTFirestoreClient.h>
 #import <GRPCClient/GRPCCall+ChannelArg.h>
 #import <GRPCClient/GRPCCall+Tests.h>
 
@@ -158,11 +159,7 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (void)shutdownFirestore:(FIRFirestore *)firestore {
-  XCTestExpectation *shutdownCompletion = [self expectationWithDescription:@"shutdown"];
-  [firestore shutdownWithCompletion:^(NSError *_Nullable error) {
-    XCTAssertNil(error);
-    [shutdownCompletion fulfill];
-  }];
+  [firestore shutdownWithCompletion:[self completionForExpectationWithName:@"shutdown"]];
   [self awaitExpectations];
 }
 
@@ -261,31 +258,29 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (void)writeDocumentRef:(FIRDocumentReference *)ref data:(NSDictionary<NSString *, id> *)data {
-  XCTestExpectation *expectation = [self expectationWithDescription:@"setData"];
-  [ref setData:data
-      completion:^(NSError *_Nullable error) {
-        XCTAssertNil(error);
-        [expectation fulfill];
-      }];
+  [ref setData:data completion:[self completionForExpectationWithName:@"setData"]];
   [self awaitExpectations];
 }
 
 - (void)updateDocumentRef:(FIRDocumentReference *)ref data:(NSDictionary<id, id> *)data {
-  XCTestExpectation *expectation = [self expectationWithDescription:@"updateData"];
-  [ref updateData:data
-       completion:^(NSError *_Nullable error) {
-         XCTAssertNil(error);
-         [expectation fulfill];
-       }];
+  [ref updateData:data completion:[self completionForExpectationWithName:@"updateData"]];
   [self awaitExpectations];
 }
 
 - (void)deleteDocumentRef:(FIRDocumentReference *)ref {
-  XCTestExpectation *expectation = [self expectationWithDescription:@"deleteDocument"];
-  [ref deleteDocumentWithCompletion:^(NSError *_Nullable error) {
-    XCTAssertNil(error);
-    [expectation fulfill];
-  }];
+  [ref deleteDocumentWithCompletion:[self completionForExpectationWithName:@"deleteDocument"]];
+  [self awaitExpectations];
+}
+
+- (void)disableNetwork {
+  [self.db.client
+      disableNetworkWithCompletion:[self completionForExpectationWithName:@"Disable Network."]];
+  [self awaitExpectations];
+}
+
+- (void)enableNetwork {
+  [self.db.client
+      enableNetworkWithCompletion:[self completionForExpectationWithName:@"Enable Network."]];
   [self awaitExpectations];
 }
 
diff --git a/Firestore/Example/Tests/Util/XCTestCase+Await.h b/Firestore/Example/Tests/Util/XCTestCase+Await.h
index 9d575f9..7a8feb8 100644
--- a/Firestore/Example/Tests/Util/XCTestCase+Await.h
+++ b/Firestore/Example/Tests/Util/XCTestCase+Await.h
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#import <Firestore/Source/Core/FSTTypes.h>
 #import <XCTest/XCTest.h>
 
 @interface XCTestCase (Await)
@@ -29,4 +30,10 @@
  */
 - (double)defaultExpectationWaitSeconds;
 
+/**
+ * Returns a completion block that fulfills a newly-created expectation with the specified
+ * name.
+ */
+- (FSTVoidErrorBlock)completionForExpectationWithName:(NSString *)expectationName;
+
 @end
diff --git a/Firestore/Example/Tests/Util/XCTestCase+Await.m b/Firestore/Example/Tests/Util/XCTestCase+Await.m
index 15c67ca..7f4356c 100644
--- a/Firestore/Example/Tests/Util/XCTestCase+Await.m
+++ b/Firestore/Example/Tests/Util/XCTestCase+Await.m
@@ -35,4 +35,12 @@ static const double kExpectationWaitSeconds = 10.0;
   return kExpectationWaitSeconds;
 }
 
+- (FSTVoidErrorBlock)completionForExpectationWithName:(NSString *)expectationName {
+  XCTestExpectation *expectation = [self expectationWithDescription:expectationName];
+  return ^(NSError *error) {
+    XCTAssertNil(error);
+    [expectation fulfill];
+  };
+}
+
 @end
diff --git a/Firestore/Source/Remote/FSTRemoteStore.m b/Firestore/Source/Remote/FSTRemoteStore.m
index 12cffd4..b1d91ea 100644
--- a/Firestore/Source/Remote/FSTRemoteStore.m
+++ b/Firestore/Source/Remote/FSTRemoteStore.m
@@ -161,8 +161,8 @@ static const int kOnlineAttemptsBeforeFailure = 2;
 }
 
 /**
-  * Updates our OnlineState to the new state, updating local state and notifying the
-  * onlineStateHandler as appropriate.
+ * Updates our OnlineState to the new state, updating local state and notifying the
+ * onlineStateHandler as appropriate.
  */
 - (void)updateOnlineState:(FSTOnlineState)newState {
   if (newState == FSTOnlineStateHealthy) {