firebase-js-sdk/packages/firestore/test/unit/specs/spec_builder.ts

/**
 * @license
 * Copyright 2017 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import { IndexConfiguration } from '../../../src/api/index_configuration';
import { ExpUserDataWriter } from '../../../src/api/user_data_writer';
import {
  ListenOptions,
  ListenerDataSource as Source
} from '../../../src/core/event_manager';
import { FieldFilter, Filter } from '../../../src/core/filter';
import { CorePipeline } from '../../../src/core/pipeline';
import {
  canonifyTargetOrPipeline,
  isPipeline,
  pipelineEq,
  QueryOrPipeline,
  TargetOrPipeline,
  targetOrPipelineEqual,
  toCorePipeline,
  toPipelineStages
} from '../../../src/core/pipeline-util';
import {
  LimitType,
  newQueryForPath,
  Query,
  queryEquals,
  queryToTarget
} from '../../../src/core/query';
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';
import { JsonObject } from '../../../src/model/object_value';
import { ResourcePath } from '../../../src/model/path';
import { BloomFilter as ProtoBloomFilter } from '../../../src/protos/firestore_proto_api';
import {
  isPermanentWriteError,
  mapCodeFromRpcCode,
  mapRpcCodeFromCode
} from '../../../src/remote/rpc_error';
import { debugAssert, fail } from '../../../src/util/assert';
import { TimerId } from '../../../src/util/async_queue';
import { Code } from '../../../src/util/error';
import { forEach } from '../../../src/util/obj';
import { ObjectMap } from '../../../src/util/obj_map';
import { isNullOrUndefined } from '../../../src/util/types';
import { firestore, newTestFirestore } from '../../util/api_helpers';
import { deletedDoc, TestSnapshotVersion } from '../../util/helpers';
import { pipelineFromStages } from '../../util/pipelines';

import { RpcError } from './spec_rpc_error';
import {
  parseQuery,
  PersistenceAction,
  runSpec,
  SpecConfig,
  SpecDocument,
  SpecQuery,
  SpecQueryFilter,
  SpecQueryOrderBy,
  SpecStep,
  SpecWatchFilter,
  SpecWriteAck,
  SpecWriteFailure
} from './spec_test_runner';

const userDataWriter = new ExpUserDataWriter(firestore());

// These types are used in a protected API by SpecBuilder and need to be
// exported.
export interface LimboMap {
  [key: string]: TargetId;
}

export interface ActiveTargetSpec {
  queries: Array<SpecQuery | CorePipeline>;
  targetPurpose?: TargetPurpose;
  resumeToken?: string;
  readTime?: TestSnapshotVersion;
  expectedCount?: number;
}

export interface ActiveTargetMap {
  [targetId: string]: ActiveTargetSpec;
}

export interface ResumeSpec {
  resumeToken?: string;
  readTime?: TestSnapshotVersion;
  expectedCount?: number;
}

/**
 * Tracks the expected memory state of a client (e.g. the expected active watch
 * targets based on userListens(), userUnlistens(), and watchRemoves()
 * as well as the expectActiveTargets() and expectLimboDocs() expectations).
 *
 * Automatically keeping track of the active targets makes writing tests
 * much simpler and the tests much easier to follow.
 *
 * Whenever the map changes, the expected state is automatically encoded in
 * the tests.
 */
export class ClientMemoryState {
  activeTargets: ActiveTargetMap = {};
  queryMapping = new ObjectMap<TargetOrPipeline, TargetId>(
    canonifyTargetOrPipeline,
    targetOrPipelineEqual
  );
  limboMapping: LimboMap = {};

  limboIdGenerator: TargetIdGenerator = TargetIdGenerator.forSyncEngine();
  injectFailures = false;

  constructor() {
    this.reset();
  }

  /** Reset all internal memory state (as done during a client restart). */
  reset(): void {
    this.queryMapping = new ObjectMap<TargetOrPipeline, TargetId>(
      canonifyTargetOrPipeline,
      targetOrPipelineEqual
    );
    this.limboMapping = {};
    this.activeTargets = {};
    this.limboIdGenerator = TargetIdGenerator.forSyncEngine();
  }

  /**
   * Reset the internal limbo mapping (as done during a primary lease failover).
   */
  resetLimboMapping(): void {
    this.limboMapping = {};
  }
}

/**
 * Generates and provides consistent cross-tab target IDs for queries that are
 * active in multiple tabs.
 */
class CachedTargetIdGenerator {
  // TODO(wuandy): rename this to targetMapping.
  private queryMapping = new ObjectMap<TargetOrPipeline, TargetId>(
    canonifyTargetOrPipeline,
    targetOrPipelineEqual
  );
  private targetIdGenerator = TargetIdGenerator.forTargetCache();

  /**
   * Returns a cached target ID for the provided Target, or a new ID if no
   * target ID has ever been assigned.
   */
  next(target: TargetOrPipeline): TargetId {
    if (this.queryMapping.has(target)) {
      return this.queryMapping.get(target)!;
    }
    const targetId = this.targetIdGenerator.next();
    this.queryMapping.set(target, targetId);
    return targetId;
  }

  /** Returns the target ID for a target that is known to exist. */
  cachedId(target: TargetOrPipeline): TargetId {
    if (!this.queryMapping.has(target)) {
      throw new Error("Target ID doesn't exists for target: " + target);
    }

    return this.queryMapping.get(target)!;
  }

