From 0955c14608943614781d50dd59ddf88ec195cd39 Mon Sep 17 00:00:00 2001
From: Lidi Zheng <lidiz@google.com>
Date: Fri, 14 Aug 2020 13:39:36 -0700
Subject: [PATCH] Make status_code_and_message interop test robust

---
 src/python/grpcio_tests/tests/interop/BUILD.bazel    | 12 ++++++++++++
 src/python/grpcio_tests/tests/interop/methods.py     |  4 ++++
 src/python/grpcio_tests/tests_aio/interop/methods.py |  4 ++++
 3 files changed, 20 insertions(+)

diff --git a/src/python/grpcio_tests/tests/interop/BUILD.bazel b/src/python/grpcio_tests/tests/interop/BUILD.bazel
index f38b82b030a..1b300600de5 100644
--- a/src/python/grpcio_tests/tests/interop/BUILD.bazel
+++ b/src/python/grpcio_tests/tests/interop/BUILD.bazel
@@ -88,6 +88,18 @@ py_library(
     ],
 )
 
+py_binary(
+    name = "server_bin",
+    srcs = ["server.py"],
+    main = "server.py",
+    python_version = "PY3",
+    deps = [
+        ":server",
+        ":service",
+        "//src/proto/grpc/testing:test_py_pb2_grpc",
+    ],
+)
+
 py2and3_test(
     name = "_insecure_intraop_test",
     size = "small",
diff --git a/src/python/grpcio_tests/tests/interop/methods.py b/src/python/grpcio_tests/tests/interop/methods.py
index 0d8c7cb0083..44a1c38bb93 100644
--- a/src/python/grpcio_tests/tests/interop/methods.py
+++ b/src/python/grpcio_tests/tests/interop/methods.py
@@ -301,6 +301,10 @@ def _status_code_and_message(stub):
             payload=messages_pb2.Payload(body=b'\x00'),
             response_status=messages_pb2.EchoStatus(code=code, message=details))
         pipe.add(request)  # sends the initial request.
+    try:
+        next(response_iterator)
+    except grpc.RpcError as rpc_error:
+        assert rpc_error.code() == status
     # Dropping out of with block closes the pipe
     _validate_status_code_and_details(response_iterator, status, details)
 
diff --git a/src/python/grpcio_tests/tests_aio/interop/methods.py b/src/python/grpcio_tests/tests_aio/interop/methods.py
index 706f4249be3..aa39976981f 100644
--- a/src/python/grpcio_tests/tests_aio/interop/methods.py
+++ b/src/python/grpcio_tests/tests_aio/interop/methods.py
@@ -271,6 +271,10 @@ async def _status_code_and_message(stub: test_pb2_grpc.TestServiceStub):
                                                 message=details))
     await call.write(request)  # sends the initial request.
     await call.done_writing()
+    try:
+        await call.read()
+    except aio.AioRpcError as rpc_error:
+        assert rpc_error.code() == status
     await _validate_status_code_and_details(call, status, details)