1 files changed, 28 insertions, 12 deletions

diff --git a/test/core/bad_client/bad_client.cc b/test/core/bad_client/bad_client.cc
index f5fb82e5f0..6a35ffcb6d 100644
--- a/test/core/bad_client/bad_client.cc
+++ b/test/core/bad_client/bad_client.cc
@@ -75,6 +75,17 @@ static void set_read_done(void* arg, grpc_error* error) {
   gpr_event_set(read_done, (void*)1);
 }
 
+/* shutdown client */
+static void shutdown_client(grpc_endpoint** client_fd) {
+  if (*client_fd != nullptr) {
+    grpc_endpoint_shutdown(
+        *client_fd, GRPC_ERROR_CREATE_FROM_STATIC_STRING("Forced Disconnect"));
+    grpc_endpoint_destroy(*client_fd);
+    grpc_core::ExecCtx::Get()->Flush();
+    *client_fd = nullptr;
+  }
+}
+
 /* Runs client side validator */
 void grpc_run_client_side_validator(grpc_bad_client_arg* arg, uint32_t flags,
                                     grpc_endpoint_pair* sfd,
@@ -115,11 +126,7 @@ void grpc_run_client_side_validator(grpc_bad_client_arg* arg, uint32_t flags,
   }
 
   if (flags & GRPC_BAD_CLIENT_DISCONNECT) {
-    grpc_endpoint_shutdown(
-        sfd->client, GRPC_ERROR_CREATE_FROM_STATIC_STRING("Forced Disconnect"));
-    grpc_endpoint_destroy(sfd->client);
-    grpc_core::ExecCtx::Get()->Flush();
-    sfd->client = nullptr;
+    shutdown_client(&sfd->client);
   }
 
   if (sfd->client != nullptr) {
@@ -158,6 +165,19 @@ void grpc_run_client_side_validator(grpc_bad_client_arg* arg, uint32_t flags,
     grpc_core::ExecCtx::Get()->Flush();
   }
 
+  /* If flags is non-zero, it is time to shutdown the client */
+  if (flags) {
+    shutdown_client(&sfd->client);
+  }
+
+  /* Make sure that the server is done writing */
+  while (!gpr_event_get(&done_write)) {
+    GPR_ASSERT(
+        grpc_completion_queue_next(
+            client_cq, grpc_timeout_milliseconds_to_deadline(100), nullptr)
+            .type == GRPC_QUEUE_TIMEOUT);
+  }
+
   grpc_slice_buffer_destroy_internal(&outgoing);
   grpc_core::ExecCtx::Get()->Flush();
 }
@@ -206,18 +226,14 @@ void grpc_run_bad_client_test(
   /* Start validator */
   gpr_thd_new(&id, "grpc_bad_client", thd_func, &a, nullptr);
   for (int i = 0; i < num_args; i++) {
-    grpc_run_client_side_validator(&args[i], flags, &sfd, client_cq);
+    grpc_run_client_side_validator(&args[i], i == (num_args - 1) ? flags : 0,
+                                   &sfd, client_cq);
   }
   /* Wait for server thread to finish */
   GPR_ASSERT(gpr_event_wait(&a.done_thd, grpc_timeout_seconds_to_deadline(1)));
 
   /* Shutdown. */
-  if (sfd.client != nullptr) {
-    grpc_endpoint_shutdown(
-        sfd.client, GRPC_ERROR_CREATE_FROM_STATIC_STRING("Test Shutdown"));
-    grpc_endpoint_destroy(sfd.client);
-    grpc_core::ExecCtx::Get()->Flush();
-  }
+  shutdown_client(&sfd.client);
 
   shutdown_cq = grpc_completion_queue_create_for_pluck(nullptr);
   grpc_server_shutdown_and_notify(a.server, shutdown_cq, nullptr);