feat: asyncio system tests (#132)

* feat: make collections call backed by async

* fix: failing asyncmock assertion

* fix: lint

* refactor: move AsyncMock to test helpers

* fix: rename transactional function to avoid collision

* feat: add async surface to firestore_v1 and firestore modules

* feat: add pytest-asyncio to noxfile installs

* feat: add transport to top level interface for client

* fix: batch_get_documents invocation

* fix: list_documents return type

* fix: run_query invocation

* fix: lint

* feat: add async system tests

* feat: remove Watch from async interface

* rebase: v2-staging

* fix: remove unused _transport property change

* fix: alpha sort module imports

* fix: dedup system test helpers

Author: rafilong

google/cloud/firestore.py

@@ -18,6 +18,13 @@
 from google.cloud.firestore_v1 import __version__
 from google.cloud.firestore_v1 import ArrayRemove
 from google.cloud.firestore_v1 import ArrayUnion
+from google.cloud.firestore_v1 import AsyncClient
+from google.cloud.firestore_v1 import AsyncCollectionReference
+from google.cloud.firestore_v1 import AsyncDocumentReference
+from google.cloud.firestore_v1 import AsyncQuery
+from google.cloud.firestore_v1 import async_transactional
+from google.cloud.firestore_v1 import AsyncTransaction
+from google.cloud.firestore_v1 import AsyncWriteBatch
 from google.cloud.firestore_v1 import Client
 from google.cloud.firestore_v1 import CollectionReference
 from google.cloud.firestore_v1 import DELETE_FIELD
@@ -45,6 +52,13 @@
     "__version__",
     "ArrayRemove",
     "ArrayUnion",
+    "AsyncClient",
+    "AsyncCollectionReference",
+    "AsyncDocumentReference",
+    "AsyncQuery",
+    "async_transactional",
+    "AsyncTransaction",
+    "AsyncWriteBatch",
     "Client",
     "CollectionReference",
     "DELETE_FIELD",

google/cloud/firestore_v1/__init__.py

@@ -29,21 +29,28 @@
 from google.cloud.firestore_v1._helpers import LastUpdateOption
 from google.cloud.firestore_v1._helpers import ReadAfterWriteError
 from google.cloud.firestore_v1._helpers import WriteOption
+from google.cloud.firestore_v1.async_batch import AsyncWriteBatch
+from google.cloud.firestore_v1.async_client import AsyncClient
+from google.cloud.firestore_v1.async_collection import AsyncCollectionReference
+from google.cloud.firestore_v1.async_document import AsyncDocumentReference
+from google.cloud.firestore_v1.async_query import AsyncQuery
+from google.cloud.firestore_v1.async_transaction import async_transactional
+from google.cloud.firestore_v1.async_transaction import AsyncTransaction
+from google.cloud.firestore_v1.base_document import DocumentSnapshot
 from google.cloud.firestore_v1.batch import WriteBatch
 from google.cloud.firestore_v1.client import Client
 from google.cloud.firestore_v1.collection import CollectionReference
+from google.cloud.firestore_v1.document import DocumentReference
+from google.cloud.firestore_v1.query import Query
+from google.cloud.firestore_v1.transaction import Transaction
+from google.cloud.firestore_v1.transaction import transactional
 from google.cloud.firestore_v1.transforms import ArrayRemove
 from google.cloud.firestore_v1.transforms import ArrayUnion
 from google.cloud.firestore_v1.transforms import DELETE_FIELD
 from google.cloud.firestore_v1.transforms import Increment
 from google.cloud.firestore_v1.transforms import Maximum
 from google.cloud.firestore_v1.transforms import Minimum
 from google.cloud.firestore_v1.transforms import SERVER_TIMESTAMP
-from google.cloud.firestore_v1.document import DocumentReference
-from google.cloud.firestore_v1.document import DocumentSnapshot
-from google.cloud.firestore_v1.query import Query
-from google.cloud.firestore_v1.transaction import Transaction
-from google.cloud.firestore_v1.transaction import transactional
 from google.cloud.firestore_v1.watch import Watch
 
 
@@ -100,6 +107,13 @@
     "__version__",
     "ArrayRemove",
     "ArrayUnion",
+    "AsyncClient",
+    "AsyncCollectionReference",
+    "AsyncDocumentReference",
+    "AsyncQuery",
+    "async_transactional",
+    "AsyncTransaction",
+    "AsyncWriteBatch",
     "Client",
     "CollectionReference",
     "DELETE_FIELD",

google/cloud/firestore_v1/async_client.py

@@ -242,7 +242,7 @@ async def get_all(self, references, field_paths=None, transaction=None):
         """
         document_paths, reference_map = _reference_info(references)
         mask = _get_doc_mask(field_paths)
-        response_iterator = self._firestore_api.batch_get_documents(
+        response_iterator = await self._firestore_api.batch_get_documents(
             request={
                 "database": self._database_string,
                 "documents": document_paths,

google/cloud/firestore_v1/async_collection.py

@@ -22,8 +22,6 @@
     _item_to_document_ref,
 )
 from google.cloud.firestore_v1 import async_query
-from google.cloud.firestore_v1.watch import Watch
-from google.cloud.firestore_v1 import async_document
 
 
 class AsyncCollectionReference(BaseCollectionReference):
@@ -119,7 +117,8 @@ async def list_documents(self, page_size=None):
             },
             metadata=self._client._rpc_metadata,
         )
-        return (_item_to_document_ref(self, i) for i in iterator)
+        async for i in iterator:
+            yield _item_to_document_ref(self, i)
 
     async def get(self, transaction=None):
         """Deprecated alias for :meth:`stream`."""
@@ -161,36 +160,3 @@ async def stream(self, transaction=None):
         query = async_query.AsyncQuery(self)
         async for d in query.stream(transaction=transaction):
             yield d
-
-    def on_snapshot(self, callback):
-        """Monitor the documents in this collection.
-
-        This starts a watch on this collection using a background thread. The
-        provided callback is run on the snapshot of the documents.
-
-        Args:
-            callback (Callable[[:class:`~google.cloud.firestore.collection.CollectionSnapshot`], NoneType]):
-                a callback to run when a change occurs.
-
-        Example:
-            from google.cloud import firestore_v1
-
-            db = firestore_v1.Client()
-            collection_ref = db.collection(u'users')
-
-            def on_snapshot(collection_snapshot, changes, read_time):
-                for doc in collection_snapshot.documents:
-                    print(u'{} => {}'.format(doc.id, doc.to_dict()))
-
-            # Watch this collection
-            collection_watch = collection_ref.on_snapshot(on_snapshot)
-
-            # Terminate this watch
-            collection_watch.unsubscribe()
-        """
-        return Watch.for_query(
-            self._query(),
-            callback,
-            async_document.DocumentSnapshot,
-            async_document.AsyncDocumentReference,
-        )

google/cloud/firestore_v1/async_document.py

@@ -23,7 +23,6 @@
 from google.api_core import exceptions
 from google.cloud.firestore_v1 import _helpers
 from google.cloud.firestore_v1.types import common
-from google.cloud.firestore_v1.watch import Watch
 
 
 class AsyncDocumentReference(BaseDocumentReference):
@@ -385,39 +384,3 @@ async def collections(self, page_size=None):
         # iterator.document = self
         # iterator.item_to_value = _item_to_collection_ref
         # return iterator
-
-    def on_snapshot(self, callback):
-        """Watch this document.
-
-        This starts a watch on this document using a background thread. The
-        provided callback is run on the snapshot.
-
-        Args:
-            callback(Callable[[:class:`~google.cloud.firestore.document.DocumentSnapshot`], NoneType]):
-                a callback to run when a change occurs
-
-        Example:
-
-        .. code-block:: python
-
-            from google.cloud import firestore_v1
-
-            db = firestore_v1.Client()
-            collection_ref = db.collection(u'users')
-
-            def on_snapshot(document_snapshot, changes, read_time):
-                doc = document_snapshot
-                print(u'{} => {}'.format(doc.id, doc.to_dict()))
-
-            doc_ref = db.collection(u'users').document(
-                u'alovelace' + unique_resource_id())
-
-            # Watch this document
-            doc_watch = doc_ref.on_snapshot(on_snapshot)
-
-            # Terminate this watch
-            doc_watch.unsubscribe()
-        """
-        return Watch.for_document(
-            self, callback, DocumentSnapshot, AsyncDocumentReference
-        )

google/cloud/firestore_v1/async_query.py

@@ -27,8 +27,6 @@
 )
 
 from google.cloud.firestore_v1 import _helpers
-from google.cloud.firestore_v1 import async_document
-from google.cloud.firestore_v1.watch import Watch
 
 
 class AsyncQuery(BaseQuery):
@@ -149,7 +147,7 @@ async def stream(self, transaction=None):
             The next document that fulfills the query.
         """
         parent_path, expected_prefix = self._parent._parent_info()
-        response_iterator = self._client._firestore_api.run_query(
+        response_iterator = await self._client._firestore_api.run_query(
             request={
                 "parent": parent_path,
                 "structured_query": self._to_protobuf(),
@@ -169,39 +167,3 @@ async def stream(self, transaction=None):
                 )
             if snapshot is not None:
                 yield snapshot
-
-    def on_snapshot(self, callback):
-        """Monitor the documents in this collection that match this query.
-
-        This starts a watch on this query using a background thread. The
-        provided callback is run on the snapshot of the documents.
-
-        Args:
-            callback(Callable[[:class:`~google.cloud.firestore.query.QuerySnapshot`], NoneType]):
-                a callback to run when a change occurs.
-
-        Example:
-
-        .. code-block:: python
-
-            from google.cloud import firestore_v1
-
-            db = firestore_v1.Client()
-            query_ref = db.collection(u'users').where("user", "==", u'Ada')
-
-            def on_snapshot(docs, changes, read_time):
-                for doc in docs:
-                    print(u'{} => {}'.format(doc.id, doc.to_dict()))
-
-            # Watch this query
-            query_watch = query_ref.on_snapshot(on_snapshot)
-
-            # Terminate this watch
-            query_watch.unsubscribe()
-        """
-        return Watch.for_query(
-            self,
-            callback,
-            async_document.DocumentSnapshot,
-            async_document.AsyncDocumentReference,
-        )

google/cloud/firestore_v1/async_transaction.py

@@ -287,7 +287,7 @@ async def __call__(self, transaction, *args, **kwargs):
         raise ValueError(msg)
 
 
-def transactional(to_wrap):
+def async_transactional(to_wrap):
     """Decorate a callable so that it runs in a transaction.
 
     Args:

noxfile.py

@@ -124,7 +124,7 @@ def system(session):
     # Install all test dependencies, then install this package into the
     # virtualenv's dist-packages.
     session.install(
-        "mock", "pytest", "google-cloud-testutils",
+        "mock", "pytest", "pytest-asyncio", "google-cloud-testutils",
     )
     session.install("-e", ".")
 

tests/system/test__helpers.py

@@ -0,0 +1,10 @@
+import os
+import re
+from test_utils.system import unique_resource_id
+
+FIRESTORE_CREDS = os.environ.get("FIRESTORE_APPLICATION_CREDENTIALS")
+FIRESTORE_PROJECT = os.environ.get("GCLOUD_PROJECT")
+RANDOM_ID_REGEX = re.compile("^[a-zA-Z0-9]{20}$")
+MISSING_DOCUMENT = "No document to update: "
+DOCUMENT_EXISTS = "Document already exists: "
+UNIQUE_RESOURCE_ID = unique_resource_id("-")

tests/system/test_system.py

@@ -15,8 +15,6 @@
 import datetime
 import math
 import operator
-import os
-import re
 
 from google.oauth2 import service_account
 import pytest
@@ -28,16 +26,16 @@
 from google.cloud._helpers import _datetime_to_pb_timestamp
 from google.cloud._helpers import UTC
 from google.cloud import firestore_v1 as firestore
-from test_utils.system import unique_resource_id
 
 from time import sleep
 
-FIRESTORE_CREDS = os.environ.get("FIRESTORE_APPLICATION_CREDENTIALS")
-FIRESTORE_PROJECT = os.environ.get("GCLOUD_PROJECT")
-RANDOM_ID_REGEX = re.compile("^[a-zA-Z0-9]{20}$")
-MISSING_DOCUMENT = "No document to update: "
-DOCUMENT_EXISTS = "Document already exists: "
-UNIQUE_RESOURCE_ID = unique_resource_id("-")
+from tests.system.test__helpers import (
+    FIRESTORE_CREDS,
+    FIRESTORE_PROJECT,
+    RANDOM_ID_REGEX,
+    MISSING_DOCUMENT,
+    UNIQUE_RESOURCE_ID,
+)
 
 
 @pytest.fixture(scope=u"module")
@@ -683,7 +681,7 @@ def test_query_stream_w_offset(query_docs):
 
 def test_query_with_order_dot_key(client, cleanup):
     db = client
-    collection_id = "collek" + unique_resource_id("-")
+    collection_id = "collek" + UNIQUE_RESOURCE_ID
     collection = db.collection(collection_id)
     for index in range(100, -1, -1):
         doc = collection.document("test_{:09d}".format(index))