  /** Remove the cached target ID for the provided target. */
  purge(target: TargetOrPipeline): void {
    if (!this.queryMapping.has(target)) {
      throw new Error("Target ID doesn't exists for target: " + target);
    }

    this.queryMapping.delete(target);
  }
}
/**
 * Provides a high-level language to construct spec tests that can be exported
 * to the spec JSON format or be run as a spec test directly.
 *
 * Exported JSON tests can be used in other clients without the need to
 * duplicate tests in every client.
 */
export class SpecBuilder {
  protected config: SpecConfig = { useEagerGCForMemory: true, numClients: 1 };
  // currentStep is built up (in particular, expectations can be added to it)
  // until nextStep() is called to append it to steps.
  protected currentStep: SpecStep | null = null;

  private steps: SpecStep[] = [];

  private queryIdGenerator = new CachedTargetIdGenerator();

  private readonly currentClientState: ClientMemoryState =
    new ClientMemoryState();

  // Accessor function that can be overridden to return a different
  // `ClientMemoryState`.
  protected get clientState(): ClientMemoryState {
    return this.currentClientState;
  }

  private get limboIdGenerator(): TargetIdGenerator {
    return this.clientState.limboIdGenerator;
  }

  private get queryMapping(): ObjectMap<TargetOrPipeline, TargetId> {
    return this.clientState.queryMapping;
  }

  private get limboMapping(): LimboMap {
    return this.clientState.limboMapping;
  }

  private get activeTargets(): ActiveTargetMap {
    return this.clientState.activeTargets;
  }

  private get injectFailures(): boolean {
    return this.clientState.injectFailures;
  }

  private set injectFailures(injectFailures: boolean) {
    this.clientState.injectFailures = injectFailures;
  }

  /**
   * Exports the spec steps as a JSON object that be used in the spec runner.
   */
  toJSON(): { config: SpecConfig; steps: SpecStep[] } {
    this.nextStep();
    return { config: this.config, steps: this.steps };
  }

  /**
   * Run the spec as a test. If persistence is available it will run it with and
   * without persistence enabled.
   */
  runAsTest(
    name: string,
    tags: string[],
    usePersistence: boolean,
    convertToPipeline: boolean
  ): Promise<void> {
    this.nextStep();
    this.config.convertToPipeline = convertToPipeline;
    return runSpec(name, tags, usePersistence, this.config, this.steps);
  }

  // Ensures manual LRU GC for both memory and indexeddb persistence.
  // In spec tests, GC is always manually triggered via triggerLruGC().
  ensureManualLruGC(): this {
    debugAssert(
      !this.currentStep,
      'ensureManualLruGC() must be called before all spec steps.'
    );
    this.config.useEagerGCForMemory = false;
    return this;
  }

  /**
   * By default the spec test runner does not allow removal of a non-active target,
   * but this scenario is possible in the real world due to a client / server race,
   * and therefore must be allowed for some spec tests.
   */
  allowUnlistedTargetRemoval(): this {
    this.config.allowUnlistedTargetRemoval = true;
    return this;
  }

  withMaxConcurrentLimboResolutions(value?: number): this {
    this.config.maxConcurrentLimboResolutions = value;
    return this;
  }

  private addUserListenStep(
    query: QueryOrPipeline,
    resume?: ResumeSpec,
    options?: ListenOptions
  ): void {
    this.nextStep();

    const target = isPipeline(query) ? query : queryToTarget(query);
    let targetId: TargetId = 0;

    if (this.injectFailures) {
      // Return a `userListens()` step but don't advance the target IDs.
      this.currentStep = {
        userListen: {
          targetId,
          query: isPipeline(query) ? query : SpecBuilder.queryToSpec(query),
          options
        }
      };
    } else {
      if (this.queryMapping.has(target)) {
        targetId = this.queryMapping.get(target)!;
      } else {
        targetId = this.queryIdGenerator.next(target);
      }

      this.queryMapping.set(target, targetId);

      if (options?.source !== Source.Cache) {
        // Active targets are created if listener is not sourcing from cache
        this.addQueryToActiveTargets(targetId, query, resume);
      }

      this.currentStep = {
        userListen: {
          targetId,
          query: isPipeline(query) ? query : SpecBuilder.queryToSpec(query),
          options
        },
        expectedState: { activeTargets: { ...this.activeTargets } }
      };
    }
  }

  userListens(query: QueryOrPipeline, resume?: ResumeSpec): this {
    this.addUserListenStep(query, resume);
    return this;
  }

  /** Listen to query using the same options as executing a getDoc or getDocs */
  userListensForGet(query: Query, resume?: ResumeSpec): this {
    this.addUserListenStep(query, resume, {
      includeMetadataChanges: true,
      waitForSyncWhenOnline: true
    });
    return this;
  }

  userListensToCache(query: QueryOrPipeline, resume?: ResumeSpec): this {
    this.addUserListenStep(query, resume, { source: Source.Cache });
    return this;
  }

