1 files changed, 53 insertions, 1 deletions
diff --git a/Firestore/Example/Tests/Util/FSTHelpers.h b/Firestore/Example/Tests/Util/FSTHelpers.h
index ccc01ca..7946c06 100644
--- a/Firestore/Example/Tests/Util/FSTHelpers.h
+++ b/Firestore/Example/Tests/Util/FSTHelpers.h
@@ -16,11 +16,11 @@
 
 #import <Foundation/Foundation.h>
 
-#include <map>
 #include <vector>
 
 #import "Firestore/Source/Core/FSTTypes.h"
 #import "Firestore/Source/Model/FSTDocumentDictionary.h"
+#import "Firestore/Source/Remote/FSTRemoteEvent.h"
 
 #include "Firestore/core/src/firebase/firestore/model/field_path.h"
 #include "Firestore/core/src/firebase/firestore/model/field_value.h"
@@ -148,6 +148,42 @@ inline NSString *FSTRemoveExceptionPrefix(NSString *exception) {
     XCTAssertTrue(didThrow, ##__VA_ARGS__);                             \
   } while (0)
 
+/**
+ * An implementation of FSTTargetMetadataProvider that provides controlled access to the
+ * `FSTTargetMetadataProvider` callbacks. Any target accessed via these callbacks must be
+ * registered beforehand via the factory methods or via `setSyncedKeys:forQueryData:`.
+ */
+@interface FSTTestTargetMetadataProvider : NSObject <FSTTargetMetadataProvider>
+
+/**
+ * Creates an FSTTestTargetMetadataProvider that behaves as if there's an established listen for
+ * each of the given targets, where each target has previously seen query results containing just
+ * the given documentKey.
+ *
+ * Internally this means that the `remoteKeysForTarget` callback for these targets will return just
+ * the documentKey and that the provided targets will be returned as active from the
+ * `queryDataForTarget` target.
+ */
++ (instancetype)providerWithSingleResultForKey:(firebase::firestore::model::DocumentKey)documentKey
+                                       targets:(NSArray<FSTBoxedTargetID *> *)targets;
+
+/**
+ * Creates an FSTTestTargetMetadataProvider that behaves as if there's an established listen for
+ * each of the given targets, where each target has not seen any previous document.
+ *
+ * Internally this means that the `remoteKeysForTarget` callback for these targets will return an
+ * empty set of document keys and that the provided targets will be returned as active from the
+ * `queryDataForTarget` target.
+ */
++ (instancetype)providerWithEmptyResultForKey:(firebase::firestore::model::DocumentKey)documentKey
+                                      targets:(NSArray<FSTBoxedTargetID *> *)targets;
+
+/** Sets or replaces the local state for the provided query data. */
+- (void)setSyncedKeys:(firebase::firestore::model::DocumentKeySet)keys
+         forQueryData:(FSTQueryData *)queryData;
+
+@end
+
 /** Creates a new FIRTimestamp from components. Note that year, month, and day are all one-based. */
 FIRTimestamp *FSTTestTimestamp(int year, int month, int day, int hour, int minute, int second);
 
@@ -250,6 +286,9 @@ FSTDeleteMutation *FSTTestDeleteMutation(NSString *path);
 /** Converts a list of documents to a sorted map. */
 FSTMaybeDocumentDictionary *FSTTestDocUpdates(NSArray<FSTMaybeDocument *> *docs);
 
+/** Creates a remote event that inserts a new document. */
+FSTRemoteEvent *FSTTestAddedRemoteEvent(FSTMaybeDocument *doc, NSArray<NSNumber *> *addedToTargets);
+
 /** Creates a remote event with changes to a document. */
 FSTRemoteEvent *FSTTestUpdateRemoteEvent(FSTMaybeDocument *doc,
                                          NSArray<NSNumber *> *updatedInTargets,
@@ -260,6 +299,19 @@ FSTLocalViewChanges *FSTTestViewChanges(FSTQuery *query,
                                         NSArray<NSString *> *addedKeys,
                                         NSArray<NSString *> *removedKeys);
 
+/** Creates a test target change that acks all 'docs' and  marks the target as CURRENT  */
+FSTTargetChange *FSTTestTargetChangeAckDocuments(firebase::firestore::model::DocumentKeySet docs);
+
+/** Creates a test target change that marks the target as CURRENT  */
+FSTTargetChange *FSTTestTargetChangeMarkCurrent();
+
+/** Creates a test target change. */
+FSTTargetChange *FSTTestTargetChange(firebase::firestore::model::DocumentKeySet added,
+                                     firebase::firestore::model::DocumentKeySet modified,
+                                     firebase::firestore::model::DocumentKeySet removed,
+                                     NSData *resumeToken,
+                                     BOOL current);
+
 /** Creates a resume token to match the given snapshot version. */
 NSData *_Nullable FSTTestResumeTokenFromSnapshotVersion(FSTTestSnapshotVersion watchSnapshot);