feat: asyncio microgen transaction (#123)

* refactor: move generated client instantiation out of base class

* feat: integrate microgen async client to client

* feat: make collections call backed by async

* fix: failing asyncmock assertion

* refactor: remove unused install

* fix: lint

* refactor: shared functionality in client to base class

* refactor: move AsyncMock to test helpers

* fix: return type in client docs

* feat: integrate microgen async client to collection

* fix: lint

* feat: integrate microgen async client to document

* feat: integrate microgen async client to batch

* fix: use AsyncMock for batch async tests:

* fix: collection and document testing batch

* feat: integrate microgen async client to transaction

* fix: remove unused imports

Author: rafilong

google/cloud/firestore_v1/async_transaction.py

@@ -85,7 +85,7 @@ async def _begin(self, retry_id=None):
             msg = _CANT_BEGIN.format(self._id)
             raise ValueError(msg)
 
-        transaction_response = self._client._firestore_api.begin_transaction(
+        transaction_response = await self._client._firestore_api.begin_transaction(
             request={
                 "database": self._client._database_string,
                 "options": self._options_protobuf(retry_id),
@@ -105,7 +105,7 @@ async def _rollback(self):
 
         try:
             # NOTE: The response is just ``google.protobuf.Empty``.
-            self._client._firestore_api.rollback(
+            await self._client._firestore_api.rollback(
                 request={
                     "database": self._client._database_string,
                     "transaction": self._id,
@@ -148,7 +148,7 @@ async def get_all(self, references):
             .DocumentSnapshot: The next document snapshot that fulfills the
             query, or :data:`None` if the document does not exist.
         """
-        return self._client.get_all(references, transaction=self)
+        return await self._client.get_all(references, transaction=self)
 
     async def get(self, ref_or_query):
         """
@@ -160,9 +160,9 @@ async def get(self, ref_or_query):
             query, or :data:`None` if the document does not exist.
         """
         if isinstance(ref_or_query, AsyncDocumentReference):
-            return self._client.get_all([ref_or_query], transaction=self)
+            return await self._client.get_all([ref_or_query], transaction=self)
         elif isinstance(ref_or_query, AsyncQuery):
-            return ref_or_query.stream(transaction=self)
+            return await ref_or_query.stream(transaction=self)
         else:
             raise ValueError(
                 'Value for argument "ref_or_query" must be a AsyncDocumentReference or a AsyncQuery.'
@@ -192,7 +192,7 @@ async def _pre_commit(self, transaction, *args, **kwargs):
 
         Args:
             transaction
-                (:class:`~google.cloud.firestore_v1.transaction.Transaction`):
+                (:class:`~google.cloud.firestore_v1.async_transaction.AsyncTransaction`):
                 A transaction to execute the callable within.
             args (Tuple[Any, ...]): The extra positional arguments to pass
                 along to the wrapped callable.
@@ -330,7 +330,7 @@ async def _commit_with_retry(client, write_pbs, transaction_id):
     current_sleep = _INITIAL_SLEEP
     while True:
         try:
-            return client._firestore_api.commit(
+            return await client._firestore_api.commit(
                 request={
                     "database": client._database_string,
                     "writes": write_pbs,

tests/unit/v1/test_async_transaction.py

@@ -14,7 +14,9 @@
 
 import pytest
 import aiounittest
+
 import mock
+from tests.unit.v1.test__helpers import AsyncMock
 
 
 class TestAsyncTransaction(aiounittest.AsyncTestCase):
@@ -80,15 +82,10 @@ def test__clean_up(self):
 
     @pytest.mark.asyncio
     async def test__begin(self):
-        from google.cloud.firestore_v1.services.firestore import (
-            client as firestore_client,
-        )
         from google.cloud.firestore_v1.types import firestore
 
         # Create a minimal fake GAPIC with a dummy result.
-        firestore_api = mock.create_autospec(
-            firestore_client.FirestoreClient, instance=True
-        )
+        firestore_api = AsyncMock()
         txn_id = b"to-begin"
         response = firestore.BeginTransactionResponse(transaction=txn_id)
         firestore_api.begin_transaction.return_value = response
@@ -128,14 +125,9 @@ async def test__begin_failure(self):
     @pytest.mark.asyncio
     async def test__rollback(self):
         from google.protobuf import empty_pb2
-        from google.cloud.firestore_v1.services.firestore import (
-            client as firestore_client,
-        )
 
         # Create a minimal fake GAPIC with a dummy result.
-        firestore_api = mock.create_autospec(
-            firestore_client.FirestoreClient, instance=True
-        )
+        firestore_api = AsyncMock()
         firestore_api.rollback.return_value = empty_pb2.Empty()
 
         # Attach the fake GAPIC to a real client.
@@ -172,14 +164,9 @@ async def test__rollback_not_allowed(self):
     @pytest.mark.asyncio
     async def test__rollback_failure(self):
         from google.api_core import exceptions
-        from google.cloud.firestore_v1.services.firestore import (
-            client as firestore_client,
-        )
 
         # Create a minimal fake GAPIC with a dummy failure.
-        firestore_api = mock.create_autospec(
-            firestore_client.FirestoreClient, instance=True
-        )
+        firestore_api = AsyncMock()
         exc = exceptions.InternalServerError("Fire during rollback.")
         firestore_api.rollback.side_effect = exc
 
@@ -207,16 +194,11 @@ async def test__rollback_failure(self):
 
     @pytest.mark.asyncio
     async def test__commit(self):
-        from google.cloud.firestore_v1.services.firestore import (
-            client as firestore_client,
-        )
         from google.cloud.firestore_v1.types import firestore
         from google.cloud.firestore_v1.types import write
 
         # Create a minimal fake GAPIC with a dummy result.
-        firestore_api = mock.create_autospec(
-            firestore_client.FirestoreClient, instance=True
-        )
+        firestore_api = AsyncMock()
         commit_response = firestore.CommitResponse(write_results=[write.WriteResult()])
         firestore_api.commit.return_value = commit_response
 
@@ -262,14 +244,9 @@ async def test__commit_not_allowed(self):
     @pytest.mark.asyncio
     async def test__commit_failure(self):
         from google.api_core import exceptions
-        from google.cloud.firestore_v1.services.firestore import (
-            client as firestore_client,
-        )
 
         # Create a minimal fake GAPIC with a dummy failure.
-        firestore_api = mock.create_autospec(
-            firestore_client.FirestoreClient, instance=True
-        )
+        firestore_api = AsyncMock()
         exc = exceptions.InternalServerError("Fire during commit.")
         firestore_api.commit.side_effect = exc
 
@@ -304,7 +281,7 @@ async def test__commit_failure(self):
 
     @pytest.mark.asyncio
     async def test_get_all(self):
-        client = mock.Mock(spec=["get_all"])
+        client = AsyncMock(spec=["get_all"])
         transaction = self._make_one(client)
         ref1, ref2 = mock.Mock(), mock.Mock()
         result = await transaction.get_all([ref1, ref2])
@@ -315,7 +292,7 @@ async def test_get_all(self):
     async def test_get_document_ref(self):
         from google.cloud.firestore_v1.async_document import AsyncDocumentReference
 
-        client = mock.Mock(spec=["get_all"])
+        client = AsyncMock(spec=["get_all"])
         transaction = self._make_one(client)
         ref = AsyncDocumentReference("documents", "doc-id")
         result = await transaction.get(ref)
@@ -326,10 +303,10 @@ async def test_get_document_ref(self):
     async def test_get_w_query(self):
         from google.cloud.firestore_v1.async_query import AsyncQuery
 
-        client = mock.Mock(spec=[])
+        client = AsyncMock(spec=[])
         transaction = self._make_one(client)
-        query = AsyncQuery(parent=mock.Mock(spec=[]))
-        query.stream = mock.MagicMock()
+        query = AsyncQuery(parent=AsyncMock(spec=[]))
+        query.stream = AsyncMock()
         result = await transaction.get(query)
         query.stream.assert_called_once_with(transaction=transaction)
         self.assertIs(result, query.stream.return_value)
@@ -804,14 +781,9 @@ async def _call_fut(client, write_pbs, transaction_id):
     @mock.patch("google.cloud.firestore_v1.async_transaction._sleep")
     @pytest.mark.asyncio
     async def test_success_first_attempt(self, _sleep):
-        from google.cloud.firestore_v1.services.firestore import (
-            client as firestore_client,
-        )
 
         # Create a minimal fake GAPIC with a dummy result.
-        firestore_api = mock.create_autospec(
-            firestore_client.FirestoreClient, instance=True
-        )
+        firestore_api = AsyncMock()
 
         # Attach the fake GAPIC to a real client.
         client = _make_client("summer")
@@ -839,14 +811,10 @@ async def test_success_first_attempt(self, _sleep):
     @pytest.mark.asyncio
     async def test_success_third_attempt(self, _sleep):
         from google.api_core import exceptions
-        from google.cloud.firestore_v1.services.firestore import (
-            client as firestore_client,
-        )
 
         # Create a minimal fake GAPIC with a dummy result.
-        firestore_api = mock.create_autospec(
-            firestore_client.FirestoreClient, instance=True
-        )
+        firestore_api = AsyncMock()
+
         # Make sure the first two requests fail and the third succeeds.
         firestore_api.commit.side_effect = [
             exceptions.ServiceUnavailable("Server sleepy."),
@@ -885,14 +853,10 @@ async def test_success_third_attempt(self, _sleep):
     @pytest.mark.asyncio
     async def test_failure_first_attempt(self, _sleep):
         from google.api_core import exceptions
-        from google.cloud.firestore_v1.services.firestore import (
-            client as firestore_client,
-        )
 
         # Create a minimal fake GAPIC with a dummy result.
-        firestore_api = mock.create_autospec(
-            firestore_client.FirestoreClient, instance=True
-        )
+        firestore_api = AsyncMock()
+
         # Make sure the first request fails with an un-retryable error.
         exc = exceptions.ResourceExhausted("We ran out of fries.")
         firestore_api.commit.side_effect = exc
@@ -923,14 +887,10 @@ async def test_failure_first_attempt(self, _sleep):
     @pytest.mark.asyncio
     async def test_failure_second_attempt(self, _sleep):
         from google.api_core import exceptions
-        from google.cloud.firestore_v1.services.firestore import (
-            client as firestore_client,
-        )
 
         # Create a minimal fake GAPIC with a dummy result.
-        firestore_api = mock.create_autospec(
-            firestore_client.FirestoreClient, instance=True
-        )
+        firestore_api = AsyncMock()
+
         # Make sure the first request fails retry-able and second
         # fails non-retryable.
         exc1 = exceptions.ServiceUnavailable("Come back next time.")
@@ -1031,15 +991,12 @@ def _make_client(project="feral-tom-cat"):
 
 def _make_transaction(txn_id, **txn_kwargs):
     from google.protobuf import empty_pb2
-    from google.cloud.firestore_v1.services.firestore import client as firestore_client
     from google.cloud.firestore_v1.types import firestore
     from google.cloud.firestore_v1.types import write
     from google.cloud.firestore_v1.async_transaction import AsyncTransaction
 
     # Create a fake GAPIC ...
-    firestore_api = mock.create_autospec(
-        firestore_client.FirestoreClient, instance=True
-    )
+    firestore_api = AsyncMock()
     # ... with a dummy ``BeginTransactionResponse`` result ...
     begin_response = firestore.BeginTransactionResponse(transaction=txn_id)
     firestore_api.begin_transaction.return_value = begin_response