Repository: mdowds/python-mock-firestore

Issue:
DocumentSnapshot.id missing
The Google Firestore client has an id property available on DocumentSnapshot.  id seems to be set regardless of whether the document exists or not.

`id = client.collection('test').document('zzz').get().id`

id should be `'zzz'` in this example, but with mock-firestore, this raises `AttributeError: 'DocumentSnapshot' object has no attribute 'id'`.

Comments:

Patch:
diff --git a/mockfirestore/__init__.py b/mockfirestore/__init__.py
index 5636753..a7f18de 100644
--- a/mockfirestore/__init__.py
+++ b/mockfirestore/__init__.py
@@ -12,3 +12,4 @@ from mockfirestore.document import DocumentSnapshot, DocumentReference
 from mockfirestore.collection import CollectionReference
 from mockfirestore.query import Query
 from mockfirestore._helpers import Timestamp
+from mockfirestore.transaction import Transaction
diff --git a/mockfirestore/client.py b/mockfirestore/client.py
index c66533a..50bbbb4 100644
--- a/mockfirestore/client.py
+++ b/mockfirestore/client.py
@@ -1,16 +1,7 @@
+from typing import Iterable
 from mockfirestore.collection import CollectionReference
-
-
-
-
-
-
-
-
-
-
-
-
+from mockfirestore.document import DocumentReference, DocumentSnapshot
+from mockfirestore.transaction import Transaction
 
 
 class MockFirestore:
@@ -26,5 +17,13 @@ class MockFirestore:
     def reset(self):
         self._data = {}
 
+    def get_all(self, references: Iterable[DocumentReference],
+                field_paths=None,
+                transaction=None) -> Iterable[DocumentSnapshot]:
+        for doc_ref in set(references):
+            yield doc_ref.get()
+
+    def transaction(self, **kwargs) -> Transaction:
+        return Transaction(self, **kwargs)
 
 
diff --git a/mockfirestore/document.py b/mockfirestore/document.py
index ccd3ba7..657143c 100644
--- a/mockfirestore/document.py
+++ b/mockfirestore/document.py
@@ -27,6 +27,15 @@ class DocumentSnapshot:
         timestamp = Timestamp.from_now()
         return timestamp
 
+    @property
+    def update_time(self) -> Timestamp:
+        return self.create_time
+
+    @property
+    def read_time(self) -> Timestamp:
+        timestamp = Timestamp.from_now()
+        return timestamp
+
     def get(self, field_path: str) -> Any:
         if not self.exists:
             return None
diff --git a/mockfirestore/query.py b/mockfirestore/query.py
index 2e94f07..7f32e25 100644
--- a/mockfirestore/query.py
+++ b/mockfirestore/query.py
@@ -24,7 +24,7 @@ class Query:
             for field_filter in field_filters:
                 self._add_field_filter(*field_filter)
 
-    def stream(self) -> Iterator[DocumentSnapshot]:
+    def stream(self, transaction=None) -> Iterator[DocumentSnapshot]:
         doc_snapshots = self.parent.stream()
 
         for field, compare, value in self._field_filters:
diff --git a/mockfirestore/transaction.py b/mockfirestore/transaction.py
new file mode 100644
index 0000000..7f06d2d
--- /dev/null
+++ b/mockfirestore/transaction.py
@@ -0,0 +1,119 @@
+from functools import partial
+import random
+from typing import Iterable, Callable
+from mockfirestore._helpers import generate_random_string, Timestamp
+from mockfirestore.document import DocumentReference, DocumentSnapshot
+from mockfirestore.query import Query
+
+MAX_ATTEMPTS = 5
+_MISSING_ID_TEMPLATE = "The transaction has no transaction ID, so it cannot be {}."
+_CANT_BEGIN = "The transaction has already begun. Current transaction ID: {!r}."
+_CANT_ROLLBACK = _MISSING_ID_TEMPLATE.format("rolled back")
+_CANT_COMMIT = _MISSING_ID_TEMPLATE.format("committed")
+
+
+class WriteResult:
+    def __init__(self):
+        self.update_time = Timestamp.from_now()
+
+
+class Transaction:
+    """
+    This mostly follows the model from
+    https://googleapis.dev/python/firestore/latest/transaction.html
+    """
+    def __init__(self, client,
+                 max_attempts=MAX_ATTEMPTS, read_only=False):
+        self._client = client
+        self._max_attempts = max_attempts
+        self._read_only = read_only
+        self._id = None
+        self._write_ops = []
+        self.write_results = None
+
+    @property
+    def in_progress(self):
+        return self._id is not None
+
+    @property
+    def id(self):
+        return self._id
+
+    def _begin(self, retry_id=None):
+        # generate a random ID to set the transaction as in_progress
+        self._id = generate_random_string()
+
+    def _clean_up(self):
+        self._write_ops.clear()
+        self._id = None
+
+    def _rollback(self):
+        if not self.in_progress:
+            raise ValueError(_CANT_ROLLBACK)
+
+        self._clean_up()
+
+    def _commit(self) -> Iterable[WriteResult]:
+        if not self.in_progress:
+            raise ValueError(_CANT_COMMIT)
+
+        results = []
+        for write_op in self._write_ops:
+            write_op()
+            results.append(WriteResult())
+        self.write_results = results
+        self._clean_up()
+        return results
+
+    def get_all(self,
+                references: Iterable[DocumentReference]) -> Iterable[DocumentSnapshot]:
+        return self._client.get_all(references)
+
+    def get(self, ref_or_query) -> Iterable[DocumentSnapshot]:
+        if isinstance(ref_or_query, DocumentReference):
+            return self._client.get_all([ref_or_query])
+        elif isinstance(ref_or_query, Query):
+            return ref_or_query.stream()
+        else:
+            raise ValueError(
+                'Value for argument "ref_or_query" must be a DocumentReference or a Query.'
+            )
+
+    # methods from
+    # https://googleapis.dev/python/firestore/latest/batch.html#google.cloud.firestore_v1.batch.WriteBatch
+
+    def _add_write_op(self, write_op: Callable):
+        if self._read_only:
+            raise ValueError(
+                "Cannot perform write operation in read-only transaction."
+            )
+        self._write_ops.append(write_op)
+
+    def create(self, reference: DocumentReference, document_data):
+        # this is a no-op, because if we have a DocumentReference
+        # it's already in the MockFirestore
+        ...
+
+    def set(self, reference: DocumentReference, document_data: dict,
+            merge=False):
+        write_op = partial(reference.set, document_data, merge=merge)
+        self._add_write_op(write_op)
+
+    def update(self, reference: DocumentReference,
+               field_updates: dict, option=None):
+        write_op = partial(reference.update, field_updates)
+        self._add_write_op(write_op)
+
+    def delete(self, reference: DocumentReference, option=None):
+        write_op = reference.delete
+        self._add_write_op(write_op)
+
+    def commit(self):
+        return self._commit()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_type is None:
+            self.commit()


Test patch:
diff --git a/tests/test_document_snapshot.py b/tests/test_document_snapshot.py
index 7d60e4b..896f45b 100644
--- a/tests/test_document_snapshot.py
+++ b/tests/test_document_snapshot.py
@@ -46,6 +46,38 @@ class TestDocumentSnapshot(TestCase):
         doc_snapshot = doc_ref.get()
         self.assertIs(doc_ref, doc_snapshot.reference)
 