  /**
   * Registers a previously active target with the test expectations after a
   * stream disconnect.
   */
  restoreListen(
    query: QueryOrPipeline,
    resumeToken: string,
    expectedCount?: number
  ): this {
    const targetId = this.queryMapping.get(
      isPipeline(query) ? query : queryToTarget(query)
    );

    if (isNullOrUndefined(targetId)) {
      throw new Error("Can't restore an unknown query: " + query);
    }

    this.addQueryToActiveTargets(targetId!, query, {
      resumeToken,
      expectedCount
    });

    const currentStep = this.currentStep!;
    currentStep.expectedState = currentStep.expectedState || {};
    currentStep.expectedState.activeTargets = { ...this.activeTargets };
    return this;
  }

  userUnlistens(
    query: QueryOrPipeline,
    shouldRemoveWatchTarget: boolean = true
  ): this {
    this.nextStep();
    const target = isPipeline(query) ? query : queryToTarget(query);
    if (!this.queryMapping.has(target)) {
      throw new Error('Unlistening to query not listened to: ' + query);
    }
    const targetId = this.queryMapping.get(target)!;
    if (shouldRemoveWatchTarget) {
      this.removeQueryFromActiveTargets(query, targetId);
    }

    if (this.config.useEagerGCForMemory && !this.activeTargets[targetId]) {
      this.queryMapping.delete(target);
      this.queryIdGenerator.purge(target);
    }

    this.currentStep = {
      userUnlisten: [
        targetId,
        isPipeline(query) ? query : SpecBuilder.queryToSpec(query)
      ],
      expectedState: { activeTargets: { ...this.activeTargets } }
    };
    return this;
  }

  userUnlistensToCache(query: QueryOrPipeline): this {
    // Listener sourced from cache do not need to close watch stream.
    return this.userUnlistens(query, /** shouldRemoveWatchTarget= */ false);
  }

  userSets(key: string, value: JsonObject<unknown>): this {
    this.nextStep();
    this.currentStep = {
      userSet: [key, value]
    };
    return this;
  }

  userPatches(key: string, value: JsonObject<unknown>): this {
    this.nextStep();
    this.currentStep = {
      userPatch: [key, value]
    };
    return this;
  }

  userDeletes(key: string): this {
    this.nextStep();
    this.currentStep = {
      userDelete: key
    };
    return this;
  }

  userAddsSnapshotsInSyncListener(): this {
    this.nextStep();
    this.currentStep = {
      addSnapshotsInSyncListener: true
    };
    return this;
  }

  userRemovesSnapshotsInSyncListener(): this {
    this.nextStep();
    this.currentStep = {
      removeSnapshotsInSyncListener: true
    };
    return this;
  }

  loadBundle(bundleContent: string): this {
    this.nextStep();
    this.currentStep = {
      loadBundle: bundleContent
    };
    // Loading a bundle implicitly creates a new target. We advance the `queryIdGenerator` to match.
    this.queryIdGenerator.next(
      queryToTarget(newQueryForPath(ResourcePath.emptyPath()))
    );
    return this;
  }

  setIndexConfiguration(
    jsonOrConfiguration: string | IndexConfiguration
  ): this {
    this.nextStep();
    this.currentStep = {
      setIndexConfiguration: jsonOrConfiguration
    };
    return this;
  }

  // PORTING NOTE: Only used by web multi-tab tests.
  becomeHidden(): this {
    this.nextStep();
    this.currentStep = {
      applyClientState: { visibility: 'hidden' }
    };
    return this;
  }

  // PORTING NOTE: Only used by web multi-tab tests.
  becomeVisible(): this {
    this.nextStep();
    this.currentStep = {
      applyClientState: { visibility: 'visible' }
    };
    return this;
  }

  runTimer(timerId: TimerId): this {
    this.nextStep();
    this.currentStep = { runTimer: timerId };
    return this;
  }

  changeUser(uid: string | null): this {
    this.nextStep();
    this.currentStep = { changeUser: uid };
    return this;
  }

  disableNetwork(): this {
    this.nextStep();
    this.currentStep = {
      enableNetwork: false,
      expectedState: {
        activeTargets: {},
        activeLimboDocs: [],
        enqueuedLimboDocs: []
      }
    };
    return this;
  }

  enableNetwork(): this {
    this.nextStep();
    this.currentStep = {
      enableNetwork: true
    };
    return this;
  }

  clearPersistence(): this {
    this.nextStep();
    this.currentStep = {
      clearPersistence: true
    };
    return this;
  }

  restart(): this {
    this.nextStep();
    this.currentStep = {
      restart: true,
      expectedState: {
        activeTargets: {},
        activeLimboDocs: [],
        enqueuedLimboDocs: []
      }
    };
    // Reset our mappings / target ids since all existing listens will be
    // forgotten.
    this.clientState.reset();
    return this;
  }

  shutdown(): this {
    this.nextStep();
    this.currentStep = {
      shutdown: true,
      expectedState: {
        activeTargets: {},
        activeLimboDocs: [],
        enqueuedLimboDocs: []
      }
    };
    // Reset our mappings / target ids since all existing listens will be
    // forgotten.
    this.clientState.reset();
    return this;
  }

  /**
   * Fails the specified database transaction until `recoverDatabase()` is
   * called.
   */
  failDatabaseTransactions(...actions: PersistenceAction[]): this {
    this.nextStep();
    this.injectFailures = true;
    this.currentStep = {
      failDatabase: actions
    };
    return this;
  }

