Closing the write and watch stream after 60s of idleness (#388)

8 files changed, 275 insertions, 51 deletions

diff --git a/Firestore/Example/Firestore.xcodeproj/project.pbxproj b/Firestore/Example/Firestore.xcodeproj/project.pbxproj
index 2de1066..0197deb 100644
--- a/Firestore/Example/Firestore.xcodeproj/project.pbxproj
+++ b/Firestore/Example/Firestore.xcodeproj/project.pbxproj
@@ -59,6 +59,8 @@
 		AFE6114F0D4DAECBA7B7C089 /* Pods_Firestore_IntegrationTests.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = B2FA635DF5D116A67A7441CD /* Pods_Firestore_IntegrationTests.framework */; };
 		C4E749275AD0FBDF9F4716A8 /* Pods_SwiftBuildTest.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 32AD40BF6B0E849B07FFD05E /* Pods_SwiftBuildTest.framework */; };
 		D5B2532E4676014F57A7EAB9 /* FSTStreamTests.m in Sources */ = {isa = PBXBuildFile; fileRef = D5B25C0D4AADFCA3ADB883E4 /* FSTStreamTests.m */; };
+		D5B25474286C9800CE42B8C2 /* FSTTestDispatchQueue.m in Sources */ = {isa = PBXBuildFile; fileRef = D5B25292CED31B81FDED0411 /* FSTTestDispatchQueue.m */; };
+		D5B259FDEE8094E8D710C5BF /* FSTTestDispatchQueue.m in Sources */ = {isa = PBXBuildFile; fileRef = D5B25292CED31B81FDED0411 /* FSTTestDispatchQueue.m */; };
 		DE03B2C91F2149D600A30B9C /* FSTTransactionTests.m in Sources */ = {isa = PBXBuildFile; fileRef = DE51B1C61F0D48AC0013853F /* FSTTransactionTests.m */; };
 		DE03B2D41F2149D600A30B9C /* XCTest.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 6003F5AF195388D20070C39A /* XCTest.framework */; };
 		DE03B2D51F2149D600A30B9C /* UIKit.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 6003F591195388D20070C39A /* UIKit.framework */; };
@@ -222,6 +224,8 @@
 		B2FA635DF5D116A67A7441CD /* Pods_Firestore_IntegrationTests.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = Pods_Firestore_IntegrationTests.framework; sourceTree = BUILT_PRODUCTS_DIR; };
 		CE00BABB5A3AAB44A4C209E2 /* Pods-Firestore_Tests.debug.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-Firestore_Tests.debug.xcconfig"; path = "Pods/Target Support Files/Pods-Firestore_Tests/Pods-Firestore_Tests.debug.xcconfig"; sourceTree = "<group>"; };
 		D3CC3DC5338DCAF43A211155 /* README.md */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = net.daringfireball.markdown; name = README.md; path = ../README.md; sourceTree = "<group>"; };
+		D5B25292CED31B81FDED0411 /* FSTTestDispatchQueue.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = FSTTestDispatchQueue.m; sourceTree = "<group>"; };
+		D5B259DAA9149B80D6245B57 /* FSTTestDispatchQueue.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = FSTTestDispatchQueue.h; sourceTree = "<group>"; };
 		D5B25C0D4AADFCA3ADB883E4 /* FSTStreamTests.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = FSTStreamTests.m; sourceTree = "<group>"; };
 		DB17FEDFB80770611A935A60 /* Pods-Firestore_IntegrationTests.release.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-Firestore_IntegrationTests.release.xcconfig"; path = "Pods/Target Support Files/Pods-Firestore_IntegrationTests/Pods-Firestore_IntegrationTests.release.xcconfig"; sourceTree = "<group>"; };
 		DE03B2E91F2149D600A30B9C /* Firestore_IntegrationTests.xctest */ = {isa = PBXFileReference; explicitFileType = wrapper.cfbundle; includeInIndex = 0; path = Firestore_IntegrationTests.xctest; sourceTree = BUILT_PRODUCTS_DIR; };
@@ -561,6 +565,8 @@
 				DE51B18A1F0D48AC0013853F /* FSTUtilTests.m */,
 				54E9282A1F339CAD00C1953E /* XCTestCase+Await.h */,
 				54E9282B1F339CAD00C1953E /* XCTestCase+Await.m */,
