Update tests to current gRPC code elements

8 files changed, 25 insertions, 23 deletions

diff --git a/src/python/grpcio_tests/tests/interop/_insecure_intraop_test.py b/src/python/grpcio_tests/tests/interop/_insecure_intraop_test.py
index 58f3b364ba..3325d54375 100644
--- a/src/python/grpcio_tests/tests/interop/_insecure_intraop_test.py
+++ b/src/python/grpcio_tests/tests/interop/_insecure_intraop_test.py
@@ -32,7 +32,7 @@ from concurrent import futures
 import unittest
 
 import grpc
-from src.proto.grpc.testing import test_pb2
+from src.proto.grpc.testing import test_pb2_grpc
 
 from tests.interop import _intraop_test_case
 from tests.interop import methods
@@ -44,11 +44,11 @@ class InsecureIntraopTest(_intraop_test_case.IntraopTestCase,
 
     def setUp(self):
         self.server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
-        test_pb2.add_TestServiceServicer_to_server(methods.TestService(),
-                                                   self.server)
+        test_pb2_grpc.add_TestServiceServicer_to_server(methods.TestService(),
+                                                        self.server)
         port = self.server.add_insecure_port('[::]:0')
         self.server.start()
-        self.stub = test_pb2.TestServiceStub(
+        self.stub = test_pb2_grpc.TestServiceStub(
             grpc.insecure_channel('localhost:{}'.format(port)))
 
 
diff --git a/src/python/grpcio_tests/tests/interop/_secure_intraop_test.py b/src/python/grpcio_tests/tests/interop/_secure_intraop_test.py
index 5fe929b99e..857e00efb3 100644
--- a/src/python/grpcio_tests/tests/interop/_secure_intraop_test.py
+++ b/src/python/grpcio_tests/tests/interop/_secure_intraop_test.py
@@ -32,7 +32,7 @@ from concurrent import futures
 import unittest
 
 import grpc
-from src.proto.grpc.testing import test_pb2
+from src.proto.grpc.testing import test_pb2_grpc
 
 from tests.interop import _intraop_test_case
 from tests.interop import methods
@@ -45,14 +45,14 @@ class SecureIntraopTest(_intraop_test_case.IntraopTestCase, unittest.TestCase):
 
     def setUp(self):
         self.server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
-        test_pb2.add_TestServiceServicer_to_server(methods.TestService(),
-                                                   self.server)
+        test_pb2_grpc.add_TestServiceServicer_to_server(methods.TestService(),
+                                                        self.server)
         port = self.server.add_secure_port(
             '[::]:0',
             grpc.ssl_server_credentials(
                 [(resources.private_key(), resources.certificate_chain())]))
         self.server.start()
-        self.stub = test_pb2.TestServiceStub(
+        self.stub = test_pb2_grpc.TestServiceStub(
             grpc.secure_channel('localhost:{}'.format(port),
                                 grpc.ssl_channel_credentials(
                                     resources.test_root_certificates()), (
diff --git a/src/python/grpcio_tests/tests/interop/methods.py b/src/python/grpcio_tests/tests/interop/methods.py
index 662ea9ce57..e1016f7c0d 100644
--- a/src/python/grpcio_tests/tests/interop/methods.py
+++ b/src/python/grpcio_tests/tests/interop/methods.py
@@ -40,7 +40,7 @@ from grpc.beta import implementations
 
 from src.proto.grpc.testing import empty_pb2
 from src.proto.grpc.testing import messages_pb2
-from src.proto.grpc.testing import test_pb2
+from src.proto.grpc.testing import test_pb2_grpc
 
 _INITIAL_METADATA_KEY = "x-grpc-test-echo-initial"
 _TRAILING_METADATA_KEY = "x-grpc-test-echo-trailing-bin"
@@ -66,7 +66,7 @@ def _maybe_echo_status_and_message(request, servicer_context):
         servicer_context.set_details(request.response_status.message)
 
 
-class TestService(test_pb2.TestServiceServicer):
+class TestService(test_pb2_grpc.TestServiceServicer):
 
     def EmptyCall(self, request, context):
         _maybe_echo_metadata(context)
diff --git a/src/python/grpcio_tests/tests/interop/server.py b/src/python/grpcio_tests/tests/interop/server.py
index 65f1604eb8..0ae2c97b42 100644
--- a/src/python/grpcio_tests/tests/interop/server.py
+++ b/src/python/grpcio_tests/tests/interop/server.py
@@ -34,7 +34,7 @@ import logging
 import time
 
 import grpc
-from src.proto.grpc.testing import test_pb2
+from src.proto.grpc.testing import test_pb2_grpc
 
 from tests.interop import methods
 from tests.interop import resources
@@ -53,7 +53,8 @@ def serve():
     args = parser.parse_args()
 
     server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
-    test_pb2.add_TestServiceServicer_to_server(methods.TestService(), server)
+    test_pb2_grpc.add_TestServiceServicer_to_server(methods.TestService(),
+                                                    server)
     if args.use_tls:
         private_key = resources.private_key()
         certificate_chain = resources.certificate_chain()
diff --git a/src/python/grpcio_tests/tests/qps/qps_worker.py b/src/python/grpcio_tests/tests/qps/qps_worker.py
index 025dfb9d4a..7cd53e7bd9 100644
--- a/src/python/grpcio_tests/tests/qps/qps_worker.py
+++ b/src/python/grpcio_tests/tests/qps/qps_worker.py
@@ -33,7 +33,7 @@ import time
 
 from concurrent import futures
 import grpc
-from src.proto.grpc.testing import services_pb2
+from src.proto.grpc.testing import services_pb2_grpc
 
 from tests.qps import worker_server
 
@@ -41,7 +41,7 @@ from tests.qps import worker_server
 def run_worker_server(port):
     server = grpc.server(futures.ThreadPoolExecutor(max_workers=5))
     servicer = worker_server.WorkerServer()
-    services_pb2.add_WorkerServiceServicer_to_server(servicer, server)
+    services_pb2_grpc.add_WorkerServiceServicer_to_server(servicer, server)
     server.add_insecure_port('[::]:{}'.format(port))
     server.start()
     servicer.wait_for_quit()
diff --git a/src/python/grpcio_tests/tests/qps/worker_server.py b/src/python/grpcio_tests/tests/qps/worker_server.py
index ca1a777611..de9535f46e 100644
--- a/src/python/grpcio_tests/tests/qps/worker_server.py
+++ b/src/python/grpcio_tests/tests/qps/worker_server.py
@@ -35,7 +35,7 @@ import time
 from concurrent import futures
 import grpc
 from src.proto.grpc.testing import control_pb2
-from src.proto.grpc.testing import services_pb2
+from src.proto.grpc.testing import services_pb2_grpc
 from src.proto.grpc.testing import stats_pb2
 
 from tests.qps import benchmark_client
@@ -45,7 +45,7 @@ from tests.qps import histogram
 from tests.unit import resources
 
 
-class WorkerServer(services_pb2.WorkerServiceServicer):
+class WorkerServer(services_pb2_grpc.WorkerServiceServicer):
     """Python Worker Server implementation."""
 
     def __init__(self):
@@ -87,8 +87,8 @@ class WorkerServer(services_pb2.WorkerServiceServicer):
             futures.ThreadPoolExecutor(max_workers=server_threads))
         if config.server_type == control_pb2.ASYNC_SERVER:
             servicer = benchmark_server.BenchmarkServer()
-            services_pb2.add_BenchmarkServiceServicer_to_server(servicer,
-                                                                server)
+            services_pb2_grpc.add_BenchmarkServiceServicer_to_server(servicer,
+                                                                     server)
         elif config.server_type == control_pb2.ASYNC_GENERIC_SERVER:
             resp_size = config.payload_config.bytebuf_params.resp_size
             servicer = benchmark_server.GenericBenchmarkServer(resp_size)
diff --git a/src/python/grpcio_tests/tests/reflection/_reflection_servicer_test.py b/src/python/grpcio_tests/tests/reflection/_reflection_servicer_test.py
index d06ff064e2..4d73be6204 100644
--- a/src/python/grpcio_tests/tests/reflection/_reflection_servicer_test.py
+++ b/src/python/grpcio_tests/tests/reflection/_reflection_servicer_test.py
@@ -34,6 +34,7 @@ import grpc
 from grpc.framework.foundation import logging_pool
 from grpc_reflection.v1alpha import reflection
 from grpc_reflection.v1alpha import reflection_pb2
+from grpc_reflection.v1alpha import reflection_pb2_grpc
 
 from google.protobuf import descriptor_pool
 from google.protobuf import descriptor_pb2
@@ -61,12 +62,12 @@ class ReflectionServicerTest(unittest.TestCase):
         server_pool = logging_pool.pool(test_constants.THREAD_CONCURRENCY)
         self._server = grpc.server(server_pool)
         port = self._server.add_insecure_port('[::]:0')
-        reflection_pb2.add_ServerReflectionServicer_to_server(servicer,
-                                                              self._server)
+        reflection_pb2_grpc.add_ServerReflectionServicer_to_server(servicer,
+                                                                   self._server)
         self._server.start()
 
         channel = grpc.insecure_channel('localhost:%d' % port)
-        self._stub = reflection_pb2.ServerReflectionStub(channel)
+        self._stub = reflection_pb2_grpc.ServerReflectionStub(channel)
 
     def testFileByName(self):
         requests = (reflection_pb2.ServerReflectionRequest(
diff --git a/src/python/grpcio_tests/tests/stress/client.py b/src/python/grpcio_tests/tests/stress/client.py
index b9dbe61d44..b7eb12bff8 100644
--- a/src/python/grpcio_tests/tests/stress/client.py
+++ b/src/python/grpcio_tests/tests/stress/client.py
@@ -34,7 +34,7 @@ import threading
 
 import grpc
 from six.moves import queue
-from src.proto.grpc.testing import metrics_pb2
+from src.proto.grpc.testing import metrics_pb2_grpc
 from src.proto.grpc.testing import test_pb2
 
 from tests.interop import methods
@@ -139,7 +139,7 @@ def run_test(args):
     runners = []
 
     server = grpc.server(futures.ThreadPoolExecutor(max_workers=25))
-    metrics_pb2.add_MetricsServiceServicer_to_server(
+    metrics_pb2_grpc.add_MetricsServiceServicer_to_server(
         metrics_server.MetricsServer(hist), server)
     server.add_insecure_port('[::]:{}'.format(args.metrics_port))
     server.start()