  /** Stops failing database operations. */
  recoverDatabase(): this {
    this.nextStep();
    this.injectFailures = false;
    this.currentStep = {
      failDatabase: false
    };
    return this;
  }

  expectIsShutdown(): this {
    this.assertStep('Active target expectation requires previous step');
    const currentStep = this.currentStep!;
    currentStep.expectedState = currentStep.expectedState || {};
    currentStep.expectedState.isShutdown = true;
    return this;
  }

  /** Expects indexes to exist (in any order) */
  expectIndexes(indexes: FieldIndex[]): this {
    this.assertStep('Indexes expectation requires previous step');
    const currentStep = this.currentStep!;
    currentStep.expectedState = currentStep.expectedState || {};
    currentStep.expectedState.indexes = indexes;
    return this;
  }

  /** Overrides the currently expected set of active targets. */
  expectActiveTargets(
    ...targets: Array<{
      query: Query;
      targetPurpose?: TargetPurpose;
      resumeToken?: string;
      readTime?: TestSnapshotVersion;
      expectedCount?: number;
    }>
  ): this {
    this.assertStep('Active target expectation requires previous step');
    const currentStep = this.currentStep!;
    this.clientState.activeTargets = {};
    targets.forEach(
      ({ query, resumeToken, readTime, expectedCount, targetPurpose }) => {
        this.addQueryToActiveTargets(
          this.getTargetId(query),
          query,
          {
            resumeToken,
            readTime,
            expectedCount
          },
          targetPurpose
        );
      }
    );
    currentStep.expectedState = currentStep.expectedState || {};
    currentStep.expectedState.activeTargets = { ...this.activeTargets };
    return this;
  }

  /**
   * Expects a document to be in limbo. A targetId is assigned if it's not in
   * limbo yet.
   */
  expectLimboDocs(...keys: DocumentKey[]): this {
    this.assertStep('Limbo expectation requires previous step');
    const currentStep = this.currentStep!;

    // Clear any preexisting limbo watch targets, which we'll re-create as
    // necessary from the provided keys below.
    forEach(this.limboMapping, (key, targetId) => {
      delete this.activeTargets[targetId];
    });

    keys.forEach(key => {
      const path = key.path.canonicalString();
      // Create limbo target ID mapping if it was not in limbo yet
      if (!this.limboMapping[path]) {
        this.limboMapping[path] = this.limboIdGenerator.next();
      }
      // Limbo doc queries are always without resume token
      this.addQueryToActiveTargets(
        this.limboMapping[path],
        newQueryForPath(key.path),
        { resumeToken: '' },
        TargetPurpose.LimboResolution
      );
    });

    currentStep.expectedState = currentStep.expectedState || {};
    currentStep.expectedState.activeLimboDocs = keys.map(k =>
      SpecBuilder.keyToSpec(k)
    );
    currentStep.expectedState.activeTargets = { ...this.activeTargets };
    return this;
  }

  /**
   * Expects a document to be in limbo, enqueued for limbo resolution, and
   * therefore *without* an active targetId.
   */
  expectEnqueuedLimboDocs(...keys: DocumentKey[]): this {
    this.assertStep('Limbo expectation requires previous step');
    const currentStep = this.currentStep!;

    currentStep.expectedState = currentStep.expectedState || {};
    currentStep.expectedState.enqueuedLimboDocs = keys.map(k =>
      SpecBuilder.keyToSpec(k)
    );

    return this;
  }

  /**
   * Special helper for limbo documents that acks with either a document or
   * with no document for NoDocument. This is translated into normal watch
   * messages.
   */
  ackLimbo(version: TestSnapshotVersion, doc: Document): this {
    const query = newQueryForPath(doc.key.path);
    this.watchAcks(query);
    if (doc.isFoundDocument()) {
      this.watchSends({ affects: [query] }, doc);
    } else if (doc.isNoDocument()) {
      // Don't send any updates
    } else {
      fail(0xd71c, 'Unknown parameter', { doc });
    }
    this.watchCurrents(query, 'resume-token-' + version);
    this.watchSnapshots(version);
    return this;
  }

  /**
   * Special helper for limbo documents that acks an unlisten for a limbo doc
   * with either a document or with no document for NoDocument. This is
   * translated into normal watch messages.
   */
  watchRemovesLimboTarget(doc: Document): this {
    const query = newQueryForPath(doc.key.path);
    this.watchRemoves(query);
    return this;
  }

  /**
   * Acks a write with a version and optional additional options.
   *
   * expectUserCallback defaults to true if omitted.
   */
  writeAcks(
    doc: string,
    version: TestSnapshotVersion,
    options?: { expectUserCallback?: boolean; keepInQueue?: boolean }
  ): this {
    this.nextStep();
    options = options || {};

    const writeAck: SpecWriteAck = { version };
    if (options.keepInQueue) {
      writeAck.keepInQueue = true;
    }
    this.currentStep = { writeAck };

    if (options.expectUserCallback !== false) {
      return this.expectUserCallbacks({ acknowledged: [doc] });
    } else {
      return this;
    }
  }

