Simplify call_test using mixin class

src/python/grpcio_tests/tests_aio/unit/call_test.py

@@ -14,19 +14,17 @@
 """Tests behavior of the grpc.aio.UnaryUnaryCall class."""
 
 import asyncio
+import datetime
 import logging
 import unittest
-import datetime
 
 import grpc
-
 from grpc.experimental import aio
-from src.proto.grpc.testing import messages_pb2
+
-from src.proto.grpc.testing import test_pb2_grpc
+from src.proto.grpc.testing import messages_pb2, test_pb2_grpc
 from tests.unit.framework.common import test_constants
-from tests_aio.unit._test_server import start_test_server
 from tests_aio.unit._test_base import AioTestBase
-from src.proto.grpc.testing import messages_pb2
+from tests_aio.unit._test_server import start_test_server
 
 _NUM_STREAM_RESPONSES = 5
 _RESPONSE_PAYLOAD_SIZE = 42
@@ -37,21 +35,22 @@ _UNREACHABLE_TARGET = '0.1:1111'
 _INFINITE_INTERVAL_US = 2**31 - 1
 
 
-class TestUnaryUnaryCall(AioTestBase):
+class _MulticallableTestMixin():
 
     async def setUp(self):
-        self._server_target, self._server = await start_test_server()
+        address, self._server = await start_test_server()
+        self._channel = aio.insecure_channel(address)
+        self._stub = test_pb2_grpc.TestServiceStub(self._channel)
 
     async def tearDown(self):
+        await self._channel.close()
         await self._server.stop(None)
 
+
+class TestUnaryUnaryCall(_MulticallableTestMixin, AioTestBase):
+
     async def test_call_ok(self):
-        async with aio.insecure_channel(self._server_target) as channel:
+        call = self._stub.UnaryCall(messages_pb2.SimpleRequest())
-            hi = channel.unary_unary(
-                '/grpc.testing.TestService/UnaryCall',
-                request_serializer=messages_pb2.SimpleRequest.SerializeToString,
-                response_deserializer=messages_pb2.SimpleResponse.FromString)
-            call = hi(messages_pb2.SimpleRequest())
 
         self.assertFalse(call.done())
 
@@ -68,13 +67,9 @@ class TestUnaryUnaryCall(AioTestBase):
 
     async def test_call_rpc_error(self):
         async with aio.insecure_channel(_UNREACHABLE_TARGET) as channel:
-            hi = channel.unary_unary(
+            stub = test_pb2_grpc.TestServiceStub(channel)
-                '/grpc.testing.TestService/UnaryCall',
-                request_serializer=messages_pb2.SimpleRequest.SerializeToString,
-                response_deserializer=messages_pb2.SimpleResponse.FromString,
-            )
 
-            call = hi(messages_pb2.SimpleRequest(), timeout=0.1)
+            call = stub.UnaryCall(messages_pb2.SimpleRequest(), timeout=0.1)
 
             with self.assertRaises(grpc.RpcError) as exception_context:
                 await call
@@ -95,48 +90,23 @@ class TestUnaryUnaryCall(AioTestBase):
                           exception_context_retry.exception)
 
     async def test_call_code_awaitable(self):
-        async with aio.insecure_channel(self._server_target) as channel:
+        call = self._stub.UnaryCall(messages_pb2.SimpleRequest())
-            hi = channel.unary_unary(
-                '/grpc.testing.TestService/UnaryCall',
-                request_serializer=messages_pb2.SimpleRequest.SerializeToString,
-                response_deserializer=messages_pb2.SimpleResponse.FromString)
-            call = hi(messages_pb2.SimpleRequest())
         self.assertEqual(await call.code(), grpc.StatusCode.OK)
 
     async def test_call_details_awaitable(self):
-        async with aio.insecure_channel(self._server_target) as channel:
+        call = self._stub.UnaryCall(messages_pb2.SimpleRequest())
-            hi = channel.unary_unary(
-                '/grpc.testing.TestService/UnaryCall',
-                request_serializer=messages_pb2.SimpleRequest.SerializeToString,
-                response_deserializer=messages_pb2.SimpleResponse.FromString)
-            call = hi(messages_pb2.SimpleRequest())
         self.assertEqual('', await call.details())
 
     async def test_call_initial_metadata_awaitable(self):
-        async with aio.insecure_channel(self._server_target) as channel:
+        call = self._stub.UnaryCall(messages_pb2.SimpleRequest())
-            hi = channel.unary_unary(
-                '/grpc.testing.TestService/UnaryCall',
-                request_serializer=messages_pb2.SimpleRequest.SerializeToString,
-                response_deserializer=messages_pb2.SimpleResponse.FromString)
-            call = hi(messages_pb2.SimpleRequest())
         self.assertEqual((), await call.initial_metadata())
 
     async def test_call_trailing_metadata_awaitable(self):
