2 files changed, 40 insertions, 11 deletions
diff --git a/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m b/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m
index d22552a..47ba95b 100644
--- a/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m
+++ b/Firestore/Example/Tests/Integration/API/FIRDatabaseTests.m
@@ -680,6 +680,40 @@
   XCTAssertEqual(q.firestore, self.db);
 }
 
+- (void)testDocumentReferenceEquality {
+  FIRFirestore *firestore = self.db;
+  FIRDocumentReference *docRef = [firestore documentWithPath:@"foo/bar"];
+  XCTAssertEqualObjects([firestore documentWithPath:@"foo/bar"], docRef);
+  XCTAssertEqualObjects([docRef collectionWithPath:@"blah"].parent, docRef);
+
+  XCTAssertNotEqualObjects([firestore documentWithPath:@"foo/BAR"], docRef);
+
+  FIRFirestore *otherFirestore = [self firestore];
+  XCTAssertNotEqualObjects([otherFirestore documentWithPath:@"foo/bar"], docRef);
+}
+
+- (void)testQueryReferenceEquality {
+  FIRFirestore *firestore = self.db;
+  FIRQuery *query =
+      [[[firestore collectionWithPath:@"foo"] queryOrderedByField:@"bar"] queryWhereField:@"baz"
+                                                                                isEqualTo:@42];
+  FIRQuery *query2 =
+      [[[firestore collectionWithPath:@"foo"] queryOrderedByField:@"bar"] queryWhereField:@"baz"
+                                                                                isEqualTo:@42];
+  XCTAssertEqualObjects(query, query2);
+
+  FIRQuery *query3 =
+      [[[firestore collectionWithPath:@"foo"] queryOrderedByField:@"BAR"] queryWhereField:@"baz"
+                                                                                isEqualTo:@42];
+  XCTAssertNotEqualObjects(query, query3);
+
+  FIRFirestore *otherFirestore = [self firestore];
+  FIRQuery *query4 = [[[otherFirestore collectionWithPath:@"foo"] queryOrderedByField:@"bar"]
+      queryWhereField:@"baz"
+            isEqualTo:@42];
+  XCTAssertNotEqualObjects(query, query4);
+}
+
 - (void)testCanTraverseCollectionsAndDocuments {
   NSString *expected = @"a/b/c/d";
   // doc path from root Firestore.
diff --git a/Firestore/Example/Tests/Integration/API/FIRTypeTests.m b/Firestore/Example/Tests/Integration/API/FIRTypeTests.m
index f3021dd..4e2c76e 100644
--- a/Firestore/Example/Tests/Integration/API/FIRTypeTests.m
+++ b/Firestore/Example/Tests/Integration/API/FIRTypeTests.m
@@ -62,18 +62,13 @@
 }
 
 - (void)testCanReadAndWriteDocumentReferences {
-  // We can't use assertSuccessfulRoundtrip since FIRDocumentReference doesn't implement isEqual.
-  FIRDocumentReference *docRef = [self.db documentWithPath:@"rooms/eros"];
-  id data = @{ @"a" : @42, @"ref" : docRef };
-  [self writeDocumentRef:docRef data:data];
-
-  FIRDocumentSnapshot *readDoc = [self readDocumentForRef:docRef];
-  XCTAssertTrue(readDoc.exists);
+  FIRDocumentReference *docRef = [self documentRef];
+  [self assertSuccessfulRoundtrip:@{ @"a" : @42, @"ref" : docRef }];
+}
 
-  XCTAssertEqualObjects(readDoc[@"a"], data[@"a"]);
-  FIRDocumentReference *readDocRef = readDoc[@"ref"];
-  XCTAssertTrue([readDocRef isKindOfClass:[FIRDocumentReference class]]);
-  XCTAssertEqualObjects(readDocRef.path, docRef.path);
+- (void)testCanReadAndWriteDocumentReferencesInArrays {
+  FIRDocumentReference *docRef = [self documentRef];
+  [self assertSuccessfulRoundtrip:@{ @"a" : @42, @"refs" : @[ docRef ] }];
 }
 
 @end