  /**
   * Fails a write with an error and optional additional options.
   *
   * expectUserCallback defaults to true if omitted.
   */
  failWrite(
    doc: string,
    error: RpcError,
    options?: { expectUserCallback?: boolean; keepInQueue?: boolean }
  ): this {
    this.nextStep();
    options = options || {};

    // If this is a permanent error, the write is not expected to be sent
    // again.
    const code = mapCodeFromRpcCode(error.code);
    const isPermanentFailure = isPermanentWriteError(code);
    const keepInQueue =
      options.keepInQueue !== undefined
        ? options.keepInQueue
        : !isPermanentFailure;

    const failWrite: SpecWriteFailure = { error };
    if (keepInQueue) {
      failWrite.keepInQueue = true;
    }
    this.currentStep = { failWrite };

    if (options.expectUserCallback !== false) {
      return this.expectUserCallbacks({ rejected: [doc] });
    } else {
      return this;
    }
  }

  // TODO(wuandy): watch* methods should really be dealing with Target, not
  // Query, make this happen.
  watchAcks(query: Query): this {
    this.nextStep();
    this.currentStep = {
      watchAck: [this.getTargetId(query)]
    };
    return this;
  }

  // Technically any target change can contain a resume token, but a CURRENT
  // target change is where it makes the most sense in our tests currently.
  // Eventually we want to make the model more generic so we can add resume
  // tokens in other places.
  // TODO(b/37254270): Handle global resume tokens
  watchCurrents(query: Query, resumeToken: string): this {
    this.nextStep();
    this.currentStep = {
      watchCurrent: [[this.getTargetId(query)], resumeToken]
    };
    return this;
  }

  watchRemoves(query: Query, cause?: RpcError): this {
    this.nextStep();
    this.currentStep = {
      watchRemove: { targetIds: [this.getTargetId(query)], cause }
    };
    if (cause) {
      if (!this.config.allowUnlistedTargetRemoval) {
        delete this.activeTargets[this.getTargetId(query)];
      }
      this.currentStep.expectedState = {
        activeTargets: { ...this.activeTargets }
      };
    }
    return this;
  }

  /**
   * After setting the remote target index, for any watch* step following
   * this step, the simulated watch message will use the Nth remote target ID
   * that was associated with this query
   *
   * .userListens(q1) // remote target id: 1002, index: 0
   * .userUnlistens(q1))
   * .userListens(q1) // remote target id: 1004, index: 1
   * .userUnlistens(q1))
   * .watchUsesTargetIndex(0)
   * .watchAcks(q1) //will ack 1002
   * .watchUsesTargetIndex(1)
   * .watchAcks(q1) //will ack 1004
   */
  watchUsesTargetIndex(remoteTargetIndex: number | 'latest' | 'last'): this {
    this.nextStep();
    this.currentStep = {
      watchUsesTargetIndex:
        remoteTargetIndex === 'last' ? 'latest' : remoteTargetIndex
    };
    return this;
  }

  private watchRemovesWithTargetIndex(
    query: Query,
    remoteTargetIndex: number | undefined,
    cause?: RpcError
  ): this {
    this.nextStep();
    this.currentStep = {
      watchRemove: {
        targetIds: [this.getTargetId(query)],
        cause
      }
    };
    if (cause) {
      if (!this.config.allowUnlistedTargetRemoval) {
        delete this.activeTargets[this.getTargetId(query)];
      }
      this.currentStep.expectedState = {
        activeTargets: { ...this.activeTargets }
      };
    }
    return this;
  }

  watchSends(
    targets: { affects?: Query[]; removed?: Query[] },
    ...docs: Document[]
  ): this {
    this.nextStep();
    const affects =
      targets.affects &&
      targets.affects.map(query => {
        return this.getTargetId(query);
      });
    const removed =
      targets.removed &&
      targets.removed.map(query => {
        return this.getTargetId(query);
      });
    const specDocs: SpecDocument[] = [];
    for (const doc of docs) {
      specDocs.push(SpecBuilder.docToSpec(doc));
    }
    this.currentStep = {
      watchEntity: {
        docs: specDocs,
        targets: affects,
        removedTargets: removed
      }
    };
    return this;
  }

  watchRemovesDoc(key: DocumentKey, ...targets: Query[]): this {
    this.nextStep();
    this.currentStep = {
      watchEntity: {
        key: SpecBuilder.keyToSpec(key),
        removedTargets: targets.map(query => this.getTargetId(query))
      }
    };
    return this;
  }

  watchDeletesDoc(
    key: DocumentKey,
    version: TestSnapshotVersion,
    ...targets: Query[]
  ): this {
    this.nextStep();
    this.currentStep = {
      watchEntity: {
        doc: SpecBuilder.docToSpec(
          deletedDoc(key.path.canonicalString(), version)
        ),
        removedTargets: targets.map(query => this.getTargetId(query))
      }
    };
    return this;
  }

  watchFilters(
    queries: Query[],
    docs: DocumentKey[] = [],
    bloomFilter?: ProtoBloomFilter
  ): this {
    this.nextStep();
    const targetIds = queries.map(query => {
      return this.getTargetId(query);
    });
    const keys = docs.map(key => {
      return key.path.canonicalString();
    });
    const filter: SpecWatchFilter = { targetIds, keys, bloomFilter };
    this.currentStep = {
      watchFilter: filter
    };
    return this;
  }

