Expose network management (#566)

* Expose network management in public API

* Clean up a few more references to the internal access of network management

* Move test

* Update comments

* Swap _Nullable for nullable

* Fix comment

* Add tests, including swift

* Styling

3 files changed, 67 insertions, 41 deletions

diff --git a/Firestore/Source/API/FIRFirestore.m b/Firestore/Source/API/FIRFirestore.m
index 7814ce1..e347911 100644
--- a/Firestore/Source/API/FIRFirestore.m
+++ b/Firestore/Source/API/FIRFirestore.m
@@ -279,6 +279,16 @@ NSString *const FIRFirestoreErrorDomain = @"FIRFirestoreErrorDomain";
   FIRSetLoggerLevel(logging ? FIRLoggerLevelDebug : FIRLoggerLevelNotice);
 }
 
+- (void)enableNetworkWithCompletion:(nullable void (^)(NSError *_Nullable error))completion {
+  [self firestoreWithConfiguredClient];
+  [self.client enableNetworkWithCompletion:completion];
+}
+
+- (void)disableNetworkWithCompletion:(nullable void (^)(NSError *_Nullable))completion {
+  [self firestoreWithConfiguredClient];
+  [self.client disableNetworkWithCompletion:completion];
+}
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/Firestore/Source/Public/FIRFirestore.h b/Firestore/Source/Public/FIRFirestore.h
index 91a96a5..4c85aba 100644
--- a/Firestore/Source/Public/FIRFirestore.h
+++ b/Firestore/Source/Public/FIRFirestore.h
@@ -139,6 +139,23 @@ NS_SWIFT_NAME(Firestore)
 + (void)enableLogging:(BOOL)logging
     DEPRECATED_MSG_ATTRIBUTE("Use FIRSetLoggerLevel(FIRLoggerLevelDebug) to enable logging");
 
+#pragma mark - Network
+
+/**
+ * Re-enables usage of the network by this Firestore instance after a prior call to
+ * `disableNetworkWithCompletion`. Completion block, if provided, will be called once network uasge
+ * has been enabled.
+ */
+- (void)enableNetworkWithCompletion:(nullable void (^)(NSError *_Nullable error))completion;
+
+/**
+ * Disables usage of the network by this Firestore instance. It can be re-enabled by via
+ * `enableNetworkWithCompletion`. While the network is disabled, any snapshot listeners or get calls
+ * will return results from cache and any write operations will be queued until the network is
+ * restored. The completion block, if provided, will be called once network usage has been disabled.
+ */
+- (void)disableNetworkWithCompletion:(nullable void (^)(NSError *_Nullable error))completion;
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/Firestore/Source/Remote/FSTRemoteStore.m b/Firestore/Source/Remote/FSTRemoteStore.m
index b1d91ea..1fec5e5 100644
--- a/Firestore/Source/Remote/FSTRemoteStore.m
+++ b/Firestore/Source/Remote/FSTRemoteStore.m
@@ -165,21 +165,20 @@ static const int kOnlineAttemptsBeforeFailure = 2;
  * onlineStateHandler as appropriate.
  */
 - (void)updateOnlineState:(FSTOnlineState)newState {
-  if (newState == FSTOnlineStateHealthy) {
-    // We've connected to watch at least once. Don't warn the developer about being offline going
-    // forward.
-    self.shouldWarnOffline = NO;
-  } else if (newState == FSTOnlineStateUnknown) {
-    // The state is set to unknown when a healthy stream is closed (e.g. due to a token timeout) or
-    // when we have no active listens and therefore there's no need to start the stream. Assuming
-    // there is (possibly in the future) an active listen, then we will eventually move to state
-    // Online or Failed, but we always want to make at least kOnlineAttemptsBeforeFailure attempts
-    // before failing, so we reset the count here.
-    self.watchStreamFailures = 0;
-  }
-
   // Update and broadcast the new state.
   if (newState != self.watchStreamOnlineState) {
+    if (newState == FSTOnlineStateHealthy) {
+      // We've connected to watch at least once. Don't warn the developer about being offline going
+      // forward.
+      self.shouldWarnOffline = NO;
+    } else if (newState == FSTOnlineStateUnknown) {
+      // The state is set to unknown when a healthy stream is closed (e.g. due to a token timeout)
+      // or when we have no active listens and therefore there's no need to start the stream.
+      // Assuming there is (possibly in the future) an active listen, then we will eventually move
+      // to state Online or Failed, but we always want to make at least kOnlineAttemptsBeforeFailure
+      // attempts before failing, so we reset the count here.
+      self.watchStreamFailures = 0;
+    }
     self.watchStreamOnlineState = newState;
     [self.onlineStateDelegate watchStreamDidChangeOnlineState:newState];
   }
@@ -214,8 +213,9 @@ static const int kOnlineAttemptsBeforeFailure = 2;
 }
 
 - (void)enableNetwork {
-  FSTAssert(self.watchStream == nil, @"enableNetwork: called with non-null watchStream.");
-  FSTAssert(self.writeStream == nil, @"enableNetwork: called with non-null writeStream.");
+  if ([self isNetworkEnabled]) {
+    return;
+  }
 
   // Create new streams (but note they're not started yet).
   self.watchStream = [self.datastore createWatchStream];
@@ -235,48 +235,47 @@ static const int kOnlineAttemptsBeforeFailure = 2;
 }
 
 - (void)disableNetwork {
+  [self disableNetworkInternal];
   // Set the FSTOnlineState to failed so get()'s return from cache, etc.
-  [self disableNetworkWithTargetOnlineState:FSTOnlineStateFailed];
+  [self updateOnlineState:FSTOnlineStateFailed];
 }
 
 /** Disables the network, setting the FSTOnlineState to the specified targetOnlineState. */
