test: improve async test coverage

Author: sinhasubham

google/cloud/aio/_cross_sync/cross_sync.py

@@ -89,6 +89,8 @@ class CrossSync(metaclass=MappingMeta):
     LifoQueue: TypeAlias = asyncio.LifoQueue
     PriorityQueue: TypeAlias = asyncio.PriorityQueue
     StopIteration: TypeAlias = StopAsyncIteration
+    QueueEmpty: TypeAlias = asyncio.QueueEmpty
+    QueueFull: TypeAlias = asyncio.QueueFull
     # provide aliases for common async type annotations
     Awaitable: TypeAlias = typing.Awaitable
     Iterable: TypeAlias = AsyncIterable
@@ -160,17 +162,29 @@ async def run_if_async(func, *args, **kwargs):
     @staticmethod
     async def queue_get(queue, block=True, timeout=None):
         if not block:
-            return queue.get_nowait()
+            try:
+                return queue.get_nowait()
+            except asyncio.QueueEmpty:
+                raise CrossSync.QueueEmpty()
         if timeout is not None:
-            return await asyncio.wait_for(queue.get(), timeout=timeout)
+            try:
+                return await asyncio.wait_for(queue.get(), timeout=timeout)
+            except asyncio.TimeoutError:
+                raise CrossSync.QueueEmpty()
         return await queue.get()
 
     @staticmethod
     async def queue_put(queue, item, block=True, timeout=None):
         if not block:
-            return queue.put_nowait(item)
+            try:
+                return queue.put_nowait(item)
+            except asyncio.QueueFull:
+                raise CrossSync.QueueFull()
         if timeout is not None:
-            await asyncio.wait_for(queue.put(item), timeout=timeout)
+            try:
+                await asyncio.wait_for(queue.put(item), timeout=timeout)
+            except asyncio.TimeoutError:
+                raise CrossSync.QueueFull()
         else:
             await queue.put(item)
 
@@ -304,6 +318,8 @@ class _Sync_Impl(metaclass=MappingMeta):
         Semaphore: TypeAlias = threading.Semaphore
         LifoQueue: TypeAlias = queue.LifoQueue
         PriorityQueue: TypeAlias = queue.PriorityQueue
+        QueueEmpty: TypeAlias = queue.Empty
+        QueueFull: TypeAlias = queue.Full
         StopIteration: TypeAlias = StopIteration
         # type annotations
         Awaitable: TypeAlias = Union[T]

google/cloud/spanner_v1/_async/_helpers.py

@@ -60,7 +60,9 @@ async def _retry(
             if retries >= retry_count:
                 raise e
             if before_next_retry:
-                before_next_retry(retries, delay)
+                res = before_next_retry(retries, delay)
+                if asyncio.iscoroutine(res) or inspect.isawaitable(res):
+                    await res
             await asyncio.sleep(delay)
             retries += 1
 

google/cloud/spanner_v1/_async/batch.py

@@ -21,14 +21,7 @@
 from google.api_core.exceptions import InternalServerError
 
 from google.cloud.aio._cross_sync import CrossSync
-from google.cloud.spanner_v1 import (
-    BatchWriteRequest,
-    CommitRequest,
-    CommitResponse,
-    Mutation,
-    RequestOptions,
-    TransactionOptions,
-)
+
 from google.cloud.spanner_v1._async._helpers import _retry, _retry_on_aborted_exception
 from google.cloud.spanner_v1._helpers import (
     AtomicCounter,
@@ -45,6 +38,14 @@
 )
 from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
 from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
+from google.cloud.spanner_v1.types.commit_response import CommitResponse
+from google.cloud.spanner_v1.types.mutation import Mutation
+from google.cloud.spanner_v1.types.spanner import (
+    BatchWriteRequest,
+    CommitRequest,
+    RequestOptions,
+)
+from google.cloud.spanner_v1.types.transaction import TransactionOptions
 
 DEFAULT_RETRY_TIMEOUT_SECS = 30
 

google/cloud/spanner_v1/_async/client.py

@@ -66,22 +66,22 @@
     ListInstanceConfigsRequest,
     ListInstancesRequest,
 )