  watchResets(...queries: Query[]): this {
    this.nextStep();
    const targetIds = queries.map(query => this.getTargetId(query));
    this.currentStep = {
      watchReset: targetIds
    };
    return this;
  }

  watchSnapshots(
    version: TestSnapshotVersion,
    targets?: Query[],
    resumeToken?: string
  ): this {
    this.nextStep();
    const targetIds = targets
      ? targets.map(query => this.getTargetId(query))
      : [];
    this.currentStep = {
      watchSnapshot: { version, targetIds, resumeToken }
    };
    return this;
  }

  watchAcksFull(
    query: Query,
    version: TestSnapshotVersion,
    ...docs: Document[]
  ): this {
    this.watchAcks(query);
    this.watchSends({ affects: [query] }, ...docs);
    this.watchCurrents(query, 'resume-token-' + version);
    this.watchSnapshots(version);
    return this;
  }

  watchStreamCloses(error: Code, opts?: { runBackoffTimer: boolean }): this {
    if (!opts) {
      opts = { runBackoffTimer: true };
    }

    this.nextStep();
    this.currentStep = {
      watchStreamClose: {
        error: {
          code: mapRpcCodeFromCode(error),
          message: 'Simulated Backend Error'
        },
        runBackoffTimer: opts.runBackoffTimer
      }
    };
    return this;
  }

  waitForPendingWrites(): this {
    this.nextStep();
    this.currentStep = {
      waitForPendingWrites: true
    };
    return this;
  }

  expectUserCallbacks(docs: {
    acknowledged?: string[];
    rejected?: string[];
  }): this {
    this.assertStep('Expectations require previous step');
    const currentStep = this.currentStep!;
    currentStep.expectedState = currentStep.expectedState || {};
    currentStep.expectedState.userCallbacks = currentStep.expectedState
      .userCallbacks || { acknowledgedDocs: [], rejectedDocs: [] };

    if (docs.acknowledged) {
      currentStep.expectedState.userCallbacks.acknowledgedDocs.push(
        ...docs.acknowledged
      );
    }

    if (docs.rejected) {
      currentStep.expectedState.userCallbacks.rejectedDocs.push(
        ...docs.rejected
      );
    }

    return this;
  }

  expectEvents(
    query: Query,
    events: {
      fromCache?: boolean;
      hasPendingWrites?: boolean;
      added?: Document[];
      modified?: Document[];
      removed?: Document[];
      metadata?: Document[];
      errorCode?: Code;
    }
  ): this {
    this.assertStep('Expectations require previous step');
    const currentStep = this.currentStep!;
    if (!currentStep.expectedSnapshotEvents) {
      currentStep.expectedSnapshotEvents = [];
    }
    debugAssert(
      !events.errorCode ||
        !(events.added || events.modified || events.removed || events.metadata),
      "Can't provide both error and events"
    );
    currentStep.expectedSnapshotEvents.push({
      query: SpecBuilder.queryToSpec(query),
      added: events.added && events.added.map(SpecBuilder.docToSpec),
      modified: events.modified && events.modified.map(SpecBuilder.docToSpec),
      removed: events.removed && events.removed.map(SpecBuilder.docToSpec),
      metadata: events.metadata && events.metadata.map(SpecBuilder.docToSpec),
      errorCode: mapRpcCodeFromCode(events.errorCode),
      fromCache: events.fromCache || false,
      hasPendingWrites: events.hasPendingWrites || false
    });
    return this;
  }

  private registerQuery(
    query: Query,
    shouldAddWatchTarget: boolean,
    resume?: ResumeSpec
  ): this {
    this.assertStep('Expectations require previous step');

    const target = queryToTarget(query);
    const targetId = this.queryIdGenerator.cachedId(target);
    this.queryMapping.set(target, targetId);

    if (shouldAddWatchTarget) {
      this.addQueryToActiveTargets(targetId, query, resume);
    }
    const currentStep = this.currentStep!;
    currentStep.expectedState = currentStep.expectedState || {};
    currentStep.expectedState.activeTargets = { ...this.activeTargets };
    return this;
  }

  /** Registers a query that is active in another tab. */
  expectListen(query: Query, resume?: ResumeSpec): this {
    return this.registerQuery(query, true, resume);
  }

  /** Registers a query that is listening to cache and active in another tab. */
  expectListenToCache(query: Query, resume?: ResumeSpec): this {
    // Listeners that source from cache would not send watch request.
    return this.registerQuery(query, false, resume);
  }

  removeQuery(query: Query, shouldRemoveWatchTarget: boolean = true): this {
    this.assertStep('Expectations require previous step');

    const target = queryToTarget(query);
    const targetId = this.queryMapping.get(target)!;

    if (shouldRemoveWatchTarget) {
      this.removeQueryFromActiveTargets(query, targetId);
    }

    if (this.config.useEagerGCForMemory && !this.activeTargets[targetId]) {
      this.queryMapping.delete(target);
      this.queryIdGenerator.purge(target);
    }

    const currentStep = this.currentStep!;
    currentStep.expectedState = currentStep.expectedState || {};
    currentStep.expectedState.activeTargets = { ...this.activeTargets };
    return this;
  }

  /** Removes a query that is no longer active in any tab. */
  expectUnlisten(query: Query): this {
    return this.removeQuery(query);
  }

