fix: Add emulator support to schema service

google/pubsub_v1/services/publisher/client.py

@@ -15,6 +15,7 @@
 #
 from collections import OrderedDict
 import functools
+import grpc
 import os
 import re
 from typing import Dict, Mapping, Optional, Sequence, Tuple, Type, Union
@@ -42,8 +43,6 @@
 from google.pubsub_v1.services.publisher import pagers
 from google.pubsub_v1.types import pubsub
 from google.pubsub_v1.types import TimeoutType
-
-import grpc
 from .transports.base import PublisherTransport, DEFAULT_CLIENT_INFO
 from .transports.grpc import PublisherGrpcTransport
 from .transports.grpc_asyncio import PublisherGrpcAsyncIOTransport

google/pubsub_v1/services/schema_service/client.py

@@ -14,6 +14,8 @@
 # limitations under the License.
 #
 from collections import OrderedDict
+import functools
+import grpc
 import os
 import re
 from typing import Dict, Mapping, Optional, Sequence, Tuple, Type, Union
@@ -411,6 +413,15 @@ def __init__(
                 )
 
             Transport = type(self).get_transport_class(transport)
+
+            emulator_host = os.environ.get("PUBSUB_EMULATOR_HOST")
+            if emulator_host:
+                if issubclass(Transport, type(self)._transport_registry["grpc"]):
+                    channel = grpc.insecure_channel(target=emulator_host)
+                else:
+                    channel = grpc.aio.insecure_channel(target=emulator_host)
+                Transport = functools.partial(Transport, channel=channel)
+
             self._transport = Transport(
                 credentials=credentials,
                 credentials_file=client_options.credentials_file,

google/pubsub_v1/services/subscriber/client.py

@@ -15,6 +15,7 @@
 #
 from collections import OrderedDict
 import functools
+import grpc
 import os
 import re
 from typing import (
@@ -52,8 +53,6 @@
 from google.protobuf import timestamp_pb2  # type: ignore
 from google.pubsub_v1.services.subscriber import pagers
 from google.pubsub_v1.types import pubsub
-
-import grpc
 from .transports.base import SubscriberTransport, DEFAULT_CLIENT_INFO
 from .transports.grpc import SubscriberGrpcTransport
 from .transports.grpc_asyncio import SubscriberGrpcAsyncIOTransport

owlbot.py

@@ -85,21 +85,13 @@
     clients_to_patch = [
         library / f"google/pubsub_{library.name}/services/publisher/client.py",
         library / f"google/pubsub_{library.name}/services/subscriber/client.py",
+        library / f"google/pubsub_{library.name}/services/schema_service/client.py",
     ]
     err_msg = (
         "Expected replacements for gRPC channel to use with the emulator not made."
     )
 
-    count = s.replace(clients_to_patch, r"import os", "import functools\n\g<0>")
-
-    if count < len(clients_to_patch):
-        raise Exception(err_msg)
-
-    count = s.replace(
-        clients_to_patch,
-        f"from google\.pubsub_{library.name}\.types import pubsub",
-        "\g<0>\n\nimport grpc",
-    )
+    count = s.replace(clients_to_patch, r"import os", "import functools\nimport grpc\n\g<0>")
 
     if count < len(clients_to_patch):
         raise Exception(err_msg)