-        async with aio.insecure_channel(self._server_target) as channel:
+        call = self._stub.UnaryCall(messages_pb2.SimpleRequest())
-            hi = channel.unary_unary(
-                '/grpc.testing.TestService/UnaryCall',
-                request_serializer=messages_pb2.SimpleRequest.SerializeToString,
-                response_deserializer=messages_pb2.SimpleResponse.FromString)
-            call = hi(messages_pb2.SimpleRequest())
         self.assertEqual((), await call.trailing_metadata())
 
     async def test_cancel_unary_unary(self):
-        async with aio.insecure_channel(self._server_target) as channel:
+        call = self._stub.UnaryCall(messages_pb2.SimpleRequest())
-            hi = channel.unary_unary(
-                '/grpc.testing.TestService/UnaryCall',
-                request_serializer=messages_pb2.SimpleRequest.SerializeToString,
-                response_deserializer=messages_pb2.SimpleResponse.FromString)
-            call = hi(messages_pb2.SimpleRequest())
 
         self.assertFalse(call.cancelled())
 
@@ -153,10 +123,8 @@ class TestUnaryUnaryCall(AioTestBase):
                             'Locally cancelled by application!')
 
     async def test_cancel_unary_unary_in_task(self):
-        async with aio.insecure_channel(self._server_target) as channel:
-            stub = test_pb2_grpc.TestServiceStub(channel)
         coro_started = asyncio.Event()
-            call = stub.EmptyCall(messages_pb2.SimpleRequest())
+        call = self._stub.EmptyCall(messages_pb2.SimpleRequest())
 
         async def another_coro():
             coro_started.set()
@@ -174,18 +142,9 @@ class TestUnaryUnaryCall(AioTestBase):
             await task
 
 
-class TestUnaryStreamCall(AioTestBase):
+class TestUnaryStreamCall(_MulticallableTestMixin, AioTestBase):
-
-    async def setUp(self):
-        self._server_target, self._server = await start_test_server()
-
-    async def tearDown(self):
-        await self._server.stop(None)
 
     async def test_cancel_unary_stream(self):
-        async with aio.insecure_channel(self._server_target) as channel:
-            stub = test_pb2_grpc.TestServiceStub(channel)
-
         # Prepares the request
         request = messages_pb2.StreamingOutputCallRequest()
         for _ in range(_NUM_STREAM_RESPONSES):
@@ -196,7 +155,7 @@ class TestUnaryStreamCall(AioTestBase):
                 ))
 
         # Invokes the actual RPC
-            call = stub.StreamingOutputCall(request)
+        call = self._stub.StreamingOutputCall(request)
         self.assertFalse(call.cancelled())
 
         response = await call.read()
@@ -215,9 +174,6 @@ class TestUnaryStreamCall(AioTestBase):
         self.assertTrue(call.cancelled())
 
     async def test_multiple_cancel_unary_stream(self):
-        async with aio.insecure_channel(self._server_target) as channel:
-            stub = test_pb2_grpc.TestServiceStub(channel)
-
         # Prepares the request
         request = messages_pb2.StreamingOutputCallRequest()
         for _ in range(_NUM_STREAM_RESPONSES):
@@ -228,7 +184,7 @@ class TestUnaryStreamCall(AioTestBase):
                 ))
 
         # Invokes the actual RPC
-            call = stub.StreamingOutputCall(request)
+        call = self._stub.StreamingOutputCall(request)
         self.assertFalse(call.cancelled())
 
         response = await call.read()
@@ -246,9 +202,6 @@ class TestUnaryStreamCall(AioTestBase):
 
     async def test_early_cancel_unary_stream(self):
         """Test cancellation before receiving messages."""
-        async with aio.insecure_channel(self._server_target) as channel:
-            stub = test_pb2_grpc.TestServiceStub(channel)
-
         # Prepares the request
         request = messages_pb2.StreamingOutputCallRequest()
         for _ in range(_NUM_STREAM_RESPONSES):
@@ -259,7 +212,7 @@ class TestUnaryStreamCall(AioTestBase):
                 ))
 
         # Invokes the actual RPC
-            call = stub.StreamingOutputCall(request)
+        call = self._stub.StreamingOutputCall(request)
 
         self.assertFalse(call.cancelled())
         self.assertTrue(call.cancel())
@@ -276,9 +229,6 @@ class TestUnaryStreamCall(AioTestBase):
 
     async def test_late_cancel_unary_stream(self):
         """Test cancellation after received all messages."""
-        async with aio.insecure_channel(self._server_target) as channel:
-            stub = test_pb2_grpc.TestServiceStub(channel)
-
         # Prepares the request
         request = messages_pb2.StreamingOutputCallRequest()
         for _ in range(_NUM_STREAM_RESPONSES):
@@ -287,7 +237,7 @@ class TestUnaryStreamCall(AioTestBase):
                     size=_RESPONSE_PAYLOAD_SIZE,))
 
         # Invokes the actual RPC
