1 files changed, 15 insertions, 17 deletions
diff --git a/Firestore/Example/Tests/Local/FSTQueryCacheTests.mm b/Firestore/Example/Tests/Local/FSTQueryCacheTests.mm
index 429a83a..44b49de 100644
--- a/Firestore/Example/Tests/Local/FSTQueryCacheTests.mm
+++ b/Firestore/Example/Tests/Local/FSTQueryCacheTests.mm
@@ -19,7 +19,6 @@
 #include <set>
 
 #import "Firestore/Source/Core/FSTQuery.h"
-#import "Firestore/Source/Core/FSTSnapshotVersion.h"
 #import "Firestore/Source/Local/FSTEagerGarbageCollector.h"
 #import "Firestore/Source/Local/FSTPersistence.h"
 #import "Firestore/Source/Local/FSTQueryData.h"
@@ -32,6 +31,8 @@
 
 namespace testutil = firebase::firestore::testutil;
 using firebase::firestore::model::DocumentKey;
+using firebase::firestore::model::SnapshotVersion;
+using firebase::firestore::model::DocumentKeySet;
 
 NS_ASSUME_NONNULL_BEGIN
 
@@ -130,9 +131,9 @@ NS_ASSUME_NONNULL_BEGIN
 
     FSTQueryData *result = [self.queryCache queryDataForQuery:_queryRooms];
     XCTAssertNotEqualObjects(queryData2.resumeToken, queryData1.resumeToken);
-    XCTAssertNotEqualObjects(queryData2.snapshotVersion, queryData1.snapshotVersion);
+    XCTAssertNotEqual(queryData2.snapshotVersion, queryData1.snapshotVersion);
     XCTAssertEqualObjects(result.resumeToken, queryData2.resumeToken);
-    XCTAssertEqualObjects(result.snapshotVersion, queryData2.snapshotVersion);
+    XCTAssertEqual(result.snapshotVersion, queryData2.snapshotVersion);
   });
 }
 
@@ -278,12 +279,12 @@ NS_ASSUME_NONNULL_BEGIN
     [self addMatchingKey:key2 forTargetID:1];
     [self addMatchingKey:key3 forTargetID:2];
 
-    FSTAssertEqualSets([self.queryCache matchingKeysForTargetID:1], (@[ key1, key2 ]));
-    FSTAssertEqualSets([self.queryCache matchingKeysForTargetID:2], @[ key3 ]);
+    XCTAssertEqual([self.queryCache matchingKeysForTargetID:1], (DocumentKeySet{key1, key2}));
+    XCTAssertEqual([self.queryCache matchingKeysForTargetID:2], (DocumentKeySet{key3}));
 
     [self addMatchingKey:key1 forTargetID:2];
-    FSTAssertEqualSets([self.queryCache matchingKeysForTargetID:1], (@[ key1, key2 ]));
-    FSTAssertEqualSets([self.queryCache matchingKeysForTargetID:2], (@[ key1, key3 ]));
+    XCTAssertEqual([self.queryCache matchingKeysForTargetID:1], (DocumentKeySet{key1, key2}));
+    XCTAssertEqual([self.queryCache matchingKeysForTargetID:2], (DocumentKeySet{key1, key3}));
   });
 }
 
@@ -386,19 +387,18 @@ NS_ASSUME_NONNULL_BEGIN
   if ([self isTestBaseClass]) return;
 
   self.persistence.run("testLastRemoteSnapshotVersion", [&]() {
-    XCTAssertEqualObjects([self.queryCache lastRemoteSnapshotVersion],
-                          [FSTSnapshotVersion noVersion]);
+    XCTAssertEqual([self.queryCache lastRemoteSnapshotVersion], SnapshotVersion::None());
 
     // Can set the snapshot version.
-    [self.queryCache setLastRemoteSnapshotVersion:FSTTestVersion(42)];
-    XCTAssertEqualObjects([self.queryCache lastRemoteSnapshotVersion], FSTTestVersion(42));
+    [self.queryCache setLastRemoteSnapshotVersion:testutil::Version(42)];
+    XCTAssertEqual([self.queryCache lastRemoteSnapshotVersion], testutil::Version(42));
   });
 
   // Snapshot version persists restarts.
   self.queryCache = [self.persistence queryCache];
   self.persistence.run("testLastRemoteSnapshotVersion restart", [&]() {
     [self.queryCache start];
-    XCTAssertEqualObjects([self.queryCache lastRemoteSnapshotVersion], FSTTestVersion(42));
+    XCTAssertEqual([self.queryCache lastRemoteSnapshotVersion], testutil::Version(42));
   });
 }
 
@@ -424,19 +424,17 @@ NS_ASSUME_NONNULL_BEGIN
                                     targetID:targetID
                         listenSequenceNumber:sequenceNumber
                                      purpose:FSTQueryPurposeListen
-                             snapshotVersion:FSTTestVersion(version)
+                             snapshotVersion:testutil::Version(version)
                                  resumeToken:resumeToken];
 }
 
 - (void)addMatchingKey:(const DocumentKey &)key forTargetID:(FSTTargetID)targetID {
-  FSTDocumentKeySet *keys = [FSTDocumentKeySet keySet];
-  keys = [keys setByAddingObject:key];
+  DocumentKeySet keys{key};
   [self.queryCache addMatchingKeys:keys forTargetID:targetID];
 }
 
 - (void)removeMatchingKey:(const DocumentKey &)key forTargetID:(FSTTargetID)targetID {
-  FSTDocumentKeySet *keys = [FSTDocumentKeySet keySet];
-  keys = [keys setByAddingObject:key];
+  DocumentKeySet keys{key};
   [self.queryCache removeMatchingKeys:keys forTargetID:targetID];
 }