Sending an empty write request before tearing down the stream (#438)

7 files changed, 315 insertions, 28 deletions

diff --git a/Firestore/Example/Tests/SpecTests/FSTMockDatastore.h b/Firestore/Example/Tests/SpecTests/FSTMockDatastore.h
index 4ff1220..7ad0d18 100644
--- a/Firestore/Example/Tests/SpecTests/FSTMockDatastore.h
+++ b/Firestore/Example/Tests/SpecTests/FSTMockDatastore.h
@@ -22,6 +22,12 @@ NS_ASSUME_NONNULL_BEGIN
 
 @interface FSTMockDatastore : FSTDatastore
 
+/** A count of the total number of requests sent to the watch stream since the beginning of the test case. */
+@property(nonatomic) int watchStreamRequestCount;
+
+/** A count of the total number of requests sent to the write stream since the beginning of the test case. */
+@property(nonatomic) int writeStreamRequestCount;
+
 + (instancetype)mockDatastoreWithWorkerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue;
 
 #pragma mark - Watch Stream manipulation.
diff --git a/Firestore/Example/Tests/SpecTests/FSTMockDatastore.m b/Firestore/Example/Tests/SpecTests/FSTMockDatastore.m
index 10bfc5b..d3b196b 100644
--- a/Firestore/Example/Tests/SpecTests/FSTMockDatastore.m
+++ b/Firestore/Example/Tests/SpecTests/FSTMockDatastore.m
@@ -37,10 +37,15 @@ NS_ASSUME_NONNULL_BEGIN
 
 @interface FSTMockWatchStream : FSTWatchStream
 
+- (instancetype)initWithDatastore:(FSTMockDatastore *)datastore
+              workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
+                      credentials:(id<FSTCredentialsProvider>)credentials
+                       serializer:(FSTSerializerBeta *)serializer NS_DESIGNATED_INITIALIZER;
+
 - (instancetype)initWithDatabase:(FSTDatabaseInfo *)database
              workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
                      credentials:(id<FSTCredentialsProvider>)credentials
-                      serializer:(FSTSerializerBeta *)serializer NS_DESIGNATED_INITIALIZER;
+                      serializer:(FSTSerializerBeta *)serializer NS_UNAVAILABLE;
 
 - (instancetype)initWithDatabase:(FSTDatabaseInfo *)database
              workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
@@ -48,6 +53,7 @@ NS_ASSUME_NONNULL_BEGIN
             responseMessageClass:(Class)responseMessageClass NS_UNAVAILABLE;
 
 @property(nonatomic, assign) BOOL open;
+@property(nonatomic, strong, readonly) FSTMockDatastore *datastore;
 @property(nonatomic, strong, readonly)
     NSMutableDictionary<FSTBoxedTargetID *, FSTQueryData *> *activeTargets;
 @property(nonatomic, weak, readwrite, nullable) id<FSTWatchStreamDelegate> delegate;
@@ -56,16 +62,17 @@ NS_ASSUME_NONNULL_BEGIN
 
 @implementation FSTMockWatchStream
 
-- (instancetype)initWithDatabase:(FSTDatabaseInfo *)database
-             workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
-                     credentials:(id<FSTCredentialsProvider>)credentials
-                      serializer:(FSTSerializerBeta *)serializer {
-  self = [super initWithDatabase:database
+- (instancetype)initWithDatastore:(FSTMockDatastore *)datastore
+              workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
+                      credentials:(id<FSTCredentialsProvider>)credentials
+                       serializer:(FSTSerializerBeta *)serializer {
+  self = [super initWithDatabase:datastore.databaseInfo
              workerDispatchQueue:workerDispatchQueue
                      credentials:credentials
                       serializer:serializer];
   if (self) {
-    FSTAssert(database, @"Database must not be nil");
+    FSTAssert(datastore, @"Datastore must not be nil");
+    _datastore = datastore;
     _activeTargets = [NSMutableDictionary dictionary];
   }
   return self;
@@ -81,6 +88,7 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (void)stop {
+  [self.activeTargets removeAllObjects];
   self.delegate = nil;
 }
 
@@ -102,6 +110,7 @@ NS_ASSUME_NONNULL_BEGIN
 
 - (void)watchQuery:(FSTQueryData *)query {
   FSTLog(@"watchQuery: %d: %@", query.targetID, query.query);
+  self.datastore.watchStreamRequestCount += 1;
   // Snapshot version is ignored on the wire
   FSTQueryData *sentQueryData =
       [query queryDataByReplacingSnapshotVersion:[FSTSnapshotVersion noVersion]
@@ -143,35 +152,51 @@ NS_ASSUME_NONNULL_BEGIN
 
 #pragma mark - FSTMockWriteStream
 
+@interface FSTWriteStream ()
+
+@property(nonatomic, weak, readwrite, nullable) id<FSTWriteStreamDelegate> delegate;
+
+- (void)notifyStreamOpen;
+- (void)notifyStreamInterruptedWithError:(nullable NSError *)error;
+
+@end
+
 @interface FSTMockWriteStream : FSTWriteStream
 
+- (instancetype)initWithDatastore:(FSTMockDatastore *)datastore
+              workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
+                      credentials:(id<FSTCredentialsProvider>)credentials
+                       serializer:(FSTSerializerBeta *)serializer NS_DESIGNATED_INITIALIZER;
+
 - (instancetype)initWithDatabase:(FSTDatabaseInfo *)database
              workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
                      credentials:(id<FSTCredentialsProvider>)credentials
-                      serializer:(FSTSerializerBeta *)serializer NS_DESIGNATED_INITIALIZER;
+                      serializer:(FSTSerializerBeta *)serializer NS_UNAVAILABLE;
 
 - (instancetype)initWithDatabase:(FSTDatabaseInfo *)database
              workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
                      credentials:(id<FSTCredentialsProvider>)credentials
             responseMessageClass:(Class)responseMessageClass NS_UNAVAILABLE;
 
+@property(nonatomic, strong, readonly) FSTMockDatastore *datastore;
 @property(nonatomic, assign) BOOL open;
 @property(nonatomic, strong, readonly) NSMutableArray<NSArray<FSTMutation *> *> *sentMutations;
-@property(nonatomic, weak, readwrite, nullable) id<FSTWriteStreamDelegate> delegate;
 
 @end
 
 @implementation FSTMockWriteStream
 
-- (instancetype)initWithDatabase:(FSTDatabaseInfo *)database
-             workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
-                     credentials:(id<FSTCredentialsProvider>)credentials
-                      serializer:(FSTSerializerBeta *)serializer {
-  self = [super initWithDatabase:database
+- (instancetype)initWithDatastore:(FSTMockDatastore *)datastore
+              workerDispatchQueue:(FSTDispatchQueue *)workerDispatchQueue
+                      credentials:(id<FSTCredentialsProvider>)credentials
+                       serializer:(FSTSerializerBeta *)serializer {
+  self = [super initWithDatabase:datastore.databaseInfo
              workerDispatchQueue:workerDispatchQueue
                      credentials:credentials
                       serializer:serializer];
   if (self) {
+    FSTAssert(datastore, @"Datastore must not be nil");
+    _datastore = datastore;
     _sentMutations = [NSMutableArray array];
   }
   return self;
@@ -187,10 +212,6 @@ NS_ASSUME_NONNULL_BEGIN
   [self notifyStreamOpen];
 }
 
-- (void)stop {
-  self.delegate = nil;
-}
-
 - (BOOL)isOpen {
   return self.open;
 }
@@ -200,22 +221,16 @@ NS_ASSUME_NONNULL_BEGIN
 }
 
 - (void)writeHandshake {
+  self.datastore.writeStreamRequestCount += 1;
   self.handshakeComplete = YES;
   [self.delegate writeStreamDidCompleteHandshake];
 }
 
 - (void)writeMutations:(NSArray<FSTMutation *> *)mutations {
+  self.datastore.writeStreamRequestCount += 1;
   [self.sentMutations addObject:mutations];
 }
 
-- (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. */
@@ -284,7 +299,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (FSTWatchStream *)createWatchStream {
   FSTAssert(self.databaseInfo, @"DatabaseInfo must not be nil");
   self.watchStream = [[FSTMockWatchStream alloc]
-         initWithDatabase:self.databaseInfo
+        initWithDatastore:self
       workerDispatchQueue:self.workerDispatchQueue
               credentials:self.credentials
                serializer:[[FSTSerializerBeta alloc]
@@ -295,7 +310,7 @@ NS_ASSUME_NONNULL_BEGIN
 - (FSTWriteStream *)createWriteStream {
   FSTAssert(self.databaseInfo, @"DatabaseInfo must not be nil");
   self.writeStream = [[FSTMockWriteStream alloc]
-         initWithDatabase:self.databaseInfo
+        initWithDatastore:self
       workerDispatchQueue:self.workerDispatchQueue
               credentials:self.credentials
                serializer:[[FSTSerializerBeta alloc]
diff --git a/Firestore/Example/Tests/SpecTests/FSTSpecTests.m b/Firestore/Example/Tests/SpecTests/FSTSpecTests.m
index f681347..08ff6c7 100644
--- a/Firestore/Example/Tests/SpecTests/FSTSpecTests.m
+++ b/Firestore/Example/Tests/SpecTests/FSTSpecTests.m
@@ -318,6 +318,14 @@ static NSString *const kNoIOSTag = @"no-ios";
   }
 }
 
+- (void)doDisableNetwork {
+  [self.driver disableNetwork];
+}
+
+- (void)doEnableNetwork {
+  [self.driver enableNetwork];
+}
+
 - (void)doChangeUser:(id)UID {
   FSTUser *user = [UID isEqual:[NSNull null]] ? [FSTUser unauthenticatedUser]
                                               : [[FSTUser alloc] initWithUID:UID];
@@ -376,6 +384,12 @@ static NSString *const kNoIOSTag = @"no-ios";
     [self doWriteAck:step[@"writeAck"]];
   } else if (step[@"failWrite"]) {
     [self doFailWrite:step[@"failWrite"]];
+  } else if (step[@"enableNetwork"]) {
+    if ([step[@"enableNetwork"] boolValue]) {
+      [self doEnableNetwork];
+    } else {
+      [self doDisableNetwork];
+    }
   } else if (step[@"changeUser"]) {
     [self doChangeUser:step[@"changeUser"]];
   } else if (step[@"restart"]) {
@@ -458,6 +472,14 @@ static NSString *const kNoIOSTag = @"no-ios";
     if (expected[@"numOutstandingWrites"]) {
       XCTAssertEqual([self.driver sentWritesCount], [expected[@"numOutstandingWrites"] intValue]);
     }
+    if (expected[@"writeStreamRequestCount"]) {
+      XCTAssertEqual([self.driver writeStreamRequestCount],
+                     [expected[@"writeStreamRequestCount"] intValue]);
+    }
+    if (expected[@"watchStreamRequestCount"]) {
+      XCTAssertEqual([self.driver watchStreamRequestCount],
+                     [expected[@"watchStreamRequestCount"] intValue]);
+    }
     if (expected[@"limboDocs"]) {
       NSMutableSet<FSTDocumentKey *> *expectedLimboDocuments = [NSMutableSet set];
       NSArray *docNames = expected[@"limboDocs"];
diff --git a/Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.h b/Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.h
index 0643d76..e78ef92 100644
--- a/Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.h
+++ b/Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.h
@@ -195,6 +195,28 @@ typedef NSDictionary<FSTUser *, NSArray<FSTOutstandingWrite *> *> FSTOutstanding
 @property(nonatomic, readonly) int sentWritesCount;
 
 /**
+ * A count of the total number of requests sent to the write stream since the beginning of the test
+ * case.
+ */
+@property(nonatomic, readonly) int writeStreamRequestCount;
+
+/**
+ * A count of the total number of requests sent to the watch stream since the beginning of the test
+ * case.
+ */
+@property(nonatomic, readonly) int watchStreamRequestCount;
+
+/**
+ * Disables RemoteStore's network connection and shuts down all streams.
+ */
+- (void)disableNetwork;
+
+/**
+ * Enables RemoteStore's network connection.
+ */
+- (void)enableNetwork;
+
+/**
  * Switches the FSTSyncEngine to a new user. The test driver tracks the outstanding mutations for
  * each user, so future receiveWriteAck/Error operations will validate the write sent to the mock
  * datastore matches the next outstanding write for that user.
diff --git a/Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.m b/Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.m
index b83942b..a4678e3 100644
--- a/Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.m
+++ b/Firestore/Example/Tests/SpecTests/FSTSyncEngineTestDriver.m
@@ -171,6 +171,25 @@ NS_ASSUME_NONNULL_BEGIN
   return [self.datastore writesSent];
 }
 
+- (int)writeStreamRequestCount {
+  return [self.datastore writeStreamRequestCount];
+}
+
+- (int)watchStreamRequestCount {
+  return [self.datastore watchStreamRequestCount];
+}
+
+- (void)disableNetwork {
+  // Make sure to execute all writes that are currently queued. This allows us
+  // to assert on the total number of requests sent before shutdown.
+  [self.remoteStore fillWritePipeline];
+  [self.remoteStore disableNetwork];
+}
+
+- (void)enableNetwork {
+  [self.remoteStore enableNetwork];
+}
+
 - (void)changeUser:(FSTUser *)user {
   self.currentUser = user;
   [self.syncEngine userDidChange:user];
diff --git a/Firestore/Example/Tests/SpecTests/json/listen_spec_test.json b/Firestore/Example/Tests/SpecTests/json/listen_spec_test.json
index 35704f2..e607710 100644
--- a/Firestore/Example/Tests/SpecTests/json/listen_spec_test.json
+++ b/Firestore/Example/Tests/SpecTests/json/listen_spec_test.json
@@ -1520,5 +1520,163 @@
         ]
       }
     ]
+  },
+  "Listens are reestablished after network disconnect": {
+    "describeName": "Listens:",
+    "itName": "Listens are reestablished after network disconnect",
+    "tags": [],
+    "config": {
+      "useGarbageCollection": true
+    },
+    "steps": [
+      {
+        "userListen": [
+          2,
+          {
+            "path": "collection",
+            "filters": [],
+            "orderBys": []
+          }
+        ],
+        "stateExpect": {
+          "activeTargets": {
+            "2": {
+              "query": {
+                "path": "collection",
+                "filters": [],
+                "orderBys": []
+              },
+              "resumeToken": ""
+            }
+          },
+          "watchStreamRequestCount": 1
+        }
+      },
+      {
+        "watchAck": [
+          2
+        ]
+      },
+      {
+        "watchEntity": {
+          "docs": [
+            [
+              "collection/a",
+              1000,
+              {
+                "key": "a"
+              }
+            ]
+          ],
+          "targets": [
+            2
+          ]
+        }
+      },
+      {
+        "watchCurrent": [
+          [
+            2
+          ],
+          "resume-token-1000"
+        ],
+        "watchSnapshot": 1000,
+        "expect": [
+          {
+            "query": {
+              "path": "collection",
+              "filters": [],
+              "orderBys": []
+            },
+            "added": [
+              [
+                "collection/a",
+                1000,
+                {
+                  "key": "a"
+                }
+              ]
+            ],
+            "errorCode": 0,
+            "fromCache": false,
+            "hasPendingWrites": false
+          }
+        ]
+      },
+      {
+        "enableNetwork": false,
+        "stateExpect": {
+          "activeTargets": {},
+          "limboDocs": []
+        }
+      },
+      {
+        "enableNetwork": true,
+        "stateExpect": {
+          "activeTargets": {
+            "2": {
+              "query": {
+                "path": "collection",
+                "filters": [],
+                "orderBys": []
+              },
+              "resumeToken": "resume-token-1000"
+            }
+          },
+          "watchStreamRequestCount": 2
+        }
+      },
+      {
+        "watchAck": [
+          2
+        ]
+      },
+      {
+        "watchEntity": {
+          "docs": [
+            [
+              "collection/b",
+              2000,
+              {
+                "key": "b"
+              }
+            ]
+          ],
+          "targets": [
+            2
+          ]
+        }
+      },
+      {
+        "watchCurrent": [
+          [
+            2
+          ],
+          "resume-token-2000"
+        ],
+        "watchSnapshot": 2000,
+        "expect": [
+          {
+            "query": {
+              "path": "collection",
+              "filters": [],
+              "orderBys": []
+            },
+            "added": [
+              [
+                "collection/b",
+                2000,
+                {
+                  "key": "b"
+                }
+              ]
+            ],
+            "errorCode": 0,
+            "fromCache": false,
+            "hasPendingWrites": false
+          }
+        ]
+      }
+    ]
   }
 }
diff --git a/Firestore/Example/Tests/SpecTests/json/write_spec_test.json b/Firestore/Example/Tests/SpecTests/json/write_spec_test.json
index 51d1aee..8e3f5d5 100644
--- a/Firestore/Example/Tests/SpecTests/json/write_spec_test.json
+++ b/Firestore/Example/Tests/SpecTests/json/write_spec_test.json
@@ -5376,5 +5376,50 @@
         ]
       }
     ]
+  },
+  "Writes are resent after network disconnect": {
+    "describeName": "Writes:",
+    "itName": "Writes are resent after network disconnect",
+    "tags": [],
+    "config": {
+      "useGarbageCollection": true
+    },
+    "steps": [
+      {
+        "userSet": [
+          "collection/key",
+          {
+            "foo": "bar"
+          }
+        ],
+        "stateExpect": {
+          "numOutstandingWrites": 1
+        }
+      },
+      {
+        "enableNetwork": false,
+        "stateExpect": {
+          "activeTargets": {},
+          "limboDocs": [],
+          "writeStreamRequestCount": 3
+        }
+      },
+      {
+        "enableNetwork": true,
+        "stateExpect": {
+          "writeStreamRequestCount": 5,
+          "numOutstandingWrites": 1
+        }
+      },
+      {
+        "writeAck": {
+          "version": 1,
+          "expectUserCallback": false
+        },
+        "stateExpect": {
+          "numOutstandingWrites": 0
+        }
+      }
+    ]
   }
 }