From 1a973f37e86207925e705d3fccdc88875d5d3ad1 Mon Sep 17 00:00:00 2001 From: Lidi Zheng Date: Thu, 18 Feb 2021 11:59:09 -0800 Subject: [PATCH] fix: patch emulator channel to be created accordingly (#288) Co-authored-by: Christopher Wilcox --- google/cloud/firestore_v1/base_client.py | 25 ++++++++++++++---------- tests/unit/v1/test_base_client.py | 13 ++++++++++-- 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/google/cloud/firestore_v1/base_client.py b/google/cloud/firestore_v1/base_client.py index 0f3c8e706..b2af21e3f 100644 --- a/google/cloud/firestore_v1/base_client.py +++ b/google/cloud/firestore_v1/base_client.py @@ -148,7 +148,7 @@ def _firestore_api_helper(self, transport, client_class, client_module) -> Any: # We need this in order to set appropriate keepalive options. if self._emulator_host is not None: - channel = self._emulator_channel() + channel = self._emulator_channel(transport) else: channel = transport.create_channel( self._target, @@ -165,25 +165,30 @@ def _firestore_api_helper(self, transport, client_class, client_module) -> Any: return self._firestore_api_internal - def _emulator_channel(self): + def _emulator_channel(self, transport): """ Creates a channel using self._credentials in a similar way to grpc.secure_channel but using grpc.local_channel_credentials() rather than grpc.ssh_channel_credentials() to allow easy connection to a local firestore emulator. This allows local testing of firestore rules if the credentials have been created from a signed custom token. - :return: grcp.Channel + :return: grpc.Channel or grpc.aio.Channel """ - return grpc._channel.Channel( - self._emulator_host, - (), - self._local_composite_credentials()._credentials, - None, - ) + # TODO: Implement a special credentials type for emulator and use + # "transport.create_channel" to create gRPC channels once google-auth + # extends it's allowed credentials types. + if "GrpcAsyncIOTransport" in str(transport.__name__): + return grpc.aio.secure_channel( + self._emulator_host, self._local_composite_credentials() + ) + else: + return grpc.secure_channel( + self._emulator_host, self._local_composite_credentials() + ) def _local_composite_credentials(self): """ - Ceates the credentials for the local emulator channel + Creates the credentials for the local emulator channel :return: grpc.ChannelCredentials """ credentials = google.auth.credentials.with_scopes_if_required( diff --git a/tests/unit/v1/test_base_client.py b/tests/unit/v1/test_base_client.py index 3dd7ff862..fd176d760 100644 --- a/tests/unit/v1/test_base_client.py +++ b/tests/unit/v1/test_base_client.py @@ -138,6 +138,13 @@ def test__rpc_metadata_property_with_emulator(self): ) def test_emulator_channel(self): + from google.cloud.firestore_v1.services.firestore.transports.grpc import ( + FirestoreGrpcTransport, + ) + from google.cloud.firestore_v1.services.firestore.transports.grpc_asyncio import ( + FirestoreGrpcAsyncIOTransport, + ) + emulator_host = "localhost:8081" with mock.patch("os.getenv") as getenv: getenv.return_value = emulator_host @@ -149,8 +156,10 @@ def test_emulator_channel(self): ) # checks that a channel is created - channel = client._emulator_channel() - self.assertTrue(isinstance(channel, grpc._channel.Channel)) + channel = client._emulator_channel(FirestoreGrpcTransport) + self.assertTrue(isinstance(channel, grpc.Channel)) + channel = client._emulator_channel(FirestoreGrpcAsyncIOTransport) + self.assertTrue(isinstance(channel, grpc.aio.Channel)) # checks that the credentials are composite ones using a local channel from grpc composite_credentials = client._local_composite_credentials() self.assertTrue(isinstance(composite_credentials, grpc.ChannelCredentials))