From 162fca110c7a5ca213fc9ea70ae865b7227dcaa4 Mon Sep 17 00:00:00 2001
From: yang-g <yangg@google.com>
Date: Wed, 1 Jul 2020 09:39:14 -0700
Subject: [PATCH] Fix message_allocator_e2e_test

---
 .../end2end/message_allocator_end2end_test.cc | 22 ++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/test/cpp/end2end/message_allocator_end2end_test.cc b/test/cpp/end2end/message_allocator_end2end_test.cc
index 3cae255d43e..b34a08ef281 100644
--- a/test/cpp/end2end/message_allocator_end2end_test.cc
+++ b/test/cpp/end2end/message_allocator_end2end_test.cc
@@ -143,7 +143,13 @@ class MessageAllocatorEnd2endTestBase
     builder.RegisterService(&callback_service_);
 
     server_ = builder.BuildAndStart();
-    is_server_started_ = true;
+  }
+
+  void DestroyServer() {
+    if (server_) {
+      server_->Shutdown();
+      server_.reset();
+    }
   }
 
   void ResetStub() {
@@ -165,9 +171,7 @@ class MessageAllocatorEnd2endTestBase
   }
 
   void TearDown() override {
-    if (is_server_started_) {
-      server_->Shutdown();
-    }
+    DestroyServer();
     if (picked_port_ > 0) {
       grpc_recycle_unused_port(picked_port_);
     }
@@ -206,7 +210,6 @@ class MessageAllocatorEnd2endTestBase
   }
 
   bool do_not_test_{false};
-  bool is_server_started_{false};
   int picked_port_{0};
   std::shared_ptr<Channel> channel_;
   std::unique_ptr<EchoTestService::Stub> stub_;
@@ -280,6 +283,9 @@ TEST_P(SimpleAllocatorTest, SimpleRpc) {
   CreateServer(allocator.get());
   ResetStub();
   SendRpcs(kRpcCount);
+  // messages_deallocaton_count is updated in Release after server side OnDone.
+  // Destroy server to make sure it has been updated.
+  DestroyServer();
   EXPECT_EQ(kRpcCount, allocator->allocation_count);
   EXPECT_EQ(kRpcCount, allocator->messages_deallocation_count);
   EXPECT_EQ(0, allocator->request_deallocation_count);
@@ -302,6 +308,9 @@ TEST_P(SimpleAllocatorTest, RpcWithEarlyFreeRequest) {
   CreateServer(allocator.get());
   ResetStub();
   SendRpcs(kRpcCount);
+  // messages_deallocaton_count is updated in Release after server side OnDone.
+  // Destroy server to make sure it has been updated.
+  DestroyServer();
   EXPECT_EQ(kRpcCount, allocator->allocation_count);
   EXPECT_EQ(kRpcCount, allocator->messages_deallocation_count);
   EXPECT_EQ(kRpcCount, allocator->request_deallocation_count);
@@ -326,6 +335,9 @@ TEST_P(SimpleAllocatorTest, RpcWithReleaseRequest) {
   CreateServer(allocator.get());
   ResetStub();
   SendRpcs(kRpcCount);
+  // messages_deallocaton_count is updated in Release after server side OnDone.
+  // Destroy server to make sure it has been updated.
+  DestroyServer();
   EXPECT_EQ(kRpcCount, allocator->allocation_count);
   EXPECT_EQ(kRpcCount, allocator->messages_deallocation_count);
   EXPECT_EQ(0, allocator->request_deallocation_count);