Clean up a few more uses of gRPC in _pb2.py files

Were this not done these would break when the default behavior of gRPC
Python Protoc Plug-In is changed to be the put-gRPC-code-elements-only-
in-_pb2_grpc.py-files behavior that currently happens only when the
grpc_2_0 flag is passed.

3 files changed, 18 insertions, 14 deletions

diff --git a/src/python/grpcio_tests/tests/protoc_plugin/_python_plugin_test.py b/src/python/grpcio_tests/tests/protoc_plugin/_python_plugin_test.py
index 71493bfec6..5b84001aab 100644
--- a/src/python/grpcio_tests/tests/protoc_plugin/_python_plugin_test.py
+++ b/src/python/grpcio_tests/tests/protoc_plugin/_python_plugin_test.py
@@ -33,7 +33,7 @@ from tests.unit.framework.common import test_constants
 import tests.protoc_plugin.protos.payload.test_payload_pb2 as payload_pb2
 import tests.protoc_plugin.protos.requests.r.test_requests_pb2 as request_pb2
 import tests.protoc_plugin.protos.responses.test_responses_pb2 as response_pb2
-import tests.protoc_plugin.protos.service.test_service_pb2 as service_pb2
+import tests.protoc_plugin.protos.service.test_service_pb2_grpc as service_pb2_grpc
 
 # Identifiers of entities we expect to find in the generated module.
 STUB_IDENTIFIER = 'TestServiceStub'
@@ -138,7 +138,7 @@ def _CreateService():
   """
     servicer_methods = _ServicerMethods()
 
-    class Servicer(getattr(service_pb2, SERVICER_IDENTIFIER)):
+    class Servicer(getattr(service_pb2_grpc, SERVICER_IDENTIFIER)):
 
         def UnaryCall(self, request, context):
             return servicer_methods.UnaryCall(request, context)
@@ -157,11 +157,12 @@ def _CreateService():
 
     server = grpc.server(
         futures.ThreadPoolExecutor(max_workers=test_constants.POOL_SIZE))
-    getattr(service_pb2, ADD_SERVICER_TO_SERVER_IDENTIFIER)(Servicer(), server)
+    getattr(service_pb2_grpc, ADD_SERVICER_TO_SERVER_IDENTIFIER)(Servicer(),
+                                                                 server)
     port = server.add_insecure_port('[::]:0')
     server.start()
     channel = grpc.insecure_channel('localhost:{}'.format(port))
-    stub = getattr(service_pb2, STUB_IDENTIFIER)(channel)
+    stub = getattr(service_pb2_grpc, STUB_IDENTIFIER)(channel)
     return _Service(servicer_methods, server, stub)
 
 
@@ -173,16 +174,17 @@ def _CreateIncompleteService():
       servicer_methods implements none of the methods required of it.
   """
 
-    class Servicer(getattr(service_pb2, SERVICER_IDENTIFIER)):
+    class Servicer(getattr(service_pb2_grpc, SERVICER_IDENTIFIER)):
         pass
 
     server = grpc.server(
         futures.ThreadPoolExecutor(max_workers=test_constants.POOL_SIZE))
-    getattr(service_pb2, ADD_SERVICER_TO_SERVER_IDENTIFIER)(Servicer(), server)
+    getattr(service_pb2_grpc, ADD_SERVICER_TO_SERVER_IDENTIFIER)(Servicer(),
+                                                                 server)
     port = server.add_insecure_port('[::]:0')
     server.start()
     channel = grpc.insecure_channel('localhost:{}'.format(port))
-    stub = getattr(service_pb2, STUB_IDENTIFIER)(channel)
+    stub = getattr(service_pb2_grpc, STUB_IDENTIFIER)(channel)
     return _Service(None, server, stub)
 
 
@@ -223,10 +225,11 @@ class PythonPluginTest(unittest.TestCase):
 
     def testImportAttributes(self):
         # check that we can access the generated module and its members.
-        self.assertIsNotNone(getattr(service_pb2, STUB_IDENTIFIER, None))
-        self.assertIsNotNone(getattr(service_pb2, SERVICER_IDENTIFIER, None))
+        self.assertIsNotNone(getattr(service_pb2_grpc, STUB_IDENTIFIER, None))
         self.assertIsNotNone(
-            getattr(service_pb2, ADD_SERVICER_TO_SERVER_IDENTIFIER, None))
+            getattr(service_pb2_grpc, SERVICER_IDENTIFIER, None))
+        self.assertIsNotNone(
+            getattr(service_pb2_grpc, ADD_SERVICER_TO_SERVER_IDENTIFIER, None))
 
     def testUpDown(self):
         service = _CreateService()
diff --git a/src/python/grpcio_tests/tests/qps/benchmark_server.py b/src/python/grpcio_tests/tests/qps/benchmark_server.py
index 05101fdc6d..bb07844491 100644
--- a/src/python/grpcio_tests/tests/qps/benchmark_server.py
+++ b/src/python/grpcio_tests/tests/qps/benchmark_server.py
@@ -13,10 +13,10 @@
 # limitations under the License.
 
 from src.proto.grpc.testing import messages_pb2
-from src.proto.grpc.testing import services_pb2
+from src.proto.grpc.testing import services_pb2_grpc
 
 
-class BenchmarkServer(services_pb2.BenchmarkServiceServicer):
+class BenchmarkServer(services_pb2_grpc.BenchmarkServiceServicer):
     """Synchronous Server implementation for the Benchmark service."""
 
     def UnaryCall(self, request, context):
@@ -29,7 +29,7 @@ class BenchmarkServer(services_pb2.BenchmarkServiceServicer):
             yield messages_pb2.SimpleResponse(payload=payload)
 
 
-class GenericBenchmarkServer(services_pb2.BenchmarkServiceServicer):
+class GenericBenchmarkServer(services_pb2_grpc.BenchmarkServiceServicer):
     """Generic Server implementation for the Benchmark service."""
 
     def __init__(self, resp_size):
diff --git a/src/python/grpcio_tests/tests/stress/metrics_server.py b/src/python/grpcio_tests/tests/stress/metrics_server.py
index 11ab6c3f4e..33a74b4a38 100644
--- a/src/python/grpcio_tests/tests/stress/metrics_server.py
+++ b/src/python/grpcio_tests/tests/stress/metrics_server.py
@@ -16,11 +16,12 @@
 import time
 
 from src.proto.grpc.testing import metrics_pb2
+from src.proto.grpc.testing import metrics_pb2_grpc
 
 GAUGE_NAME = 'python_overall_qps'
 
 
-class MetricsServer(metrics_pb2.MetricsServiceServicer):
+class MetricsServer(metrics_pb2_grpc.MetricsServiceServicer):
 
     def __init__(self, histogram):
         self._start_time = time.time()