From 35185a849053877c9cc561e75cdb4cd7338cc508 Mon Sep 17 00:00:00 2001 From: Raphael Long Date: Tue, 28 Jul 2020 15:35:58 -0500 Subject: [PATCH] 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 --- .../cloud/firestore_v1/async_transaction.py | 14 ++-- tests/unit/v1/test_async_transaction.py | 83 +++++-------------- 2 files changed, 27 insertions(+), 70 deletions(-) diff --git a/google/cloud/firestore_v1/async_transaction.py b/google/cloud/firestore_v1/async_transaction.py index f572c173f..0b1e83788 100644 --- a/google/cloud/firestore_v1/async_transaction.py +++ b/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, diff --git a/tests/unit/v1/test_async_transaction.py b/tests/unit/v1/test_async_transaction.py index b27f30e9c..6f12c3394 100644 --- a/tests/unit/v1/test_async_transaction.py +++ b/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