1 files changed, 8 insertions, 7 deletions

diff --git a/Firestore/Example/Tests/Util/FSTHelpers.m b/Firestore/Example/Tests/Util/FSTHelpers.m
index f01bddb..f2b3605 100644
--- a/Firestore/Example/Tests/Util/FSTHelpers.m
+++ b/Firestore/Example/Tests/Util/FSTHelpers.m
@@ -24,6 +24,7 @@
 #import "Firestore/Source/Core/FSTSnapshotVersion.h"
 #import "Firestore/Source/Core/FSTTimestamp.h"
 #import "Firestore/Source/Core/FSTView.h"
+#import "Firestore/Source/Core/FSTViewSnapshot.h"
 #import "Firestore/Source/Local/FSTLocalViewChanges.h"
 #import "Firestore/Source/Local/FSTQueryData.h"
 #import "Firestore/Source/Model/FSTDatabaseID.h"
@@ -138,7 +139,7 @@ FSTDocument *FSTTestDoc(NSString *path,
                         FSTTestSnapshotVersion version,
                         NSDictionary<NSString *, id> *data,
                         BOOL hasMutations) {
-  FSTDocumentKey *key = [FSTDocumentKey keyWithPathString:path];
+  FSTDocumentKey *key = FSTTestDocKey(path);
   return [FSTDocument documentWithData:FSTTestObjectValue(data)
                                    key:key
                                version:FSTTestVersion(version)
@@ -146,7 +147,7 @@ FSTDocument *FSTTestDoc(NSString *path,
 }
 
 FSTDeletedDocument *FSTTestDeletedDoc(NSString *path, FSTTestSnapshotVersion version) {
-  FSTDocumentKey *key = [FSTDocumentKey keyWithPathString:path];
+  FSTDocumentKey *key = FSTTestDocKey(path);
   return [FSTDeletedDocument documentWithKey:key version:FSTTestVersion(version)];
 }
 
@@ -214,7 +215,7 @@ FSTSortOrder *FSTTestOrderBy(NSString *field, NSString *direction) {
 }
 
 NSComparator FSTTestDocComparator(NSString *fieldPath) {
-  FSTQuery *query = [[FSTQuery queryWithPath:[FSTResourcePath pathWithSegments:@[ @"docs" ]]]
+  FSTQuery *query = [FSTTestQuery(@"docs")
       queryByAddingSortOrder:[FSTSortOrder sortOrderWithFieldPath:FSTTestFieldPath(fieldPath)
                                                         ascending:YES]];
   return [query comparator];
@@ -229,7 +230,7 @@ FSTDocumentSet *FSTTestDocSet(NSComparator comp, NSArray<FSTDocument *> *docs) {
 }
 
 FSTSetMutation *FSTTestSetMutation(NSString *path, NSDictionary<NSString *, id> *values) {
-  return [[FSTSetMutation alloc] initWithKey:[FSTDocumentKey keyWithPathString:path]
+  return [[FSTSetMutation alloc] initWithKey:FSTTestDocKey(path)
                                        value:FSTTestObjectValue(values)
                                 precondition:[FSTPrecondition none]];
 }
@@ -274,7 +275,7 @@ FSTTransformMutation *FSTTestTransformMutation(NSString *path,
 }
 
 FSTDeleteMutation *FSTTestDeleteMutation(NSString *path) {
-  return [[FSTDeleteMutation alloc] initWithKey:[FSTDocumentKey keyWithPathString:path]
+  return [[FSTDeleteMutation alloc] initWithKey:FSTTestDocKey(path)
                                    precondition:[FSTPrecondition none]];
 }
 
@@ -334,12 +335,12 @@ FSTLocalViewChanges *FSTTestViewChanges(FSTQuery *query,
                                         NSArray<NSString *> *removedKeys) {
   FSTDocumentKeySet *added = [FSTDocumentKeySet keySet];
   for (NSString *keyPath in addedKeys) {
-    FSTDocumentKey *key = [FSTDocumentKey keyWithPathString:keyPath];
+    FSTDocumentKey *key = FSTTestDocKey(keyPath);
     added = [added setByAddingObject:key];
   }
   FSTDocumentKeySet *removed = [FSTDocumentKeySet keySet];
   for (NSString *keyPath in removedKeys) {
-    FSTDocumentKey *key = [FSTDocumentKey keyWithPathString:keyPath];
+    FSTDocumentKey *key = FSTTestDocKey(keyPath);
     removed = [removed setByAddingObject:key];
   }
   return [FSTLocalViewChanges changesForQuery:query addedKeys:added removedKeys:removed];