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

8 files changed, 144 insertions, 76 deletions

diff --git a/.gitignore b/.gitignore
index 93e3974..61fc41b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -57,3 +57,4 @@ Podfile.lock
 
 # CMake
 .downloads
+.idea/
diff --git a/Firestore/Example/SwiftBuildTest/main.swift b/Firestore/Example/SwiftBuildTest/main.swift
index c05f378..f62bf48 100644
--- a/Firestore/Example/SwiftBuildTest/main.swift
+++ b/Firestore/Example/SwiftBuildTest/main.swift
@@ -39,6 +39,8 @@ func main() {
 
     listenToDocuments(matching: query);
 
+    enableDisableNetwork(db: db);
+
     types();
 }
 
@@ -131,6 +133,23 @@ func writeDocument(at docRef: DocumentReference) {
     }
 }
 
+func enableDisableNetwork(db db: Firestore) {
+    // closure syntax
+    db.disableNetwork(completion: { (error) in
+        if let e = error {
+            print("Uh oh! \(e)")
+            return
+        }
+    })
+    // trailing block syntax
+    db.enableNetwork { (error) in
+        if let e = error {
+            print("Uh oh! \(e)")
+            return
+        }
+    }
+}
+
 func writeDocuments(at docRef: DocumentReference, database db: Firestore) {
   var batch: WriteBatch;
 
diff --git a/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m b/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m
index 1c31242..f557ee6 100644
--- a/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m
+++ b/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m
@@ -16,6 +16,7 @@
 
 @import FirebaseFirestore;
 
+#import <FirebaseFirestore/FIRFirestore.h>
 #import <XCTest/XCTest.h>
 
 #import "Firestore/Example/Tests/Util/FSTIntegrationTestCase.h"
@@ -851,7 +852,7 @@
   FIRFirestore *firestore = doc.firestore;
   NSDictionary<NSString *, id> *data = @{@"a" : @"b"};
 
-  [firestore.client disableNetworkWithCompletion:^(NSError *error) {
+  [firestore disableNetworkWithCompletion:^(NSError *error) {
     XCTAssertNil(error);
 
     [doc setData:data
@@ -860,7 +861,7 @@
           [writeEpectation fulfill];
         }];
 
-    [firestore.client enableNetworkWithCompletion:^(NSError *error) {
+    [firestore enableNetworkWithCompletion:^(NSError *error) {
       XCTAssertNil(error);
       [networkExpectation fulfill];
     }];
@@ -890,7 +891,7 @@
 
   __weak FIRDocumentReference *weakDoc = doc;
 
-  [firestore.client disableNetworkWithCompletion:^(NSError *error) {
+  [firestore disableNetworkWithCompletion:^(NSError *error) {
     XCTAssertNil(error);
     [doc setData:data
         completion:^(NSError *_Nullable error) {
@@ -911,7 +912,7 @@
       // Verify that we are reading from cache.
       XCTAssertTrue(snapshot.metadata.fromCache);
       XCTAssertEqualObjects(snapshot.data, data);
-      [firestore.client enableNetworkWithCompletion:^(NSError *error) {
+      [firestore enableNetworkWithCompletion:^(NSError *error) {
         [networkExpectation fulfill];
       }];
     }];
@@ -938,4 +939,25 @@
   [self readSnapshotForRef:[self documentRef] requireOnline:YES];
 }
 
+- (void)testCanDisableNetwork {
+  FIRDocumentReference *doc = [self documentRef];
+  FIRFirestore *firestore = doc.firestore;
+
+  [firestore enableNetworkWithCompletion:[self completionForExpectationWithName:@"Enable network"]];
+  [self awaitExpectations];
+  [firestore
+      enableNetworkWithCompletion:[self completionForExpectationWithName:@"Enable network again"]];
+  [self awaitExpectations];
+  [firestore
+      disableNetworkWithCompletion:[self completionForExpectationWithName:@"Disable network"]];
+  [self awaitExpectations];
+  [firestore
+      disableNetworkWithCompletion:[self
+                                       completionForExpectationWithName:@"Disable network again"]];
+  [self awaitExpectations];
+  [firestore
+      enableNetworkWithCompletion:[self completionForExpectationWithName:@"Final enable network"]];
+  [self awaitExpectations];
+}
+
 @end
diff --git a/Firestore/Example/Tests/Integration/API/FIRListenerRegistrationTests.m b/Firestore/Example/Tests/Integration/API/FIRListenerRegistrationTests.m
index 9751844..52d73b1 100644
--- a/Firestore/Example/Tests/Integration/API/FIRListenerRegistrationTests.m
+++ b/Firestore/Example/Tests/Integration/API/FIRListenerRegistrationTests.m
@@ -128,35 +128,4 @@
   [two remove];
 }
 
-- (void)testWatchSurvivesNetworkDisconnect {
-  XCTestExpectation *testExpectiation =
-      [self expectationWithDescription:@"testWatchSurvivesNetworkDisconnect"];
-
-  FIRCollectionReference *collectionRef = [self collectionRef];
-  FIRDocumentReference *docRef = [collectionRef documentWithAutoID];
-
-  FIRFirestore *firestore = collectionRef.firestore;
-
-  FIRQueryListenOptions *options = [[[FIRQueryListenOptions options]
-      includeDocumentMetadataChanges:YES] includeQueryMetadataChanges:YES];
-
-  [collectionRef addSnapshotListenerWithOptions:options
-                                       listener:^(FIRQuerySnapshot *snapshot, NSError *error) {
-                                         XCTAssertNil(error);
-                                         if (!snapshot.empty && !snapshot.metadata.fromCache) {
-                                           [testExpectiation fulfill];
-                                         }
-                                       }];
-
-  [firestore.client disableNetworkWithCompletion:^(NSError *error) {
-    XCTAssertNil(error);
-    [docRef setData:@{@"foo" : @"bar"}];
-    [firestore.client enableNetworkWithCompletion:^(NSError *error) {
-      XCTAssertNil(error);
-    }];
-  }];
-
-  [self awaitExpectations];
-}
-
 @end
diff --git a/Firestore/Example/Tests/Integration/API/FIRQueryTests.m b/Firestore/Example/Tests/Integration/API/FIRQueryTests.m
index e92bbcf..251270a 100644
--- a/Firestore/Example/Tests/Integration/API/FIRQueryTests.m
+++ b/Firestore/Example/Tests/Integration/API/FIRQueryTests.m
@@ -212,6 +212,37 @@
   XCTAssertEqualObjects(FIRQuerySnapshotGetData(docs), (@[ testDocs[@"ab"], testDocs[@"ba"] ]));
 }
 
+- (void)testWatchSurvivesNetworkDisconnect {
+  XCTestExpectation *testExpectiation =
+      [self expectationWithDescription:@"testWatchSurvivesNetworkDisconnect"];
+
+  FIRCollectionReference *collectionRef = [self collectionRef];
+  FIRDocumentReference *docRef = [collectionRef documentWithAutoID];
+
+  FIRFirestore *firestore = collectionRef.firestore;
+
+  FIRQueryListenOptions *options = [[[FIRQueryListenOptions options]
+      includeDocumentMetadataChanges:YES] includeQueryMetadataChanges:YES];
+
+  [collectionRef addSnapshotListenerWithOptions:options
+                                       listener:^(FIRQuerySnapshot *snapshot, NSError *error) {
+                                         XCTAssertNil(error);
+                                         if (!snapshot.empty && !snapshot.metadata.fromCache) {
+                                           [testExpectiation fulfill];
+                                         }
+                                       }];
+
+  [firestore disableNetworkWithCompletion:^(NSError *error) {
+    XCTAssertNil(error);
+    [docRef setData:@{@"foo" : @"bar"}];
+    [firestore enableNetworkWithCompletion:^(NSError *error) {
+      XCTAssertNil(error);
+    }];
+  }];
+
+  [self awaitExpectations];
+}
+
 - (void)testQueriesFireFromCacheWhenOffline {
   NSDictionary *testDocs = @{
     @"a" : @{@"foo" : @1},
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