Verify large write batches support (#2321)

* Reorder tests to match Android
* Add missing NS_ASSUME_NONNULL_BEGIN
* Add missing [self awaitExpectations]
* Port can write very large batches from Android

Author: wilhuff

Firestore/Example/Tests/Integration/API/FIRWriteBatchTests.mm

@@ -24,6 +24,14 @@
 #import "Firestore/Example/Tests/Util/FSTEventAccumulator.h"
 #import "Firestore/Example/Tests/Util/FSTIntegrationTestCase.h"
 
+#include "Firestore/core/src/firebase/firestore/util/autoid.h"
+#include "Firestore/core/src/firebase/firestore/util/string_apple.h"
+
+using firebase::firestore::util::CreateAutoId;
+using firebase::firestore::util::WrapNSString;
+
+NS_ASSUME_NONNULL_BEGIN
+
 @interface FIRWriteBatchTests : FSTIntegrationTestCase
 @end
 
@@ -124,6 +132,52 @@ - (void)testCannotUpdateNonexistentDocuments {
   XCTAssertFalse(result.exists);
 }
 
+- (void)testUpdateFieldsWithDots {
+  FIRDocumentReference *doc = [self documentRef];
+
+  XCTestExpectation *expectation = [self expectationWithDescription:@"testUpdateFieldsWithDots"];
+  FIRWriteBatch *batch = [doc.firestore batch];
+  [batch setData:@{@"a.b" : @"old", @"c.d" : @"old"} forDocument:doc];
+  [batch updateData:@{[[FIRFieldPath alloc] initWithFields:@[ @"a.b" ]] : @"new"} forDocument:doc];
+
+  [batch commitWithCompletion:^(NSError *_Nullable error) {
+    XCTAssertNil(error);
+    [doc getDocumentWithCompletion:^(FIRDocumentSnapshot *snapshot, NSError *error) {
+      XCTAssertNil(error);
+      XCTAssertEqualObjects(snapshot.data, (@{@"a.b" : @"new", @"c.d" : @"old"}));
+    }];
+    [expectation fulfill];
+  }];
+
+  [self awaitExpectations];
+}
+
+- (void)testUpdateNestedFields {
+  FIRDocumentReference *doc = [self documentRef];
+
+  XCTestExpectation *expectation = [self expectationWithDescription:@"testUpdateNestedFields"];
+  FIRWriteBatch *batch = [doc.firestore batch];
+  [batch setData:@{@"a" : @{@"b" : @"old"}, @"c" : @{@"d" : @"old"}, @"e" : @{@"f" : @"old"}}
+      forDocument:doc];
+  [batch
+       updateData:@{@"a.b" : @"new", [[FIRFieldPath alloc] initWithFields:@[ @"c", @"d" ]] : @"new"}
+      forDocument:doc];
+  [batch commitWithCompletion:^(NSError *_Nullable error) {
+    XCTAssertNil(error);
+    [doc getDocumentWithCompletion:^(FIRDocumentSnapshot *snapshot, NSError *error) {
+      XCTAssertNil(error);
+      XCTAssertEqualObjects(snapshot.data, (@{
+                              @"a" : @{@"b" : @"new"},
+                              @"c" : @{@"d" : @"new"},
+                              @"e" : @{@"f" : @"old"}
+                            }));
+    }];
+    [expectation fulfill];
+  }];
+
+  [self awaitExpectations];
+}
+
 - (void)testDeleteDocuments {
   FIRDocumentReference *doc = [self documentRef];
   [self writeDocumentRef:doc data:@{@"foo" : @"bar"}];
@@ -161,6 +215,7 @@ - (void)testBatchesCommitAtomicallyRaisingCorrectEvents {
     XCTAssertNil(error);
     [expectation fulfill];
   }];
+  [self awaitExpectations];
 
   FIRQuerySnapshot *localSnap = [accumulator awaitEventWithName:@"local event"];
   XCTAssertTrue(localSnap.metadata.hasPendingWrites);
@@ -192,6 +247,7 @@ - (void)testBatchesFailAtomicallyRaisingCorrectEvents {
     XCTAssertEqual(error.code, FIRFirestoreErrorCodeNotFound);
     [expectation fulfill];
   }];
+  [self awaitExpectations];
 
   // Local event with the set document.
   FIRQuerySnapshot *localSnap = [accumulator awaitEventWithName:@"local event"];
@@ -223,6 +279,7 @@ - (void)testWriteTheSameServerTimestampAcrossWrites {
     XCTAssertNil(error);
     [expectation fulfill];
   }];
+  [self awaitExpectations];
 
   FIRQuerySnapshot *localSnap = [accumulator awaitEventWithName:@"local event"];
   XCTAssertTrue(localSnap.metadata.hasPendingWrites);
