diff --git a/mockfirestore/query.py b/mockfirestore/query.py index a8641d7..6cac8fa 100644 --- a/mockfirestore/query.py +++ b/mockfirestore/query.py @@ -71,7 +71,7 @@ def stream(self, transaction=None) -> Iterator[DocumentSnapshot]: for key, direction in reversed(self.orders): # Process least significant first doc_snapshots = sorted( doc_snapshots, - key=lambda doc: doc.to_dict()[key], + key=lambda doc: doc.get(key), reverse=direction == "DESCENDING", ) if self._start_at: diff --git a/tests/test_collection_reference.py b/tests/test_collection_reference.py index 9edc4c2..8dba919 100644 --- a/tests/test_collection_reference.py +++ b/tests/test_collection_reference.py @@ -547,3 +547,31 @@ def test_collection_useDocumentIdKwarg(self): }} doc = fs.collection('foo').document(document_id='first').get() self.assertEqual({'id': 1}, doc.to_dict()) + + def test_collection_orderBy_with_nested_field(self): + fs = MockFirestore() + fs._data = {'foo': { + 'first': {'user': {'name': 'Charlie', 'age': 30}}, + 'second': {'user': {'name': 'Alice', 'age': 25}}, + 'third': {'user': {'name': 'Bob', 'age': 35}} + }} + # Test ordering by nested field using dot notation + docs = list(fs.collection('foo').order_by('user.name').stream()) + self.assertEqual(3, len(docs)) + self.assertEqual('Alice', docs[0].get('user.name')) + self.assertEqual('Bob', docs[1].get('user.name')) + self.assertEqual('Charlie', docs[2].get('user.name')) + + # Test ordering by another nested field + docs = list(fs.collection('foo').order_by('user.age').stream()) + self.assertEqual(3, len(docs)) + self.assertEqual(25, docs[0].get('user.age')) + self.assertEqual(30, docs[1].get('user.age')) + self.assertEqual(35, docs[2].get('user.age')) + + # Test ordering in descending order + docs = list(fs.collection('foo').order_by('user.age', direction='DESCENDING').stream()) + self.assertEqual(3, len(docs)) + self.assertEqual(35, docs[0].get('user.age')) + self.assertEqual(30, docs[1].get('user.age')) + self.assertEqual(25, docs[2].get('user.age'))