  /** Removes a query that is listening to cache and no longer active in any tab. */
  expectUnlistenToCache(query: Query): this {
    // Listeners that source from cache did not establish watch connection, so no active targets to remove.
    return this.removeQuery(query, false);
  }

  /**
   * Verifies the total number of requests sent to the write backend since test
   * initialization.
   */
  expectWriteStreamRequestCount(num: number): this {
    this.assertStep('Expectations require previous step');
    const currentStep = this.currentStep!;
    currentStep.expectedState = currentStep.expectedState || {};
    currentStep.expectedState.writeStreamRequestCount = num;
    return this;
  }

  /**
   * Verifies the total number of requests sent to the watch backend since test
   * initialization.
   */
  expectWatchStreamRequestCount(num: number): this {
    this.assertStep('Expectations require previous step');
    const currentStep = this.currentStep!;
    currentStep.expectedState = currentStep.expectedState || {};
    currentStep.expectedState.watchStreamRequestCount = num;
    return this;
  }

  expectNumOutstandingWrites(num: number): this {
    this.assertStep('Expectations require previous step');
    const currentStep = this.currentStep!;
    currentStep.expectedState = currentStep.expectedState || {};
    currentStep.expectedState.numOutstandingWrites = num;
    return this;
  }

  expectNumActiveClients(num: number): this {
    this.assertStep('Expectations require previous step');
    const currentStep = this.currentStep!;
    currentStep.expectedState = currentStep.expectedState || {};
    currentStep.expectedState.numActiveClients = num;
    return this;
  }

  expectPrimaryState(isPrimary: boolean): this {
    this.assertStep('Expectations require previous step');
    const currentStep = this.currentStep!;
    currentStep.expectedState = currentStep.expectedState || {};
    currentStep.expectedState.isPrimary = isPrimary;
    return this;
  }

  expectSnapshotsInSyncEvent(count = 1): this {
    this.assertStep('Expectations require previous step');
    const currentStep = this.currentStep!;
    currentStep.expectedSnapshotsInSyncEvents = count;
    return this;
  }

  expectWaitForPendingWritesEvent(count = 1): this {
    this.assertStep('Expectations require previous step');
    const currentStep = this.currentStep!;
    currentStep.expectedWaitForPendingWritesEvents = count;
    return this;
  }

  triggerLruGC(cacheThreshold: number): this {
    this.nextStep();
    this.currentStep = {
      triggerLruGC: cacheThreshold
    };
    return this;
  }

  removeExpectedTargetMapping(query: Query): this {
    const target = queryToTarget(query);
    this.queryMapping.delete(target);
    this.queryIdGenerator.purge(target);
    return this;
  }

  private static queryToSpec(query: Query): SpecQuery {
    // TODO(dimond): full query support
    const spec: SpecQuery = { path: query.path.canonicalString() };
    if (query.collectionGroup !== null) {
      spec.collectionGroup = query.collectionGroup;
    }
    if (query.limit !== null) {
      spec.limit = query.limit;
      spec.limitType =
        query.limitType === LimitType.First ? 'LimitToFirst' : 'LimitToLast';
    }
    if (query.filters) {
      spec.filters = query.filters.map((filter: Filter) => {
        if (filter instanceof FieldFilter) {
          // TODO(dimond): Support non-JSON primitive values?
          return [
            filter.field.canonicalString(),
            filter.op,
            userDataWriter.convertValue(filter.value)
          ] as SpecQueryFilter;
        } else {
          return fail(0x0e51, 'Unknown filter', { filter });
        }
      });
    }
    if (query.explicitOrderBy) {
      spec.orderBys = query.explicitOrderBy.map(orderBy => {
        return [
          orderBy.field.canonicalString(),
          orderBy.dir
        ] as SpecQueryOrderBy;
      });
    }
    return spec;
  }

  private static docToSpec(doc: Document): SpecDocument {
    if (doc.isFoundDocument()) {
      return {
        key: SpecBuilder.keyToSpec(doc.key),
        version: doc.version.toMicroseconds(),
        createTime: doc.createTime.toMicroseconds(),
        value: userDataWriter.convertValue(
          doc.data.value
        ) as JsonObject<unknown>,
        options: {
          hasLocalMutations: doc.hasLocalMutations,
          hasCommittedMutations: doc.hasCommittedMutations
        }
      };
    } else {
      return {
        key: SpecBuilder.keyToSpec(doc.key),
        version: doc.version.toMicroseconds(),
        createTime: doc.createTime.toMicroseconds(),
        value: null
      };
    }
  }

  private static keyToSpec(key: DocumentKey): string {
    return key.path.canonicalString();
  }

  protected nextStep(): void {
    if (this.currentStep !== null) {
      this.steps.push(this.currentStep);
      this.currentStep = null;
    }
  }