+				D5B259DAA9149B80D6245B57 /* FSTTestDispatchQueue.h */,
+				D5B25292CED31B81FDED0411 /* FSTTestDispatchQueue.m */,
 			);
 			path = Util;
 			sourceTree = "<group>";
@@ -1175,6 +1181,7 @@
 				DE51B1CD1F0D48CD0013853F /* FSTDatabaseInfoTests.m in Sources */,
 				DE51B1F21F0D49140013853F /* FSTPathTests.m in Sources */,
 				DE51B1DD1F0D490D0013853F /* FSTLocalStoreTests.m in Sources */,
+				D5B25474286C9800CE42B8C2 /* FSTTestDispatchQueue.m in Sources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
@@ -1199,6 +1206,7 @@
 				DE03B2C91F2149D600A30B9C /* FSTTransactionTests.m in Sources */,
 				54DA12B11F315F3800DD57A1 /* FIRValidationTests.m in Sources */,
 				D5B2532E4676014F57A7EAB9 /* FSTStreamTests.m in Sources */,
+				D5B259FDEE8094E8D710C5BF /* FSTTestDispatchQueue.m in Sources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
diff --git a/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m b/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m
index 68692cc..6a6e49a 100644
--- a/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m
+++ b/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m
@@ -913,4 +913,22 @@
   [self awaitExpectations];
 }
 
+- (void)testWriteStreamReconnectsAfterIdle {
+  FIRDocumentReference *doc = [self documentRef];
+  FIRFirestore *firestore = doc.firestore;
+
+  [self writeDocumentRef:doc data:@{@"foo" : @"bar"}];
+  [self waitForIdleFirestore:firestore];
+  [self writeDocumentRef:doc data:@{@"foo" : @"bar"}];
+}
+
+- (void)testWatchStreamReconnectsAfterIdle {
+  FIRDocumentReference *doc = [self documentRef];
+  FIRFirestore *firestore = doc.firestore;
+
+  [self readSnapshotForRef:[self documentRef] requireOnline:YES];
+  [self waitForIdleFirestore:firestore];
+  [self readSnapshotForRef:[self documentRef] requireOnline:YES];
+}
+
 @end
diff --git a/Firestore/Example/Tests/Integration/FSTStreamTests.m b/Firestore/Example/Tests/Integration/FSTStreamTests.m
index 7ca123d..dccaa70 100644
--- a/Firestore/Example/Tests/Integration/FSTStreamTests.m
+++ b/Firestore/Example/Tests/Integration/FSTStreamTests.m
@@ -22,10 +22,10 @@
 #import "Core/FSTDatabaseInfo.h"
 #import "FSTHelpers.h"
 #import "FSTIntegrationTestCase.h"
+#import "FSTTestDispatchQueue.h"
 #import "Model/FSTDatabaseID.h"
 #import "Remote/FSTDatastore.h"
 #import "Util/FSTAssert.h"
-#import "Util/FSTDispatchQueue.h"
 
 /** Exposes otherwise private methods for testing. */
 @interface FSTStream (Testing)
@@ -79,14 +79,14 @@
   _expectation = nil;
 }
 
