got unit tests working

Author: daniel-sanche

google/cloud/bigtable/data/_async/client.py

@@ -238,19 +238,15 @@ def __init__(
                     stacklevel=2,
                 )
 
-    def _build_grpc_channel(self, *args, **kwargs):
+    def _build_grpc_channel(self, *args, **kwargs) -> _AsyncReplaceableChannel:
         if self._emulator_host is not None:
             # emulators use insecure channel
-            return insecure_channel(self._emulator_host)
-        create_channel_fn = partial(TransportType.create_channel, *args, **kwargs)
-        # interceptors are handled differently between async and sync, because of differences in the grpc and gapic layers
+            create_channel_fn = partial(insecure_channel, self._emulator_host)
+        else:
+            create_channel_fn = partial(TransportType.create_channel, *args, **kwargs)
         if CrossSync.is_async:
-            # for async, add interceptors to the creation function
-            create_channel_fn = partial(create_channel_fn, interceptors=[_LoggingClientAIOInterceptor()])
             return _AsyncReplaceableChannel(create_channel_fn)
         else:
-            # for sync, chain interceptors using grpc.channel.intercept
-            # LoggingClientInterceptor not needed, since it is chained in the gapic layer
             return TransportType.create_channel(*args, **kwargs)
 
 
@@ -377,13 +373,18 @@ async def _manage_channel(
             grace_period: time to allow previous channel to serve existing
                 requests before closing, in seconds
         """
+        channel = self.transport.grpc_channel
+        if not isinstance(self.transport.grpc_channel, _AsyncReplaceableChannel):
+            warnings.warn("Channel does not support auto-refresh.")
+            return
+        channel: _AsyncReplaceableChannel = self.transport.grpc_channel
         first_refresh = self._channel_init_time + random.uniform(
             refresh_interval_min, refresh_interval_max
         )
         next_sleep = max(first_refresh - time.monotonic(), 0)
         if next_sleep > 0:
             # warm the current channel immediately
-            await self._ping_and_warm_instances(channel=self.transport.grpc_channel)
+            await self._ping_and_warm_instances(channel=channel)
         # continuously refresh the channel every `refresh_interval` seconds
         while not self._is_closed.is_set():
             await CrossSync.event_wait(
@@ -396,10 +397,10 @@ async def _manage_channel(
                 break
             start_timestamp = time.monotonic()
             # prepare new channel for use
-            new_channel = self.transport.grpc_channel.create_channel()
-            await self._ping_and_warm_instances(channel=new_channel)
+            new_sub_channel = channel.create_channel()
+            await self._ping_and_warm_instances(channel=new_sub_channel)
             # cycle channel out of use, with long grace window before closure
-            await self.transport.grpc_channel.replace_channel(new_channel, grace_period)
+            await channel.replace_wrapped_channel(new_sub_channel, grace_period)
             # subtract the time spent waiting for the channel to be replaced
             next_refresh = random.uniform(refresh_interval_min, refresh_interval_max)
             next_sleep = max(next_refresh - (time.monotonic() - start_timestamp), 0)

google/cloud/bigtable/data/_async/replaceable_channel.py

@@ -62,15 +62,19 @@ def get_state(self, try_to_connect: bool = False) -> grpc.ChannelConnectivity:
     async def wait_for_state_change(self, last_observed_state):
         return await self._channel.wait_for_state_change(last_observed_state)
 
-    @property
-    def wrapped_channel(self):
-        return self._channel
-
     def create_channel(self) -> aio.Channel:
         return self._channel_fn()
 
-    async def replace_channel(self, new_channel: aio.Channel, grace_period: float | None) -> aio.Channel:
+    async def replace_wrapped_channel(self, new_channel: aio.Channel, grace_period: float | None, copy_async_interceptors: bool=True) -> aio.Channel:
         old_channel = self._channel
+        if CrossSync.is_async and copy_async_interceptors:
+            # copy over interceptors
+            # this is needed because of how gapic attaches the LoggingClientAIOInterceptor
+            # sync channels add interceptors by wrapping, so this step isn't needed
+            new_channel._unary_unary_interceptors = old_channel._unary_unary_interceptors
+            new_channel._unary_stream_interceptors = old_channel._unary_stream_interceptors
+            new_channel._stream_unary_interceptors = old_channel._stream_unary_interceptors
+            new_channel._stream_stream_interceptors = old_channel._stream_stream_interceptors
         self._channel = new_channel
         # give old_channel a chance to complete existing rpcs
         if CrossSync.is_async:
@@ -81,7 +85,5 @@ async def replace_channel(self, new_channel: aio.Channel, grace_period: float |
             old_channel.close()  # type: ignore
         return old_channel
 
-    @property
-    def _unary_unary_interceptors(self):
-        # return empty list for compatibility with gapic layer
-        return []
+    def __getattr__(self, name):
+        return getattr(self._channel, name)

google/cloud/bigtable_v2/services/bigtable/transports/grpc_asyncio.py

@@ -290,7 +290,6 @@ def __init__(
             always_use_jwt_access=always_use_jwt_access,
             api_audience=api_audience,
         )
-
         if not self._grpc_channel:
             # initialize with the provided callable or the default channel
             channel_init = channel or type(self).create_channel

tests/system/data/test_system_async.py

@@ -252,13 +252,14 @@ async def test_channel_refresh(self, table_id, instance_id, temp_rows):
             await CrossSync.yield_to_event_loop()
             async with client.get_table(instance_id, table_id) as table:
                 rows = await table.read_rows({})
-                first_channel = client.transport.grpc_channel
+                channel_wrapper = client.transport.grpc_channel
+                first_channel = client.transport.grpc_channel._channel
                 assert len(rows) == 2
                 await CrossSync.sleep(2)
                 rows_after_refresh = await table.read_rows({})
                 assert len(rows_after_refresh) == 2
-                assert client.transport.grpc_channel is not first_channel
-                print(table)
+                assert client.transport.grpc_channel is channel_wrapper
+                assert client.transport.grpc_channel._channel is not first_channel
         finally:
             await client.close()
 

tests/unit/data/_async/test_client.py

@@ -223,7 +223,7 @@ async def test__start_background_channel_refresh_task_exists(self):
     @CrossSync.pytest
     async def test__start_background_channel_refresh(self):
         # should create background tasks for each channel
-        client = self._make_client(project="project-id")
+        client = self._make_client(project="project-id", use_emulator=False)
         with mock.patch.object(
             client, "_ping_and_warm_instances", CrossSync.Mock()
         ) as ping_and_warm:
@@ -366,7 +366,7 @@ async def test__manage_channel_first_sleep(
             with mock.patch.object(CrossSync, "event_wait") as sleep:
                 sleep.side_effect = asyncio.CancelledError
                 try:
-                    client = self._make_client(project="project-id")
+                    client = self._make_client(project="project-id", use_emulator=False)
                     client._channel_init_time = -wait_time
                     await client._manage_channel(refresh_interval, refresh_interval)
                 except asyncio.CancelledError:
@@ -395,32 +395,29 @@ async def test__manage_channel_ping_and_warm(self):
                 _LoggingClientInterceptor as Interceptor,
             )
 
-        client_mock = mock.Mock()
-        client_mock.transport._interceptor = Interceptor()
-        client_mock._is_closed.is_set.return_value = False
-        client_mock._channel_init_time = time.monotonic()
-        orig_channel = client_mock.transport.grpc_channel
+        client = self._make_client(project="project-id", use_emulator=True)
+        orig_channel = client.transport.grpc_channel
         # should ping an warm all new channels, and old channels if sleeping
         sleep_tuple = (
             (asyncio, "sleep") if CrossSync.is_async else (threading.Event, "wait")
         )
-        with mock.patch.object(*sleep_tuple):
-            # stop process after close is called
-            orig_channel.close.side_effect = asyncio.CancelledError
-            ping_and_warm = client_mock._ping_and_warm_instances = CrossSync.Mock()
+        with mock.patch.object(*sleep_tuple) as sleep_mock:
+            # stop process after loop
+            sleep_mock.side_effect = [None, asyncio.CancelledError]
+            ping_and_warm = client._ping_and_warm_instances = CrossSync.Mock()
             # should ping and warm old channel then new if sleep > 0
             try:
-                await self._get_target_class()._manage_channel(client_mock, 10)
+                await client._manage_channel(10)
             except asyncio.CancelledError:
                 pass
             # should have called at loop start, and after replacement
             assert ping_and_warm.call_count == 2
             # should have replaced channel once
-            assert client_mock.transport._grpc_channel != orig_channel
+            assert client.transport.grpc_channel._channel != orig_channel
             # make sure new and old channels were warmed
             called_with = [call[1]["channel"] for call in ping_and_warm.call_args_list]
             assert orig_channel in called_with
-            assert client_mock.transport.grpc_channel in called_with
+            assert client.transport.grpc_channel._channel in called_with
 
     @CrossSync.pytest
     @pytest.mark.parametrize(
@@ -438,8 +435,6 @@ async def test__manage_channel_sleeps(
         import time
         import random
 
-        channel = mock.Mock()
-        channel.close = CrossSync.Mock()
         with mock.patch.object(random, "uniform") as uniform:
             uniform.side_effect = lambda min_, max_: min_
             with mock.patch.object(time, "time") as time_mock:
@@ -448,8 +443,7 @@ async def test__manage_channel_sleeps(
                     sleep.side_effect = [None for i in range(num_cycles - 1)] + [
                         asyncio.CancelledError
                     ]
-                    client = self._make_client(project="project-id")
-                    client.transport._grpc_channel = channel
+                    client = self._make_client(project="project-id", use_emulator=True)
                     with mock.patch.object(
                         client.transport, "create_channel", CrossSync.Mock
                     ):
@@ -478,7 +472,7 @@ async def test__manage_channel_random(self):
                 uniform.return_value = 0
                 try:
                     uniform.side_effect = asyncio.CancelledError
-                    client = self._make_client(project="project-id")
+                    client = self._make_client(project="project-id", use_emulator=False)
                 except asyncio.CancelledError:
                     uniform.side_effect = None
                     uniform.reset_mock()
@@ -512,7 +506,7 @@ async def test__manage_channel_refresh(self, num_cycles):
                 CrossSync.grpc_helpers, "create_channel"
             ) as create_channel:
                 create_channel.return_value = new_channel
-                client = self._make_client(project="project-id")
+                client = self._make_client(project="project-id", use_emulator=False)
                 create_channel.reset_mock()
                 try:
                     await client._manage_channel(