-            call = stub.StreamingOutputCall(request)
+        call = self._stub.StreamingOutputCall(request)
 
         for _ in range(_NUM_STREAM_RESPONSES):
             response = await call.read()
@@ -305,9 +255,6 @@ class TestUnaryStreamCall(AioTestBase):
 
     async def test_too_many_reads_unary_stream(self):
         """Test calling read after received all messages fails."""
-        async with aio.insecure_channel(self._server_target) as channel:
-            stub = test_pb2_grpc.TestServiceStub(channel)
-
         # Prepares the request
         request = messages_pb2.StreamingOutputCallRequest()
         for _ in range(_NUM_STREAM_RESPONSES):
@@ -316,7 +263,7 @@ class TestUnaryStreamCall(AioTestBase):
                     size=_RESPONSE_PAYLOAD_SIZE,))
 
         # Invokes the actual RPC
-            call = stub.StreamingOutputCall(request)
+        call = self._stub.StreamingOutputCall(request)
 
         for _ in range(_NUM_STREAM_RESPONSES):
             response = await call.read()
@@ -332,9 +279,6 @@ class TestUnaryStreamCall(AioTestBase):
 
     async def test_unary_stream_async_generator(self):
         """Sunny day test case for unary_stream."""
-        async with aio.insecure_channel(self._server_target) as channel:
-            stub = test_pb2_grpc.TestServiceStub(channel)
-
         # Prepares the request
         request = messages_pb2.StreamingOutputCallRequest()
         for _ in range(_NUM_STREAM_RESPONSES):
@@ -343,7 +287,7 @@ class TestUnaryStreamCall(AioTestBase):
                     size=_RESPONSE_PAYLOAD_SIZE,))
 
         # Invokes the actual RPC
-            call = stub.StreamingOutputCall(request)
+        call = self._stub.StreamingOutputCall(request)
         self.assertFalse(call.cancelled())
 
         async for response in call:
@@ -355,8 +299,6 @@ class TestUnaryStreamCall(AioTestBase):
         self.assertEqual(await call.code(), grpc.StatusCode.OK)
 
     async def test_cancel_unary_stream_in_task_using_read(self):
-        async with aio.insecure_channel(self._server_target) as channel:
-            stub = test_pb2_grpc.TestServiceStub(channel)
         coro_started = asyncio.Event()
 
         # Configs the server method to block forever
@@ -368,7 +310,7 @@ class TestUnaryStreamCall(AioTestBase):
             ))
 
         # Invokes the actual RPC
-            call = stub.StreamingOutputCall(request)
+        call = self._stub.StreamingOutputCall(request)
 
         async def another_coro():
             coro_started.set()
@@ -386,8 +328,6 @@ class TestUnaryStreamCall(AioTestBase):
             await task
 
     async def test_cancel_unary_stream_in_task_using_async_for(self):
-        async with aio.insecure_channel(self._server_target) as channel:
-            stub = test_pb2_grpc.TestServiceStub(channel)
         coro_started = asyncio.Event()
 
         # Configs the server method to block forever
@@ -399,7 +339,7 @@ class TestUnaryStreamCall(AioTestBase):
             ))
 
         # Invokes the actual RPC
-            call = stub.StreamingOutputCall(request)
+        call = self._stub.StreamingOutputCall(request)
 
         async def another_coro():
             coro_started.set()
@@ -445,16 +385,7 @@ class TestUnaryStreamCall(AioTestBase):
         self.loop.run_until_complete(coro())
 
 
-class TestStreamUnaryCall(AioTestBase):
+class TestStreamUnaryCall(_MulticallableTestMixin, AioTestBase):
-
-    async def setUp(self):
-        self._server_target, self._server = await start_test_server()
-        self._channel = aio.insecure_channel(self._server_target)
-        self._stub = test_pb2_grpc.TestServiceStub(self._channel)
-
-    async def tearDown(self):
-        await self._channel.close()
-        await self._server.stop(None)
 
     async def test_cancel_stream_unary(self):
         call = self._stub.StreamingInputCall()
@@ -564,16 +495,7 @@ _STREAM_OUTPUT_REQUEST_ONE_RESPONSE.response_parameters.append(
     messages_pb2.ResponseParameters(size=_RESPONSE_PAYLOAD_SIZE))
 
 
-class TestStreamStreamCall(AioTestBase):
+class TestStreamStreamCall(_MulticallableTestMixin, AioTestBase):
-
-    async def setUp(self):
-        self._server_target, self._server = await start_test_server()
-        self._channel = aio.insecure_channel(self._server_target)
-        self._stub = test_pb2_grpc.TestServiceStub(self._channel)
-
-    async def tearDown(self):
-        await self._channel.close()
-        await self._server.stop(None)
 
     async def test_cancel(self):
         # Invokes the actual RPC