-- (void)writeStreamDidClose:(NSError *_Nullable)error {
-  [_states addObject:@"writeStreamDidClose"];
+- (void)writeStreamWasInterruptedWithError:(nullable NSError *)error {
+  [_states addObject:@"writeStreamWasInterrupted"];
   [_expectation fulfill];
   _expectation = nil;
 }
 
-- (void)watchStreamDidClose:(NSError *_Nullable)error {
-  [_states addObject:@"watchStreamDidClose"];
+- (void)watchStreamWasInterruptedWithError:(nullable NSError *)error {
+  [_states addObject:@"watchStreamWasInterrupted"];
   [_expectation fulfill];
   _expectation = nil;
 }
@@ -126,10 +126,10 @@
 
 @implementation FSTStreamTests {
   dispatch_queue_t _testQueue;
+  FSTTestDispatchQueue *_workerDispatchQueue;
   FSTDatabaseInfo *_databaseInfo;
   FSTEmptyCredentialsProvider *_credentials;
   FSTStreamStatusDelegate *_delegate;
-  FSTDispatchQueue *_workerDispatchQueue;
 
   /** Single mutation to send to the write stream. */
   NSArray<FSTMutation *> *_mutations;
@@ -138,38 +138,37 @@
 - (void)setUp {
   [super setUp];
 
-  _mutations = @[ FSTTestSetMutation(@"foo/bar", @{}) ];
-
   FIRFirestoreSettings *settings = [FSTIntegrationTestCase settings];
   FSTDatabaseID *databaseID =
       [FSTDatabaseID databaseIDWithProject:[FSTIntegrationTestCase projectID]
                                   database:kDefaultDatabaseID];
 
+  _testQueue = dispatch_queue_create("FSTStreamTestWorkerQueue", DISPATCH_QUEUE_SERIAL);
+  _workerDispatchQueue = [[FSTTestDispatchQueue alloc] initWithQueue:_testQueue];
+
   _databaseInfo = [FSTDatabaseInfo databaseInfoWithDatabaseID:databaseID
                                                persistenceKey:@"test-key"
                                                          host:settings.host
                                                    sslEnabled:settings.sslEnabled];
-  _testQueue = dispatch_queue_create("FSTStreamTestWorkerQueue", DISPATCH_QUEUE_SERIAL);
-  _workerDispatchQueue = [FSTDispatchQueue queueWith:_testQueue];
   _credentials = [[FSTEmptyCredentialsProvider alloc] init];
+
+  _delegate = [[FSTStreamStatusDelegate alloc] initWithTestCase:self queue:_workerDispatchQueue];
+
+  _mutations = @[ FSTTestSetMutation(@"foo/bar", @{}) ];
 }
 
 - (FSTWriteStream *)setUpWriteStream {
   FSTDatastore *datastore = [[FSTDatastore alloc] initWithDatabaseInfo:_databaseInfo
                                                    workerDispatchQueue:_workerDispatchQueue
                                                            credentials:_credentials];
-
-  _delegate = [[FSTStreamStatusDelegate alloc] initWithTestCase:self queue:_workerDispatchQueue];
-  return [datastore createWriteStreamWithDelegate:_delegate];
+  return [datastore createWriteStream];
 }
 
 - (FSTWatchStream *)setUpWatchStream {
   FSTDatastore *datastore = [[FSTDatastore alloc] initWithDatabaseInfo:_databaseInfo
                                                    workerDispatchQueue:_workerDispatchQueue
                                                            credentials:_credentials];
-
-  _delegate = [[FSTStreamStatusDelegate alloc] initWithTestCase:self queue:_workerDispatchQueue];
-  return [datastore createWatchStreamWithDelegate:_delegate];
+  return [datastore createWatchStream];
 }
 
 /**
@@ -190,7 +189,7 @@
   FSTWatchStream *watchStream = [self setUpWatchStream];
 
   [_delegate awaitNotificationFromBlock:^{
-    [watchStream start];
+    [watchStream startWithDelegate:_delegate];
   }];
 
   // Stop must not call watchStreamDidClose because the full implementation of the delegate could
@@ -210,7 +209,7 @@
   FSTWriteStream *writeStream = [self setUpWriteStream];
 
   [_delegate awaitNotificationFromBlock:^{
-    [writeStream start];
+    [writeStream startWithDelegate:_delegate];
   }];
 
   // Don't start the handshake.
@@ -231,7 +230,7 @@
   FSTWriteStream *writeStream = [self setUpWriteStream];
 
   [_delegate awaitNotificationFromBlock:^{
-    [writeStream start];
+    [writeStream startWithDelegate:_delegate];
   }];
 
   // Writing before the handshake should throw
@@ -258,4 +257,55 @@
   ]];
 }
 
+- (void)testStreamClosesWhenIdle {
+  FSTWriteStream *writeStream = [self setUpWriteStream];
+
+  [_delegate awaitNotificationFromBlock:^{
+    [writeStream startWithDelegate:_delegate];
+  }];
+
+  [_delegate awaitNotificationFromBlock:^{
+    [writeStream writeHandshake];
+  }];
+
+  [_delegate awaitNotificationFromBlock:^{
+    [writeStream markIdle];
+  }];
+
+  dispatch_sync(_testQueue, ^{
+    XCTAssertFalse([writeStream isOpen]);
+  });
+
+  [self verifyDelegateObservedStates:@[
+    @"writeStreamDidOpen", @"writeStreamDidCompleteHandshake", @"writeStreamWasInterrupted"
+  ]];
+}
+
+- (void)testStreamCancelsIdleOnWrite {
+  FSTWriteStream *writeStream = [self setUpWriteStream];
+
+  [_delegate awaitNotificationFromBlock:^{
+    [writeStream startWithDelegate:_delegate];
+  }];
+
+  [_delegate awaitNotificationFromBlock:^{
+    [writeStream writeHandshake];
+  }];
+
+  // Mark the stream idle, but immediately cancel the idle timer by issuing another write.
+  [_delegate awaitNotificationFromBlock:^{
+    [writeStream markIdle];
+    [writeStream writeMutations:_mutations];
+  }];
+
+  dispatch_sync(_testQueue, ^{
+    XCTAssertTrue([writeStream isOpen]);
+  });
+
+  [self verifyDelegateObservedStates:@[
+    @"writeStreamDidOpen", @"writeStreamDidCompleteHandshake",
+    @"writeStreamDidReceiveResponseWithVersion"
+  ]];
+}
+
 @end
diff --git a/Firestore/Example/Tests/SpecTests/FSTMockDatastore.m b/Firestore/Example/Tests/SpecTests/FSTMockDatastore.m
index 8b95286..6af2053 100644
--- a/Firestore/Example/Tests/SpecTests/FSTMockDatastore.m
+++ b/Firestore/Example/Tests/SpecTests/FSTMockDatastore.m
@@ -39,19 +39,17 @@ NS_ASSUME_NONNULL_BEGIN
 - (instancetype)initWithDatabase:(FSTDatabaseInfo *)database
              workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
                      credentials:(id<FSTCredentialsProvider>)credentials
-                      serializer:(FSTSerializerBeta *)serializer
-                        delegate:(id<FSTWatchStreamDelegate>)delegate NS_DESIGNATED_INITIALIZER;
+                      serializer:(FSTSerializerBeta *)serializer NS_DESIGNATED_INITIALIZER;
 
 - (instancetype)initWithDatabase:(FSTDatabaseInfo *)database
              workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
                      credentials:(id<FSTCredentialsProvider>)credentials
-            responseMessageClass:(Class)responseMessageClass
-                        delegate:(id<FSTWatchStreamDelegate>)delegate NS_UNAVAILABLE;
+            responseMessageClass:(Class)responseMessageClass NS_UNAVAILABLE;
 
 @property(nonatomic, assign) BOOL open;
-
 @property(nonatomic, strong, readonly)
     NSMutableDictionary<FSTBoxedTargetID *, FSTQueryData *> *activeTargets;
+@property(nonatomic, weak, readwrite, nullable) id<FSTWatchStreamDelegate> delegate;
 
 @end
 
@@ -60,13 +58,11 @@ NS_ASSUME_NONNULL_BEGIN
 - (instancetype)initWithDatabase:(FSTDatabaseInfo *)database
              workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
                      credentials:(id<FSTCredentialsProvider>)credentials
-                      serializer:(FSTSerializerBeta *)serializer
-                        delegate:(id<FSTWatchStreamDelegate>)delegate {
+                      serializer:(FSTSerializerBeta *)serializer {
   self = [super initWithDatabase:database
              workerDispatchQueue:workerDispatchQueue
                      credentials:credentials
-                      serializer:serializer
-                        delegate:delegate];
+                      serializer:serializer];
   if (self) {
     FSTAssert(database, @"Database must not be nil");
     _activeTargets = [NSMutableDictionary dictionary];
@@ -76,10 +72,15 @@ NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark - Overridden FSTWatchStream methods.
 
-- (void)start {
+- (void)startWithDelegate:(id<FSTWatchStreamDelegate>)delegate {
   FSTAssert(!self.open, @"Trying to start already started watch stream");
   self.open = YES;
-  [self handleStreamOpen];
+  self.delegate = delegate;
+  [self notifyStreamOpen];
+}
+
+- (void)stop {
+  self.delegate = nil;
 }
 
 - (BOOL)isOpen {
@@ -90,10 +91,14 @@ NS_ASSUME_NONNULL_BEGIN
   return self.open;
 }
 
-- (void)handleStreamOpen {
+- (void)notifyStreamOpen {
   [self.delegate watchStreamDidOpen];
 }
 
+- (void)notifyStreamInterruptedWithError:(nullable NSError *)error {
+  [self.delegate watchStreamWasInterruptedWithError:error];
+}
+
 - (void)watchQuery:(FSTQueryData *)query {
   FSTLog(@"watchQuery: %d: %@", query.targetID, query.query);
   // Snapshot version is ignored on the wire
@@ -110,7 +115,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)failStreamWithError:(NSError *)error {
   self.open = NO;
-  [self.delegate watchStreamDidClose:error];
+  [self notifyStreamInterruptedWithError:error];
 }
 
 #pragma mark - Helper methods.
@@ -142,17 +147,16 @@ NS_ASSUME_NONNULL_BEGIN
 - (instancetype)initWithDatabase:(FSTDatabaseInfo *)database
              workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
                      credentials:(id<FSTCredentialsProvider>)credentials
-                      serializer:(FSTSerializerBeta *)serializer
-                        delegate:(id<FSTWriteStreamDelegate>)delegate NS_DESIGNATED_INITIALIZER;
+                      serializer:(FSTSerializerBeta *)serializer NS_DESIGNATED_INITIALIZER;
 
 - (instancetype)initWithDatabase:(FSTDatabaseInfo *)database
              workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
                      credentials:(id<FSTCredentialsProvider>)credentials
-            responseMessageClass:(Class)responseMessageClass
-                        delegate:(id<FSTWriteStreamDelegate>)delegate NS_UNAVAILABLE;
+            responseMessageClass:(Class)responseMessageClass NS_UNAVAILABLE;
 
 @property(nonatomic, assign) BOOL open;
 @property(nonatomic, strong, readonly) NSMutableArray<NSArray<FSTMutation *> *> *sentMutations;
+@property(nonatomic, weak, readwrite, nullable) id<FSTWriteStreamDelegate> delegate;
 
 @end
 
@@ -161,13 +165,11 @@ NS_ASSUME_NONNULL_BEGIN
 - (instancetype)initWithDatabase:(FSTDatabaseInfo *)database
              workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
                      credentials:(id<FSTCredentialsProvider>)credentials
-                      serializer:(FSTSerializerBeta *)serializer
-                        delegate:(id<FSTWriteStreamDelegate>)delegate {
+                      serializer:(FSTSerializerBeta *)serializer {
   self = [super initWithDatabase:database
              workerDispatchQueue:workerDispatchQueue
                      credentials:credentials
-                      serializer:serializer
-                        delegate:delegate];
+                      serializer:serializer];
   if (self) {
     _sentMutations = [NSMutableArray array];
   }
@@ -176,11 +178,16 @@ NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark - Overridden FSTWriteStream methods.
 
-- (void)start {
+- (void)startWithDelegate:(id<FSTWriteStreamDelegate>)delegate {
   FSTAssert(!self.open, @"Trying to start already started write stream");
   self.open = YES;
   [self.sentMutations removeAllObjects];
-  [self handleStreamOpen];
+  self.delegate = delegate;
+  [self notifyStreamOpen];
+}
+
+- (void)stop {
+  self.delegate = nil;
 }
 
 - (BOOL)isOpen {
@@ -200,10 +207,14 @@ NS_ASSUME_NONNULL_BEGIN
   [self.sentMutations addObject:mutations];
 }
 
-- (void)handleStreamOpen {
+- (void)notifyStreamOpen {
   [self.delegate writeStreamDidOpen];
 }
 
+- (void)notifyStreamInterruptedWithError:(nullable NSError *)error {
+  [self.delegate writeStreamWasInterruptedWithError:error];
+}
+
 #pragma mark - Helper methods.
 
 /** Injects a write ack as though it had come from the backend in response to a write. */
@@ -215,7 +226,7 @@ NS_ASSUME_NONNULL_BEGIN
 /** Injects a failed write response as though it had come from the backend. */
 - (void)failStreamWithError:(NSError *)error {
   self.open = NO;
-  [self.delegate writeStreamDidClose:error];
+  [self notifyStreamInterruptedWithError:error];
 }
 
 /**
@@ -269,27 +280,25 @@ NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark - Overridden FSTDatastore methods.
 
-- (FSTWatchStream *)createWatchStreamWithDelegate:(id<FSTWatchStreamDelegate>)delegate {
+- (FSTWatchStream *)createWatchStream {
   FSTAssert(self.databaseInfo, @"DatabaseInfo must not be nil");
   self.watchStream = [[FSTMockWatchStream alloc]
          initWithDatabase:self.databaseInfo
       workerDispatchQueue:self.workerDispatchQueue
               credentials:self.credentials
                serializer:[[FSTSerializerBeta alloc]
-                              initWithDatabaseID:self.databaseInfo.databaseID]
-                 delegate:delegate];
+                              initWithDatabaseID:self.databaseInfo.databaseID]];
   return self.watchStream;
 }
 
-- (FSTWriteStream *)createWriteStreamWithDelegate:(id<FSTWriteStreamDelegate>)delegate {
+- (FSTWriteStream *)createWriteStream {
   FSTAssert(self.databaseInfo, @"DatabaseInfo must not be nil");
   self.writeStream = [[FSTMockWriteStream alloc]
          initWithDatabase:self.databaseInfo
       workerDispatchQueue:self.workerDispatchQueue
               credentials:self.credentials
                serializer:[[FSTSerializerBeta alloc]
-                              initWithDatabaseID:self.databaseInfo.databaseID]
-                 delegate:delegate];
+                              initWithDatabaseID:self.databaseInfo.databaseID]];
   return self.writeStream;
 }
 
diff --git a/Firestore/Example/Tests/Util/FSTIntegrationTestCase.h b/Firestore/Example/Tests/Util/FSTIntegrationTestCase.h
index 3dd5464..a2c08ec 100644
--- a/Firestore/Example/Tests/Util/FSTIntegrationTestCase.h
+++ b/Firestore/Example/Tests/Util/FSTIntegrationTestCase.h
@@ -56,6 +56,8 @@ NS_ASSUME_NONNULL_BEGIN
 - (FIRCollectionReference *)collectionRefWithDocuments:
     (NSDictionary<NSString *, NSDictionary<NSString *, id> *> *)documents;
 
+- (void)waitForIdleFirestore:(FIRFirestore *)firestore;
+
 - (void)writeAllDocuments:(NSDictionary<NSString *, NSDictionary<NSString *, id> *> *)documents
              toCollection:(FIRCollectionReference *)collection;
 
@@ -67,6 +69,9 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (FIRQuerySnapshot *)readDocumentSetForRef:(FIRQuery *)query;
 
+- (FIRDocumentSnapshot *)readSnapshotForRef:(FIRDocumentReference *)query
+                              requireOnline:(BOOL)online;
+
 - (void)writeDocumentRef:(FIRDocumentReference *)ref data:(NSDictionary<NSString *, id> *)data;
 
 - (void)updateDocumentRef:(FIRDocumentReference *)ref data:(NSDictionary<NSString *, id> *)data;
diff --git a/Firestore/Example/Tests/Util/FSTIntegrationTestCase.m b/Firestore/Example/Tests/Util/FSTIntegrationTestCase.m
index 87a78c3..2e1e0a9 100644
--- a/Firestore/Example/Tests/Util/FSTIntegrationTestCase.m
+++ b/Firestore/Example/Tests/Util/FSTIntegrationTestCase.m
@@ -30,9 +30,14 @@
 #import "Util/FSTUtil.h"
 
 #import "FSTEventAccumulator.h"
+#import "FSTTestDispatchQueue.h"
 
 NS_ASSUME_NONNULL_BEGIN
 
+@interface FIRFirestore (Testing)
+@property(nonatomic, strong) FSTDispatchQueue *workerDispatchQueue;
+@end
+
 @implementation FSTIntegrationTestCase {
   NSMutableArray<FIRFirestore *> *_firestores;
 }
@@ -121,7 +126,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (FIRFirestore *)firestoreWithProjectID:(NSString *)projectID {
   NSString *persistenceKey = [NSString stringWithFormat:@"db%lu", (unsigned long)_firestores.count];
 
-  FSTDispatchQueue *workerDispatchQueue = [FSTDispatchQueue
+  FSTTestDispatchQueue *workerDispatchQueue = [FSTTestDispatchQueue
       queueWith:dispatch_queue_create("com.google.firebase.firestore", DISPATCH_QUEUE_SERIAL)];
 
   FSTEmptyCredentialsProvider *credentialsProvider = [[FSTEmptyCredentialsProvider alloc] init];
@@ -142,6 +147,14 @@ NS_ASSUME_NONNULL_BEGIN
   return firestore;
 }
 
+- (void)waitForIdleFirestore:(FIRFirestore *)firestore {
+  XCTestExpectation *expectation = [self expectationWithDescription:@"idle"];
+  // Note that we wait on any task that is scheduled with a delay of 60s. Currently, the idle
+  // timeout is the only task that uses this delay.
+  [((FSTTestDispatchQueue *)firestore.workerDispatchQueue) fulfillOnExecution:expectation];
+  [self awaitExpectations];
+}
+
 - (void)shutdownFirestore:(FIRFirestore *)firestore {
   XCTestExpectation *shutdownCompletion = [self expectationWithDescription:@"shutdown"];
   [firestore shutdownWithCompletion:^(NSError *_Nullable error) {
@@ -222,6 +235,27 @@ NS_ASSUME_NONNULL_BEGIN
   return result;
 }
 
+- (FIRDocumentSnapshot *)readSnapshotForRef:(FIRDocumentReference *)ref
+                              requireOnline:(BOOL)requireOnline {
+  __block FIRDocumentSnapshot *result;
+
+  XCTestExpectation *expectation = [self expectationWithDescription:@"listener"];
+  id<FIRListenerRegistration> listener = [ref
+      addSnapshotListenerWithOptions:[[FIRDocumentListenOptions options] includeMetadataChanges:YES]
+                            listener:^(FIRDocumentSnapshot *snapshot, NSError *error) {
+                              XCTAssertNil(error);
+                              if (!requireOnline || !snapshot.metadata.fromCache) {
+                                result = snapshot;
+                                [expectation fulfill];
+                              }
+                            }];
+
+  [self awaitExpectations];
+  [listener remove];
+
+  return result;
+}
+
 - (void)writeDocumentRef:(FIRDocumentReference *)ref data:(NSDictionary<NSString *, id> *)data {
   XCTestExpectation *expectation = [self expectationWithDescription:@"setData"];
   [ref setData:data
diff --git a/Firestore/Example/Tests/Util/FSTTestDispatchQueue.h b/Firestore/Example/Tests/Util/FSTTestDispatchQueue.h
new file mode 100644
index 0000000..4f4e13e
--- /dev/null
+++ b/Firestore/Example/Tests/Util/FSTTestDispatchQueue.h
@@ -0,0 +1,39 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import "Util/FSTDispatchQueue.h"
+
+@class XCTestExpectation;
+
+NS_ASSUME_NONNULL_BEGIN
+
+/**
+ * Dispatch queue used in the integration tests that caps delayed executions at 1.0 seconds.
+ */
+@interface FSTTestDispatchQueue : FSTDispatchQueue
+
+/** Creates and returns an FSTTestDispatchQueue wrapping the specified dispatch_queue_t. */
++ (instancetype)queueWith:(dispatch_queue_t)dispatchQueue;
+
+/**
+ * Registers a test expectation that is fulfilled when the next delayed callback finished
+ * executing.
+ */
+- (void)fulfillOnExecution:(XCTestExpectation *)expectation;
+
+@end
+
+NS_ASSUME_NONNULL_END
diff --git a/Firestore/Example/Tests/Util/FSTTestDispatchQueue.m b/Firestore/Example/Tests/Util/FSTTestDispatchQueue.m
new file mode 100644
index 0000000..27b62bc
--- /dev/null
+++ b/Firestore/Example/Tests/Util/FSTTestDispatchQueue.m
@@ -0,0 +1,61 @@
+/*
+ * Copyright 2017 Google
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#import "FSTTestDispatchQueue.h"
+
+#import <XCTest/XCTestExpectation.h>
+
+#import "Util/FSTAssert.h"
+
+@interface FSTTestDispatchQueue ()
+
+@property(nonatomic, weak) XCTestExpectation* expectation;
+
+@end
+
+@implementation FSTTestDispatchQueue
+
+/** The delay used by the idle timeout */
+static const NSTimeInterval kIdleDispatchDelay = 60.0;
+
+/** The maximum delay we use in a test run. */
+static const NSTimeInterval kTestDispatchDelay = 1.0;
+
++ (instancetype)queueWith:(dispatch_queue_t)dispatchQueue {
+  return [[FSTTestDispatchQueue alloc] initWithQueue:dispatchQueue];
+}
+
+- (instancetype)initWithQueue:(dispatch_queue_t)dispatchQueue {
+  return (self = [super initWithQueue:dispatchQueue]);
+}
+
+- (void)dispatchAfterDelay:(NSTimeInterval)delay block:(void (^)(void))block {
+  [super dispatchAfterDelay:MIN(delay, kTestDispatchDelay)
+                      block:^() {
+                        block();
+                        if (delay == kIdleDispatchDelay) {
+                          [_expectation fulfill];
+                          _expectation = nil;
+                        }
+                      }];
+}
+
+- (void)fulfillOnExecution:(XCTestExpectation*)expectation {
+  FSTAssert(_expectation == nil, @"Previous expectation still active");
+  _expectation = expectation;
+}
+
+@end