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 24aa433..edb0f5f 100644 --- a/mockfirestore/document.py +++ b/mockfirestore/document.py @@ -63,7 +63,11 @@ def __init__(self, data: Store, path: List[str], def id(self): return self._path[-1] - def get(self) -> DocumentSnapshot: + @property + def path(self): + return '/'.join(self._path) + + def get(self, **kwargs) -> DocumentSnapshot: return DocumentSnapshot(self, get_by_path(self._data, self._path)) def delete(self): @@ -92,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 c582b36..4cc5310 100644 --- a/tests/test_document_reference.py +++ b/tests/test_document_reference.py @@ -16,6 +16,23 @@ def test_get_document_by_path(self): self.assertEqual({'id': 1}, doc.to_dict()) self.assertEqual('first', doc.id) + def test_get_document_by_path_with_transaction_kwarg(self): + fs = MockFirestore() + fs._data = {'foo': { + 'first': {'id': 1} + }} + doc = fs.document('foo/first').get(transaction='tx') + self.assertEqual({'id': 1}, doc.to_dict()) + self.assertEqual('first', doc.id) + + def test_document_path_property(self): + fs = MockFirestore() + fs._data = {'foo': { + 'first': {'id': 1} + }} + doc = fs.document('foo/first') + self.assertEqual('foo/first', doc.path) + def test_set_document_by_path(self): fs = MockFirestore() fs._data = {} @@ -23,7 +40,7 @@ def test_set_document_by_path(self): fs.document('foo/doc1/bar/doc2').set(doc_content) doc = fs.document('foo/doc1/bar/doc2').get().to_dict() self.assertEqual(doc_content, doc) - + def test_document_get_returnsDocument(self): fs = MockFirestore() fs._data = {'foo': { @@ -72,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': {