-from google.cloud.spanner_v1.gapic_version import __version__
-from google.cloud.spanner_v1.transaction import DefaultTransactionOptions
-from google.cloud.spanner_v1.types.spanner import ExecuteSqlRequest
+from google.cloud.spanner_v1._async.instance import Instance
 from google.cloud.spanner_v1._helpers import (
     _create_experimental_host_transport,
     _validate_client_context,
 )
-from google.cloud.spanner_v1._async.instance import Instance
 from google.cloud.spanner_v1._helpers import _merge_query_options, _metadata_with_prefix
+from google.cloud.spanner_v1.gapic_version import __version__
 from google.cloud.spanner_v1.metrics.constants import METRIC_EXPORT_INTERVAL_MS
 from google.cloud.spanner_v1.metrics.metrics_exporter import (
     CloudMonitoringMetricsExporter,
 )
 from google.cloud.spanner_v1.metrics.spanner_metrics_tracer_factory import (
     SpannerMetricsTracerFactory,
 )
+from google.cloud.spanner_v1.transaction import DefaultTransactionOptions
+from google.cloud.spanner_v1.types.spanner import ExecuteSqlRequest
 
 try:
     from opentelemetry import metrics

google/cloud/spanner_v1/_async/database.py

@@ -43,13 +43,6 @@
 from google.cloud.spanner_admin_database_v1 import CreateDatabaseRequest
 from google.cloud.spanner_admin_database_v1 import Database as DatabasePB
 from google.cloud.spanner_admin_database_v1.types import DatabaseDialect
