@@ -32,7 +32,9 @@ import {
32
32
PrivateSettings ,
33
33
SnapshotListenOptions ,
34
34
newTestFirestore ,
35
- newTestApp
35
+ newTestApp ,
36
+ writeBatch ,
37
+ WriteBatch
36
38
} from './firebase_export' ;
37
39
import {
38
40
ALT_PROJECT_ID ,
@@ -317,11 +319,34 @@ export function withTestCollectionSettings(
317
319
const collectionId = 'test-collection-' + doc ( collection ( testDb , 'x' ) ) . id ;
318
320
const testCollection = collection ( testDb , collectionId ) ;
319
321
const setupCollection = collection ( setupDb , collectionId ) ;
320
- const sets : Array < Promise < void > > = [ ] ;
321
- Object . keys ( docs ) . forEach ( key => {
322
- sets . push ( setDoc ( doc ( setupCollection , key ) , docs [ key ] ) ) ;
323
- } ) ;
324
- return Promise . all ( sets ) . then ( ( ) => fn ( testCollection , testDb ) ) ;
322
+
323
+ const writeBatchCommits : Array < Promise < void > > = [ ] ;
324
+ let writeBatch_ : WriteBatch | null = null ;
325
+ let writeBatchSize = 0 ;
326
+
327
+ for ( const key of Object . keys ( docs ) ) {
328
+ if ( writeBatch_ === null ) {
329
+ writeBatch_ = writeBatch ( setupDb ) ;
330
+ }
331
+
332
+ writeBatch_ . set ( doc ( setupCollection , key ) , docs [ key ] ) ;
333
+ writeBatchSize ++ ;
334
+
335
+ // Write batches are capped at 500 writes. Use 400 just to be safe.
336
+ if ( writeBatchSize === 400 ) {
337
+ writeBatchCommits . push ( writeBatch_ . commit ( ) ) ;
338
+ writeBatch_ = null ;
339
+ writeBatchSize = 0 ;
340
+ }
341
+ }
342
+
343
+ if ( writeBatch_ !== null ) {
344
+ writeBatchCommits . push ( writeBatch_ . commit ( ) ) ;
345
+ }
346
+
347
+ return Promise . all ( writeBatchCommits ) . then ( ( ) =>
348
+ fn ( testCollection , testDb )
349
+ ) ;
325
350
}
326
351
) ;
327
352
}
0 commit comments