+    def test_documentSnapshot_id(self):
+        fs = MockFirestore()
+        fs._data = {'foo': {
+            'first': {'id': 1}
+        }}
+        doc = fs.collection('foo').document('first').get()
+        self.assertIsInstance(doc.id, str)
+
+    def test_documentSnapshot_create_time(self):
+        fs = MockFirestore()
+        fs._data = {'foo': {
+            'first': {'id': 1}
+        }}
+        doc = fs.collection('foo').document('first').get()
+        self.assertIsNotNone(doc.create_time)
+
+    def test_documentSnapshot_update_time(self):
+        fs = MockFirestore()
+        fs._data = {'foo': {
+            'first': {'id': 1}
+        }}
+        doc = fs.collection('foo').document('first').get()
+        self.assertIsNotNone(doc.update_time)
+
+    def test_documentSnapshot_read_time(self):
+        fs = MockFirestore()
+        fs._data = {'foo': {
+            'first': {'id': 1}
+        }}
+        doc = fs.collection('foo').document('first').get()
+        self.assertIsNotNone(doc.read_time)
+
     def test_documentSnapshot_get_by_existing_field_path(self):
         fs = MockFirestore()
         fs._data = {'foo': {
diff --git a/tests/test_mock_client.py b/tests/test_mock_client.py
new file mode 100644
index 0000000..227c5c7
--- /dev/null
+++ b/tests/test_mock_client.py
@@ -0,0 +1,19 @@
+from unittest import TestCase
+
+from mockfirestore import MockFirestore
+
+
+class TestMockFirestore(TestCase):
+    def test_client_get_all(self):
+        fs = MockFirestore()
+        fs._data = {'foo': {
+            'first': {'id': 1},
+            'second': {'id': 2}
+        }}
+        doc = fs.collection('foo').document('first')
+        results = list(fs.get_all([doc]))
+        returned_doc_snapshot = results[0].to_dict()
+        expected_doc_snapshot = doc.get().to_dict()
+        self.assertEqual(returned_doc_snapshot, expected_doc_snapshot)
+
+
diff --git a/tests/test_transaction.py b/tests/test_transaction.py
new file mode 100644
index 0000000..72031fa
--- /dev/null
+++ b/tests/test_transaction.py
@@ -0,0 +1,73 @@
+from unittest import TestCase
+from mockfirestore import MockFirestore, Transaction
+
+
+class TestTransaction(TestCase):
+    def setUp(self) -> None:
+        self.fs = MockFirestore()
+        self.fs._data = {'foo': {
+                'first': {'id': 1},
+                'second': {'id': 2}
+            }}
+
+    def test_transaction_getAll(self):
+        with Transaction(self.fs) as transaction:
+            transaction._begin()
+            docs = [self.fs.collection('foo').document(doc_name)
+                    for doc_name in self.fs._data['foo']]
+            results = list(transaction.get_all(docs))
+            returned_docs_snapshots = [result.to_dict() for result in results]
+            expected_doc_snapshots = [doc.get().to_dict() for doc in docs]
+            for expected_snapshot in expected_doc_snapshots:
+                self.assertIn(expected_snapshot, returned_docs_snapshots)
+
+    def test_transaction_getDocument(self):
+        with Transaction(self.fs) as transaction:
+            transaction._begin()
+            doc = self.fs.collection('foo').document('first')
+            returned_doc = next(transaction.get(doc))
+            self.assertEqual(doc.get().to_dict(), returned_doc.to_dict())
+
+    def test_transaction_getQuery(self):
+        with Transaction(self.fs) as transaction:
+            transaction._begin()
+            query = self.fs.collection('foo').order_by('id')
+            returned_docs = [doc.to_dict() for doc in transaction.get(query)]
+            query = self.fs.collection('foo').order_by('id')
+            expected_docs = [doc.to_dict() for doc in query.stream()]
+            self.assertEqual(returned_docs, expected_docs)
+
+    def test_transaction_set_setsContentOfDocument(self):
+        doc_content = {'id': '3'}
+        doc_ref = self.fs.collection('foo').document('third')
+        with Transaction(self.fs) as transaction:
+            transaction._begin()
+            transaction.set(doc_ref, doc_content)
+        self.assertEqual(doc_ref.get().to_dict(), doc_content)
+
+    def test_transaction_set_mergeNewValue(self):
+        doc = self.fs.collection('foo').document('first')
+        with Transaction(self.fs) as transaction:
+            transaction._begin()
+            transaction.set(doc, {'updated': True}, merge=True)
+        updated_doc = {'id': 1, 'updated': True}
+        self.assertEqual(doc.get().to_dict(), updated_doc)
+
+    def test_transaction_update_changeExistingValue(self):
+        doc = self.fs.collection('foo').document('first')
+        with Transaction(self.fs) as transaction:
+            transaction._begin()
+            transaction.update(doc, {'updated': False})
+        updated_doc = {'id': 1, 'updated': False}
+        self.assertEqual(doc.get().to_dict(), updated_doc)
+
+    def test_transaction_delete_documentDoesNotExistAfterDelete(self):
+        doc = self.fs.collection('foo').document('first')
+        with Transaction(self.fs) as transaction:
+            transaction._begin()
+            transaction.delete(doc)
+        doc = self.fs.collection('foo').document('first').get()
+        self.assertEqual(False, doc.exists)
+
+
+