@@ -254,6 +311,7 @@ - (void)testCanWriteTheSameDocumentMultipleTimes {
     XCTAssertNil(error);
     [expectation fulfill];
   }];
+  [self awaitExpectations];
 
   FIRDocumentSnapshot *localSnap = [accumulator awaitEventWithName:@"local event"];
   XCTAssertTrue(localSnap.metadata.hasPendingWrites);
@@ -265,50 +323,39 @@ - (void)testCanWriteTheSameDocumentMultipleTimes {
   XCTAssertEqualObjects(serverSnap.data, (@{@"a" : @1, @"b" : @2, @"when" : when}));
 }
 
-- (void)testUpdateFieldsWithDots {
-  FIRDocumentReference *doc = [self documentRef];
+- (void)testCanWriteVeryLargeBatches {
+  // On Android, SQLite Cursors are limited reading no more than 2 MB per row (despite being able
+  // to write very large values). This test verifies that the local MutationQueue is not subject
+  // to this limitation.
+
+  // Create a map containing nearly 1 MB of data. Note that if you use 1024 below this will create
+  // a document larger than 1 MB, which will be rejected by the backend as too large.
+  NSString *kb = [@"" stringByPaddingToLength:1000 withString:@"a" startingAtIndex:0];
+  NSMutableDictionary<NSString *, id> *values = [NSMutableDictionary dictionary];
+  for (int i = 0; i < 1000; i++) {
+    values[WrapNSString(CreateAutoId())] = kb;
+  }
 
-  XCTestExpectation *expectation = [self expectationWithDescription:@"testUpdateFieldsWithDots"];
+  FIRDocumentReference *doc = [self documentRef];
   FIRWriteBatch *batch = [doc.firestore batch];
-  [batch setData:@{@"a.b" : @"old", @"c.d" : @"old"} forDocument:doc];
-  [batch updateData:@{[[FIRFieldPath alloc] initWithFields:@[ @"a.b" ]] : @"new"} forDocument:doc];
-
-  [batch commitWithCompletion:^(NSError *_Nullable error) {
-    XCTAssertNil(error);
-    [doc getDocumentWithCompletion:^(FIRDocumentSnapshot *snapshot, NSError *error) {
-      XCTAssertNil(error);
-      XCTAssertEqualObjects(snapshot.data, (@{@"a.b" : @"new", @"c.d" : @"old"}));
-    }];
-    [expectation fulfill];
-  }];
 
-  [self awaitExpectations];
-}
-
-- (void)testUpdateNestedFields {
-  FIRDocumentReference *doc = [self documentRef];
+  // Write a batch containing 3 copies of the data, creating a ~3 MB batch. Writing to the same
+  // document in a batch is allowed and so long as the net size of the document is under 1 MB the
+  // batch is allowed.
+  [batch setData:values forDocument:doc];
+  for (int i = 0; i < 2; i++) {
+    [batch updateData:values forDocument:doc];
+  }
 
-  XCTestExpectation *expectation = [self expectationWithDescription:@"testUpdateNestedFields"];
-  FIRWriteBatch *batch = [doc.firestore batch];
-  [batch setData:@{@"a" : @{@"b" : @"old"}, @"c" : @{@"d" : @"old"}, @"e" : @{@"f" : @"old"}}
-      forDocument:doc];
-  [batch
-       updateData:@{@"a.b" : @"new", [[FIRFieldPath alloc] initWithFields:@[ @"c", @"d" ]] : @"new"}
-      forDocument:doc];
+  XCTestExpectation *expectation = [self expectationWithDescription:@"batch written"];
   [batch commitWithCompletion:^(NSError *_Nullable error) {
     XCTAssertNil(error);
-    [doc getDocumentWithCompletion:^(FIRDocumentSnapshot *snapshot, NSError *error) {
-      XCTAssertNil(error);
-      XCTAssertEqualObjects(snapshot.data, (@{
-                              @"a" : @{@"b" : @"new"},
-                              @"c" : @{@"d" : @"new"},
-                              @"e" : @{@"f" : @"old"}
-                            }));
-    }];
     [expectation fulfill];
   }];
-
   [self awaitExpectations];
+
+  FIRDocumentSnapshot *snap = [self readDocumentForRef:doc];
+  XCTAssertEqualObjects(values, snap.data);
 }
 
 // Returns how much memory the test application is currently using, in megabytes (fractional part is
@@ -363,3 +410,5 @@ - (void)testReasonableMemoryUsageForLotsOfMutations {
 }
 
 @end
+
+NS_ASSUME_NONNULL_END