-from google.cloud.spanner_v1.transaction import DefaultTransactionOptions
-from google.cloud.spanner_v1.types.spanner import ExecuteSqlRequest
-from google.cloud.spanner_v1.types.spanner import RequestOptions
-from google.cloud.spanner_v1.types.transaction import TransactionOptions
-from google.cloud.spanner_v1.types.transaction import TransactionSelector
-from google.cloud.spanner_v1.types.type import Type
-from google.cloud.spanner_v1.types.type import TypeCode
 from google.cloud.spanner_v1._async.batch import Batch, MutationGroups
 from google.cloud.spanner_v1._async.database_sessions_manager import (
     DatabaseSessionsManager,
@@ -72,7 +65,16 @@
 from google.cloud.spanner_v1.services.spanner.async_client import (
     SpannerAsyncClient as SpannerClient,
 )
-from google.cloud.spanner_v1.transaction import BatchTransactionId
+from google.cloud.spanner_v1.transaction import (
+    BatchTransactionId,
+    DefaultTransactionOptions,
+)
+from google.cloud.spanner_v1.types.spanner import ExecuteSqlRequest, RequestOptions
+from google.cloud.spanner_v1.types.transaction import (
+    TransactionOptions,
+    TransactionSelector,
+)
+from google.cloud.spanner_v1.types.type import Type, TypeCode
 
 if CrossSync.is_async:
     from google.cloud.spanner_v1.services.spanner.transports.grpc_asyncio import (
@@ -198,13 +200,22 @@ def __init__(
         self._encryption_config = encryption_config
         self._database_dialect = database_dialect
         self._database_role = database_role
-        self._route_to_leader_enabled = self._instance._client.route_to_leader_enabled
+        if self._instance and self._instance._client:
+            self._route_to_leader_enabled = (
+                self._instance._client.route_to_leader_enabled
+            )
+        else:
+            self._route_to_leader_enabled = False
         self._enable_drop_protection = enable_drop_protection
         self._reconciling = False
-        self._directed_read_options = self._instance._client.directed_read_options
-        self.default_transaction_options: DefaultTransactionOptions = (
-            self._instance._client.default_transaction_options
-        )
+        if self._instance and self._instance._client:
+            self._directed_read_options = self._instance._client.directed_read_options
+            self.default_transaction_options: DefaultTransactionOptions = (
+                self._instance._client.default_transaction_options
+            )
+        else:
+            self._directed_read_options = None
+            self.default_transaction_options = None
         self._proto_descriptors = proto_descriptors
         self._channel_id = 0  # It'll be created when _spanner_api is created.
 
@@ -220,7 +231,9 @@ def __init__(
         except RuntimeError:
             # No running loop, bind should have been sync or will be failed later
             pass
-        self._experimental_host = self._instance._client._experimental_host
+        self._experimental_host = (
+            self._instance.experimental_host if self._instance else None
+        )
         is_experimental_host = self._experimental_host is not None
 
         self._sessions_manager = DatabaseSessionsManager(
@@ -231,8 +244,12 @@ def __init__(
     def _resource_info(self):
         """Resource information for metrics labels."""
         return {
-            "project": self._instance._client.project,
-            "instance": self._instance.instance_id,
+            "project": (
+                self._instance._client.project
+                if self._instance and self._instance._client
+                else None
+            ),
+            "instance": self._instance.instance_id if self._instance else None,
             "database": self.database_id,
         }
 
@@ -1351,7 +1368,8 @@ async def list_tables(self, schema="_default"):
             async for row in results:
                 yield self.table(row[0])
 
-    def get_iam_policy(self, policy_version=None):
+    @CrossSync.convert
+    async def get_iam_policy(self, policy_version=None):
         """Gets the access control policy for a database resource.
 
         :type policy_version: int
@@ -1374,13 +1392,14 @@ def get_iam_policy(self, policy_version=None):
                 requested_policy_version=policy_version
             ),
         )
-        response = api.get_iam_policy(
+        response = await api.get_iam_policy(
             request=request,
             metadata=self.metadata_with_request_id(self._next_nth_request, 1, metadata),
         )
         return response
 
-    def set_iam_policy(self, policy):
+    @CrossSync.convert
+    async def set_iam_policy(self, policy):
         """Sets the access control policy on a database resource.
         Replaces any existing policy.
 
@@ -1399,7 +1418,7 @@ def set_iam_policy(self, policy):
             resource=self.name,
             policy=policy,
         )
-        response = api.set_iam_policy(
+        response = await api.set_iam_policy(
             request=request,
             metadata=self.metadata_with_request_id(self._next_nth_request, 1, metadata),
         )
@@ -1430,6 +1449,11 @@ def sessions_manager(self) -> DatabaseSessionsManager:
         """
         return self._sessions_manager
 
+    @CrossSync.convert
+    async def close(self):
+        """Clean up underlying session manager and background tasks."""
+        await self._sessions_manager.close()
+
 
 class BatchCheckout(object):
     """Context manager for using a batch from a database.

google/cloud/spanner_v1/_async/database_sessions_manager.py

@@ -15,6 +15,7 @@
 """Manage sessions for a database."""
 __CROSS_SYNC_OUTPUT__ = "google.cloud.spanner_v1.database_sessions_manager"
 
+import asyncio
 from datetime import timedelta
 from enum import Enum
 from os import getenv
@@ -215,3 +216,20 @@ def _getenv(cls, env_var_name: str) -> bool:
         """Returns the value of the given environment variable as a boolean."""
         env_var_value = getenv(env_var_name, "true").lower().strip()
         return env_var_value != "false"
+
+    @CrossSync.convert
+    async def close(self) -> None:
+        """Closes the database session manager and stops all background tasks."""
+        self._multiplexed_session_terminate_event.set()
+        if self._multiplexed_session_thread is not None:
+            if CrossSync.is_async:
+                self._multiplexed_session_thread.cancel()
+                try:
+                    await self._multiplexed_session_thread
+                except CrossSync.rm_aio(asyncio.CancelledError):
+                    pass
+            else:
+                self._multiplexed_session_thread.join()
+        if self._multiplexed_session is not None:
+            await self._multiplexed_session.delete()
+            self._multiplexed_session = None