From 5a917af7ec2fd474a97b5504c5a58982ce4effe9 Mon Sep 17 00:00:00 2001
From: Ben Schreck <ben@benschreck.com>
Date: Mon, 5 Feb 2024 15:39:18 -0800
Subject: [PATCH] add list collections to doc reference

---
 mockfirestore/collection.py      |  4 ++++
 mockfirestore/document.py        | 10 ++++++++++
 tests/test_document_reference.py | 17 +++++++++++++++++
 3 files changed, 31 insertions(+)

diff --git a/mockfirestore/collection.py b/mockfirestore/collection.py
index 431c074..0b3f9d7 100644
--- a/mockfirestore/collection.py
+++ b/mockfirestore/collection.py
@@ -14,6 +14,10 @@ def __init__(self, data: Store, path: List[str],
         self._path = path
         self.parent = parent
 
+    @property
+    def id(self):
+        return self._path[-1]
+
     def document(self, document_id: Optional[str] = None) -> DocumentReference:
         collection = get_by_path(self._data, self._path)
         if document_id is None:
diff --git a/mockfirestore/document.py b/mockfirestore/document.py
index 76f428b..edb0f5f 100644
--- a/mockfirestore/document.py
+++ b/mockfirestore/document.py
@@ -96,3 +96,13 @@ def collection(self, name) -> 'CollectionReference':
         if name not in document:
             set_by_path(self._data, new_path, {})
         return CollectionReference(self._data, new_path, parent=self)
+
+    def collections(self) -> list['CollectionReference']:
+        from mockfirestore.collection import CollectionReference
+        document = get_by_path(self._data, self._path)
+        _collections = []
+        for collname, values in document.items():
+            if isinstance(values, dict):
+                new_path = self._path + [collname]
+                _collections.append(CollectionReference(self._data, new_path, parent=self))
+        return _collections
diff --git a/tests/test_document_reference.py b/tests/test_document_reference.py
index 6bfe225..4cc5310 100644
--- a/tests/test_document_reference.py
+++ b/tests/test_document_reference.py
@@ -89,6 +89,23 @@ def test_get_nestedDocument(self):
 
         self.assertEqual({'id': 1.1}, doc)
 
+    def test_get_nestedCollections(self):
+        fs = MockFirestore()
+        fs._data = {'top_collection': {
+            'top_document': {
+                'id': 1,
+                'nested_collection': {
+                    'nested_document': {'id': 1.1}
+                }
+            }
+        }}
+        collections = fs.collection('top_collection')\
+            .document('top_document')\
+            .collections()
+
+        self.assertEqual(1, len(collections))
+        self.assertEqual(collections[0].id, 'nested_collection')
+
     def test_get_nestedDocument_documentDoesNotExist(self):
         fs = MockFirestore()
         fs._data = {'top_collection': {