Validate target purpose in spec tests (#7113)

milaGGL · web-flow · commit a2505bd7282e · 2023-03-10T16:52:09.000-08:00
diff --git a/packages/firestore/test/unit/specs/existence_filter_spec.test.ts b/packages/firestore/test/unit/specs/existence_filter_spec.test.ts
@@ -16,6 +16,7 @@
  */
 
 import { newQueryForPath } from '../../../src/core/query';
+import { TargetPurpose } from '../../../src/local/target_data';
 import { Code } from '../../../src/util/error';
 import { deletedDoc, doc, query } from '../../util/helpers';
 
@@ -61,7 +62,11 @@ describeSpec('Existence Filters:', [], () => {
       .expectEvents(query1, {})
       .watchFilters([query1], [doc1.key])
       .watchSnapshots(2000)
-      .expectEvents(query1, { fromCache: true });
+      .expectEvents(query1, { fromCache: true })
+      .expectActiveTargets({
+        query: query1,
+        targetPurpose: TargetPurpose.ExistenceFilterMismatch
+      });
   });
 
   specTest('Existence filter ignored with pending target', [], () => {
@@ -100,7 +105,11 @@ describeSpec('Existence Filters:', [], () => {
         .watchSnapshots(2000)
         // query is now marked as "inconsistent" because of filter mismatch
         .expectEvents(query1, { fromCache: true })
-        .expectActiveTargets({ query: query1, resumeToken: '' })
+        .expectActiveTargets({
+          query: query1,
+          targetPurpose: TargetPurpose.ExistenceFilterMismatch,
+          resumeToken: ''
+        })
         .watchRemoves(query1) // Acks removal of query
         .watchAcksFull(query1, 2000, doc1)
         .expectLimboDocs(doc2.key) // doc2 is now in limbo
@@ -134,7 +143,11 @@ describeSpec('Existence Filters:', [], () => {
         .watchSnapshots(2000)
         // query is now marked as "inconsistent" because of filter mismatch
         .expectEvents(query1, { fromCache: true })
-        .expectActiveTargets({ query: query1, resumeToken: '' })
+        .expectActiveTargets({
+          query: query1,
+          targetPurpose: TargetPurpose.ExistenceFilterMismatch,
+          resumeToken: ''
+        })
         .watchRemoves(query1) // Acks removal of query
         .watchAcksFull(query1, 2000, doc1)
         .expectLimboDocs(doc2.key) // doc2 is now in limbo
@@ -165,7 +178,11 @@ describeSpec('Existence Filters:', [], () => {
         // The query result includes doc3, but is marked as "inconsistent"
         // due to the filter mismatch
         .expectEvents(query1, { added: [doc3], fromCache: true })
-        .expectActiveTargets({ query: query1, resumeToken: '' })
+        .expectActiveTargets({
+          query: query1,
+          targetPurpose: TargetPurpose.ExistenceFilterMismatch,
+          resumeToken: ''
+        })
         .watchRemoves(query1) // Acks removal of query
         .watchAcksFull(query1, 3000, doc1, doc2, doc3)
         .expectEvents(query1, { added: [doc2] })
@@ -197,7 +214,11 @@ describeSpec('Existence Filters:', [], () => {
         .watchSnapshots(2000)
         // query is now marked as "inconsistent" because of filter mismatch
         .expectEvents(query1, { fromCache: true })
-        .expectActiveTargets({ query: query1, resumeToken: '' })
+        .expectActiveTargets({
+          query: query1,
+          targetPurpose: TargetPurpose.ExistenceFilterMismatch,
+          resumeToken: ''
+        })
         .watchRemoves(query1) // Acks removal of query
         .watchAcksFull(query1, 2000, doc1)
         .expectLimboDocs(doc2.key) // doc2 is now in limbo
@@ -232,6 +253,10 @@ describeSpec('Existence Filters:', [], () => {
           .watchFilters([query1], [doc1.key]) // doc2 was deleted
           .watchSnapshots(2000)
           .expectEvents(query1, { fromCache: true })
+          .expectActiveTargets({
+            query: query1,
+            targetPurpose: TargetPurpose.ExistenceFilterMismatch
+          })
           // The SDK is unable to re-run the query, and does not remove doc2
           .restart()
           .userListens(query1)
diff --git a/packages/firestore/test/unit/specs/limbo_spec.test.ts b/packages/firestore/test/unit/specs/limbo_spec.test.ts
@@ -20,6 +20,7 @@ import {
   newQueryForPath,
   queryWithLimit
 } from '../../../src/core/query';
+import { TargetPurpose } from '../../../src/local/target_data';
 import { TimerId } from '../../../src/util/async_queue';
 import { Code } from '../../../src/util/error';
 import { deletedDoc, doc, filter, orderBy, query } from '../../util/helpers';
@@ -889,7 +890,11 @@ describeSpec('Limbo Documents:', [], () => {
           // The view now contains the docAs and the docBs (6 documents), but
           // the existence filter indicated only 3 should match. This causes
           // the client to re-listen without a resume token.
-          .expectActiveTargets({ query: query1, resumeToken: '' })
+          .expectActiveTargets({
+            query: query1,
+            targetPurpose: TargetPurpose.ExistenceFilterMismatch,
+            resumeToken: ''
+          })
           // When the existence filter mismatch was detected, the client removed
           // then re-added the target. Watch needs to acknowledge the removal.
           .watchRemoves(query1)
diff --git a/packages/firestore/test/unit/specs/limit_spec.test.ts b/packages/firestore/test/unit/specs/limit_spec.test.ts
@@ -16,6 +16,7 @@
  */
 
 import { LimitType, queryWithLimit } from '../../../src/core/query';
+import { TargetPurpose } from '../../../src/local/target_data';
 import { deletedDoc, doc, filter, orderBy, query } from '../../util/helpers';
 
 import { describeSpec, specTest } from './describe_spec';
@@ -343,7 +344,11 @@ describeSpec('Limits:', [], () => {
           .watchSends({ affects: [limitQuery] }, secondDocument)
           .watchFilters([limitQuery], [secondDocument.key])
           .watchSnapshots(1004)
-          .expectActiveTargets({ query: limitQuery, resumeToken: '' })
+          .expectActiveTargets({
+            query: limitQuery,
+            targetPurpose: TargetPurpose.ExistenceFilterMismatch,
+            resumeToken: ''
+          })
           .watchRemoves(limitQuery)
           .watchAcksFull(limitQuery, 1005, secondDocument)
           // The snapshot after the existence filter mismatch triggers limbo
diff --git a/packages/firestore/test/unit/specs/recovery_spec.test.ts b/packages/firestore/test/unit/specs/recovery_spec.test.ts
@@ -16,6 +16,7 @@
  */
 
 import { newQueryForPath } from '../../../src/core/query';
+import { TargetPurpose } from '../../../src/local/target_data';
 import { TimerId } from '../../../src/util/async_queue';
 import { Code } from '../../../src/util/error';
 import { deletedDoc, doc, filter, query } from '../../util/helpers';
@@ -137,7 +138,9 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
           // The primary client 0 receives a notification that the query can
           // be released, but it can only process the change after we recover
           // the database.
-          .expectActiveTargets({ query: query1 })
+          .expectActiveTargets({
+            query: query1
+          })
           .recoverDatabase()
           .runTimer(TimerId.AsyncQueueRetry)
           .expectActiveTargets()
@@ -641,7 +644,10 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
             query: filteredQuery,
             resumeToken: 'resume-token-2000'
           },
-          { query: limboQuery }
+          {
+            query: limboQuery,
+            targetPurpose: TargetPurpose.LimboResolution
+          }
         )
         .watchAcksFull(filteredQuery, 4000)
         .watchAcksFull(limboQuery, 4000, doc1b)
@@ -682,7 +688,7 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
       .runTimer(TimerId.AsyncQueueRetry)
       .expectActiveTargets(
         { query: filteredQuery, resumeToken: 'resume-token-2000' },
-        { query: limboQuery }
+        { query: limboQuery, targetPurpose: TargetPurpose.LimboResolution }
       )
       .watchAcksFull(filteredQuery, 3000)
       .watchRemoves(
@@ -719,15 +725,19 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
           .expectActiveTargets()
           .recoverDatabase()
           .runTimer(TimerId.AsyncQueueRetry)
-          .expectActiveTargets({ query: query1 })
+          .expectActiveTargets({
+            query: query1
+          })
           .expectEvents(query1, { removed: [doc1], fromCache: true })
           .failDatabaseTransactions('Handle user change')
           .changeUser('user1')
           // The network is offline due to the failed user change
           .expectActiveTargets()
           .recoverDatabase()
           .runTimer(TimerId.AsyncQueueRetry)
-          .expectActiveTargets({ query: query1 })
+          .expectActiveTargets({
+            query: query1
+          })
           .expectEvents(query1, {
             added: [doc1],
             fromCache: true,
@@ -763,7 +773,9 @@ describeSpec('Persistence Recovery', ['no-ios', 'no-android'], () => {
           .changeUser('user1')
           .recoverDatabase()
           .runTimer(TimerId.AsyncQueueRetry)
-          .expectActiveTargets({ query: query1 })
+          .expectActiveTargets({
+            query: query1
+          })
           // We are now user 2
           .expectEvents(query1, { removed: [doc1], fromCache: true })
           .runTimer(TimerId.AsyncQueueRetry)
diff --git a/packages/firestore/test/unit/specs/spec_builder.ts b/packages/firestore/test/unit/specs/spec_builder.ts
@@ -28,6 +28,7 @@ import {
 import { canonifyTarget, Target, targetEquals } from '../../../src/core/target';
 import { TargetIdGenerator } from '../../../src/core/target_id_generator';
 import { TargetId } from '../../../src/core/types';
+import { TargetPurpose } from '../../../src/local/target_data';
 import { Document } from '../../../src/model/document';
 import { DocumentKey } from '../../../src/model/document_key';
 import { FieldIndex } from '../../../src/model/field_index';
@@ -73,6 +74,7 @@ export interface LimboMap {
 
 export interface ActiveTargetSpec {
   queries: SpecQuery[];
+  targetPurpose?: TargetPurpose;
   resumeToken?: string;
   readTime?: TestSnapshotVersion;
 }
@@ -299,7 +301,9 @@ export class SpecBuilder {
       throw new Error("Can't restore an unknown query: " + query);
     }
 
-    this.addQueryToActiveTargets(targetId!, query, { resumeToken });
+    this.addQueryToActiveTargets(targetId!, query, {
+      resumeToken
+    });
 
     const currentStep = this.currentStep!;
     currentStep.expectedState = currentStep.expectedState || {};
@@ -525,18 +529,24 @@ export class SpecBuilder {
   expectActiveTargets(
     ...targets: Array<{
       query: Query;
+      targetPurpose?: TargetPurpose;
       resumeToken?: string;
       readTime?: TestSnapshotVersion;
     }>
   ): this {
     this.assertStep('Active target expectation requires previous step');
     const currentStep = this.currentStep!;
     this.clientState.activeTargets = {};
-    targets.forEach(({ query, resumeToken, readTime }) => {
-      this.addQueryToActiveTargets(this.getTargetId(query), query, {
-        resumeToken,
-        readTime
-      });
+    targets.forEach(({ query, targetPurpose, resumeToken, readTime }) => {
+      this.addQueryToActiveTargets(
+        this.getTargetId(query),
+        query,
+        {
+          resumeToken,
+          readTime
+        },
+        targetPurpose
+      );
     });
     currentStep.expectedState = currentStep.expectedState || {};
     currentStep.expectedState.activeTargets = { ...this.activeTargets };
@@ -567,7 +577,8 @@ export class SpecBuilder {
       this.addQueryToActiveTargets(
         this.limboMapping[path],
         newQueryForPath(key.path),
-        { resumeToken: '' }
+        { resumeToken: '' },
+        TargetPurpose.LimboResolution
       );
     });
 
@@ -1077,7 +1088,8 @@ export class SpecBuilder {
   private addQueryToActiveTargets(
     targetId: number,
     query: Query,
-    resume?: ResumeSpec
+    resume: ResumeSpec = {},
+    targetPurpose?: TargetPurpose
   ): void {
     if (this.activeTargets[targetId]) {
       const activeQueries = this.activeTargets[targetId].queries;
@@ -1089,21 +1101,24 @@ export class SpecBuilder {
         // `query` is not added yet.
         this.activeTargets[targetId] = {
           queries: [SpecBuilder.queryToSpec(query), ...activeQueries],
-          resumeToken: resume?.resumeToken || '',
-          readTime: resume?.readTime
+          targetPurpose,
+          resumeToken: resume.resumeToken || '',
+          readTime: resume.readTime
         };
       } else {
         this.activeTargets[targetId] = {
           queries: activeQueries,
-          resumeToken: resume?.resumeToken || '',
-          readTime: resume?.readTime
+          targetPurpose,
+          resumeToken: resume.resumeToken || '',
+          readTime: resume.readTime
         };
       }
     } else {
       this.activeTargets[targetId] = {
         queries: [SpecBuilder.queryToSpec(query)],
-        resumeToken: resume?.resumeToken || '',
-        readTime: resume?.readTime
+        targetPurpose,
+        resumeToken: resume.resumeToken || '',
+        readTime: resume.readTime
       };
     }
   }
@@ -1115,6 +1130,7 @@ export class SpecBuilder {
     if (queriesAfterRemoval.length > 0) {
       this.activeTargets[targetId] = {
         queries: queriesAfterRemoval,
+        targetPurpose: this.activeTargets[targetId].targetPurpose,
         resumeToken: this.activeTargets[targetId].resumeToken
       };
     } else {
diff --git a/packages/firestore/test/unit/specs/spec_test_components.ts b/packages/firestore/test/unit/specs/spec_test_components.ts
@@ -53,6 +53,7 @@ import { Mutation } from '../../../src/model/mutation';
 import { encodeBase64 } from '../../../src/platform/base64';
 import { newSerializer } from '../../../src/platform/serializer';
 import * as api from '../../../src/protos/firestore_proto_api';
+import { ApiClientObjectMap } from '../../../src/protos/firestore_proto_api';
 import { Connection, Stream } from '../../../src/remote/connection';
 import { Datastore, newDatastore } from '../../../src/remote/datastore';
 import { WriteRequest } from '../../../src/remote/persistent_stream';
@@ -259,7 +260,12 @@ export class MockConnection implements Connection {
    * Tracks the currently active watch targets as detected by the mock watch
    * stream, as a mapping from target ID to query Target.
    */
-  activeTargets: { [targetId: number]: api.Target } = {};
+  activeTargets: {
+    [targetId: number]: {
+      target: api.Target;
+      labels?: ApiClientObjectMap<string>;
+    };
+  } = {};
 
   /** A Deferred that is resolved once watch opens. */
   watchOpen = new Deferred<void>();
@@ -398,7 +404,10 @@ export class MockConnection implements Connection {
           ++this.watchStreamRequestCount;
           if (request.addTarget) {
             const targetId = request.addTarget.targetId!;
-            this.activeTargets[targetId] = request.addTarget;
+            this.activeTargets[targetId] = {
+              target: request.addTarget,
+              labels: request.labels
+            };
           } else if (request.removeTarget) {
             delete this.activeTargets[request.removeTarget];
           } else {
diff --git a/packages/firestore/test/unit/specs/spec_test_runner.ts b/packages/firestore/test/unit/specs/spec_test_runner.ts
@@ -105,6 +105,7 @@ import {
 import { mapCodeFromRpcCode } from '../../../src/remote/rpc_error';
 import {
   JsonProtoSerializer,
+  toListenRequestLabels,
   toMutation,
   toTarget,
   toVersion
@@ -1095,15 +1096,13 @@ abstract class TestRunner {
         undefined,
         'Expected active target not found: ' + JSON.stringify(expected)
       );
-      const actualTarget = actualTargets[targetId];
+      const { target: actualTarget, labels: actualLabels } =
+        actualTargets[targetId];
 
-      // TODO(mcg): populate the purpose of the target once it's possible to
-      // encode that in the spec tests. For now, hard-code that it's a listen
-      // despite the fact that it's not always the right value.
       let targetData = new TargetData(
         queryToTarget(parseQuery(expected.queries[0])),
         targetId,
-        TargetPurpose.Listen,
+        expected.targetPurpose ?? TargetPurpose.Listen,
         ARBITRARY_SEQUENCE_NUMBER
       );
       if (expected.resumeToken && expected.resumeToken !== '') {
@@ -1117,6 +1116,11 @@ abstract class TestRunner {
           version(expected.readTime!)
         );
       }
+
+      const expectedLabels =
+        toListenRequestLabels(this.serializer, targetData) ?? undefined;
+      expect(actualLabels).to.deep.equal(expectedLabels);
+
       const expectedTarget = toTarget(this.serializer, targetData);
       expect(actualTarget.query).to.deep.equal(expectedTarget.query);
       expect(actualTarget.targetId).to.equal(expectedTarget.targetId);
