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

4 files changed, 76 insertions, 35 deletions

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},