-- (void)disableNetworkWithTargetOnlineState:(FSTOnlineState)targetOnlineState {
-  // NOTE: We're guaranteed not to get any further events from these streams (not even a close
-  // event).
-  [self.watchStream stop];
-  [self.writeStream stop];
-
-  [self cleanUpWatchStreamState];
-  [self cleanUpWriteStreamState];
+- (void)disableNetworkInternal {
+  if ([self isNetworkEnabled]) {
+    // NOTE: We're guaranteed not to get any further events from these streams (not even a close
+    // event).
+    [self.watchStream stop];
+    [self.writeStream stop];
 
-  self.writeStream = nil;
-  self.watchStream = nil;
+    [self cleanUpWatchStreamState];
+    [self cleanUpWriteStreamState];
 
-  [self updateOnlineState:targetOnlineState];
+    self.writeStream = nil;
+    self.watchStream = nil;
+  }
 }
 
 #pragma mark Shutdown
 
 - (void)shutdown {
   FSTLog(@"FSTRemoteStore %p shutting down", (__bridge void *)self);
-
-  // For now, all shutdown logic is handled by disableNetwork(). We might expand on this in the
-  // future.
-  if ([self isNetworkEnabled]) {
-    // Set the FSTOnlineState to Unknown (rather than Failed) to avoid potentially triggering
-    // spurious listener events with cached data, etc.
-    [self disableNetworkWithTargetOnlineState:FSTOnlineStateUnknown];
-  }
+  [self disableNetworkInternal];
+  // Set the FSTOnlineState to Unknown (rather than Failed) to avoid potentially triggering
+  // spurious listener events with cached data, etc.
+  [self updateOnlineState:FSTOnlineStateUnknown];
 }
 
 - (void)userDidChange:(FSTUser *)user {
   FSTLog(@"FSTRemoteStore %p changing users: %@", (__bridge void *)self, user);
-
-  // Tear down and re-create our network streams. This will ensure we get a fresh auth token
-  // for the new user and re-fill the write pipeline with new mutations from the LocalStore
-  // (since mutations are per-user).
-  [self disableNetworkWithTargetOnlineState:FSTOnlineStateUnknown];
-  [self enableNetwork];
+  if ([self isNetworkEnabled]) {
+    // Tear down and re-create our network streams. This will ensure we get a fresh auth token
+    // for the new user and re-fill the write pipeline with new mutations from the LocalStore
+    // (since mutations are per-user).
+    [self disableNetworkInternal];
+    [self updateOnlineState:FSTOnlineStateUnknown];
+    [self enableNetwork];
+  }
 }
 
 #pragma mark Watch Stream