Fix test code and add a new test.

Author: ehsannas

packages/firestore/test/unit/local/document_overlay_cache.test.ts

@@ -170,14 +170,67 @@ function genericDocumentOverlayCacheTests(): void {
   it('can delete overlays repeatedly', async () => {
     const m = patchMutation('coll/doc1', { 'foo': 'bar' });
     await saveOverlaysForMutations(2, m);
-    await overlayCache.removeOverlaysForBatchId(2);
+    await overlayCache.removeOverlaysForBatchId(
+      documentKeySet(key('coll/doc1')),
+      2
+    );
     expect(await overlayCache.getOverlay(key('coll/doc1'))).to.equal(null);
 
     // Repeat
-    await overlayCache.removeOverlaysForBatchId(2);
+    await overlayCache.removeOverlaysForBatchId(
+      documentKeySet(key('coll/doc1')),
+      2
+    );
     expect(await overlayCache.getOverlay(key('coll/doc1'))).to.equal(null);
   });
 
+  it('can delete overlays', async () => {
+    const m1 = patchMutation('coll1/doc1', { 'a': 'b' });
+    const m2 = patchMutation('coll1/doc2', { 'a': 'b' });
+    const m3 = patchMutation('coll1/doc3', { 'a': 'b' });
+    const m4 = patchMutation('coll1/doc4', { 'a': 'b' });
+    const m5 = patchMutation('coll2/doc1/coll3/doc1', { 'a': 'b' });
+    const m6 = patchMutation('coll2/doc2/coll3/doc2', { 'a': 'b' });
+    await saveOverlaysForMutations(2, m1, m2);
+    await saveOverlaysForMutations(3, m3);
+    await saveOverlaysForMutations(4, m4);
+    await saveOverlaysForMutations(2, m5, m6);
+
+    // Remove documents with batch id 2.
+    const set1 = documentKeySet(
+      key('coll1/doc1'),
+      key('coll1/doc2'),
+      key('coll2/doc1/coll3/doc1'),
+      key('coll2/doc2/coll3/doc2')
+    );
+    await overlayCache.removeOverlaysForBatchId(set1, 2);
+    expect(await overlayCache.getOverlay(key('coll1/doc1'))).to.equal(null);
+    expect(await overlayCache.getOverlay(key('coll1/doc2'))).to.equal(null);
+    expect(
+      await overlayCache.getOverlay(key('coll2/doc1/coll3/doc1'))
+    ).to.equal(null);
+    expect(
+      await overlayCache.getOverlay(key('coll2/doc2/coll3/doc2'))
+    ).to.equal(null);
+    expect(await overlayCache.getOverlay(key('coll1/doc3'))).to.not.equal(null);
+    expect(await overlayCache.getOverlay(key('coll1/doc4'))).to.not.equal(null);
+
+    // Remove documents with batch id 3.
+    await overlayCache.removeOverlaysForBatchId(
+      documentKeySet(key('coll1/doc3')),
+      3
+    );
+    expect(await overlayCache.getOverlay(key('coll1/doc3'))).to.equal(null);
+    expect(await overlayCache.getOverlay(key('coll1/doc4'))).to.not.equal(null);
+
+    // Remove documents with batch id 4.
+    await overlayCache.removeOverlaysForBatchId(
+      documentKeySet(key('coll1/doc4')),
+      4
+    );
+    expect(await overlayCache.getOverlay(key('coll1/doc4'))).to.equal(null);
+  });
+
   it('can get all overlays for collection', async () => {
     const m1 = patchMutation('coll/doc1', { 'foo': 'bar' });
     const m2 = setMutation('coll/doc2', { 'foo': 'bar' });

packages/firestore/test/unit/local/test_document_overlay_cache.ts

@@ -17,6 +17,7 @@
 
 import { DocumentOverlayCache } from '../../../src/local/document_overlay_cache';
 import { Persistence } from '../../../src/local/persistence';
+import { DocumentKeySet } from '../../../src/model/collections';
 import { DocumentKey } from '../../../src/model/document_key';
 import { Mutation } from '../../../src/model/mutation';
 import { Overlay } from '../../../src/model/overlay';
@@ -89,12 +90,15 @@ export class TestDocumentOverlayCache {
     );
   }
 
-  removeOverlaysForBatchId(batchId: number): Promise<void> {
+  removeOverlaysForBatchId(
+    documentKeys: DocumentKeySet,
+    batchId: number
+  ): Promise<void> {
     return this.persistence.runTransaction(
       'removeOverlaysForBatchId',
       'readwrite-primary',
       txn => {
-        return this.cache.removeOverlaysForBatchId(txn, batchId);
+        return this.cache.removeOverlaysForBatchId(txn, documentKeys, batchId);
       }
     );
   }