  /**
   * Add the specified `Query` under give active targe id. If it is already
   * added, this is a no-op.
   */
  private addQueryToActiveTargets(
    targetId: number,
    query: QueryOrPipeline,
    resume: ResumeSpec = {},
    targetPurpose?: TargetPurpose
  ): void {
    if (!(resume?.resumeToken || resume?.readTime) && resume?.expectedCount) {
      fail(
        0xc9a1,
        'Expected count is present without a resume token or read time.'
      );
    }

    if (this.activeTargets[targetId]) {
      const activeQueries = this.activeTargets[targetId].queries;
      if (
        !activeQueries.some(specQuery =>
          this.specQueryOrPipelineEq(specQuery, query)
        )
      ) {
        // `query` is not added yet.
        this.activeTargets[targetId] = {
          queries: [
            isPipeline(query) ? query : SpecBuilder.queryToSpec(query),
            ...activeQueries
          ],
          targetPurpose,
          resumeToken: resume.resumeToken || '',
          readTime: resume.readTime
        };
      } else {
        this.activeTargets[targetId] = {
          queries: activeQueries,
          targetPurpose,
          resumeToken: resume.resumeToken || '',
          readTime: resume.readTime
        };
      }
    } else {
      this.activeTargets[targetId] = {
        queries: [isPipeline(query) ? query : SpecBuilder.queryToSpec(query)],
        targetPurpose,
        resumeToken: resume.resumeToken || '',
        readTime: resume.readTime
      };
    }
  }

  private specQueryOrPipelineEq(
    spec: SpecQuery | CorePipeline,
    query: QueryOrPipeline
  ): boolean {
    if (isPipeline(query) && spec instanceof CorePipeline) {
      return pipelineEq(spec as CorePipeline, query);
    } else if (!isPipeline(query) && spec instanceof CorePipeline) {
      return pipelineEq(
        spec as CorePipeline,
        toCorePipeline(
          pipelineFromStages(
            toPipelineStages(query as Query, newTestFirestore())
          )
        )
      );
    } else {
      return queryEquals(parseQuery(spec as SpecQuery), query as Query);
    }
  }

  private removeQueryFromActiveTargets(
    query: QueryOrPipeline,
    targetId: number
  ): void {
    const queriesAfterRemoval = this.activeTargets[targetId].queries.filter(
      specQuery => !this.specQueryOrPipelineEq(specQuery, query)
    );
    if (queriesAfterRemoval.length > 0) {
      this.activeTargets[targetId] = {
        queries: queriesAfterRemoval,
        resumeToken: this.activeTargets[targetId].resumeToken,
        expectedCount: this.activeTargets[targetId].expectedCount,
        targetPurpose: this.activeTargets[targetId].targetPurpose
      };
    } else {
      delete this.activeTargets[targetId];
    }
  }

  private assertStep(msg: string): void {
    if (this.currentStep === null) {
      throw new Error('Expected a previous step: ' + msg);
    }
  }

  private getTargetId(query: Query): TargetId {
    const queryTargetId = this.queryMapping.get(queryToTarget(query));
    const limboTargetId = this.limboMapping[query.path.canonicalString()];
    if (queryTargetId && limboTargetId) {
      // TODO(dimond): add support for query for doc and limbo doc at the same
      // time?
      fail(
        0x6e17,
        'Found both query and limbo doc with target ID, not supported yet'
      );
    }
    const targetId = queryTargetId || limboTargetId;
    debugAssert(
      !isNullOrUndefined(targetId),
      'No target ID found for query/limbo doc in spec'
    );
    return targetId;
  }
}

/**
 * SpecBuilder that supports serialized interactions between different clients.
 *
 * Use `client(clientIndex)` to switch between clients.
 */
// PORTING NOTE: Only used by web multi-tab tests.
export class MultiClientSpecBuilder extends SpecBuilder {
  private activeClientIndex = -1;
  private clientStates: ClientMemoryState[] = [];

  protected get clientState(): ClientMemoryState {
    if (!this.clientStates[this.activeClientIndex]) {
      this.clientStates[this.activeClientIndex] = new ClientMemoryState();
    }
    return this.clientStates[this.activeClientIndex];
  }

  client(clientIndex: number): MultiClientSpecBuilder {
    // Since `currentStep` is fully self-contained and does not rely on previous
    // state, we don't need to use a different SpecBuilder instance for each
    // client.
    this.nextStep();
    this.currentStep = {
      drainQueue: true
    };

    this.activeClientIndex = clientIndex;
    this.config.numClients = Math.max(
      this.config.numClients,
      this.activeClientIndex + 1
    );

    return this;
  }

  /**
   * Take the primary lease, even if another client has already obtained the
   * lease.
   */
  stealPrimaryLease(): this {
    this.nextStep();
    this.currentStep = {
      applyClientState: {
        primary: true
      },
      expectedState: {
        isPrimary: true
      }
    };

    // HACK: SyncEngine resets its limbo mapping when it gains the primary
    // lease. The SpecTests need to also clear their mapping, but when we parse
    // the spec tests, we don't know when the primary lease transition happens.
    // It is likely going to happen right after `stealPrimaryLease`, so we are
    // clearing the limbo mapping here.
    this.clientState.resetLimboMapping();

    return this;
  }

  protected nextStep(): void {
    if (this.currentStep !== null) {
      this.currentStep.clientIndex = this.activeClientIndex;
    }
    super.nextStep();
  }
}

/** Starts a new single-client SpecTest. */
export function spec(): SpecBuilder {
  return new SpecBuilder();
}

/** Starts a new multi-client SpecTest. */
// PORTING NOTE: Only used by web multi-tab tests.
export function client(num: number): MultiClientSpecBuilder {
  const specBuilder = new MultiClientSpecBuilder();
  specBuilder.ensureManualLruGC();
  return specBuilder.client(num);
}