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

4 files changed, 140 insertions, 1 deletions

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