Fix protocol handler unit test

1 files changed, 4 insertions, 2 deletions

diff --git a/src/components/protocol_handler/test/protocol_handler_tm_test.cc b/src/components/protocol_handler/test/protocol_handler_tm_test.cc
index 77de1705da..8525b1c5b9 100644
--- a/src/components/protocol_handler/test/protocol_handler_tm_test.cc
+++ b/src/components/protocol_handler/test/protocol_handler_tm_test.cc
@@ -1989,7 +1989,8 @@ TEST_F(ProtocolHandlerImplTest, SendHeartBeat_Successful) {
       transport_manager_mock,
       SendMessageToDevice(ExpectedMessage(
           FRAME_TYPE_CONTROL, FRAME_DATA_HEART_BEAT, PROTECTION_OFF, kControl)))
-      .WillOnce(Return(E_SUCCESS));
+      .WillOnce(DoAll(NotifyTestAsyncWaiter(waiter), Return(E_SUCCESS)));
+  times++;
 
   // Act
   protocol_handler_impl->SendHeartBeat(connection_id, session_id);
@@ -2086,7 +2087,8 @@ TEST_F(ProtocolHandlerImplTest,
       transport_manager_mock,
       SendMessageToDevice(ExpectedMessage(
           FRAME_TYPE_SINGLE, FRAME_DATA_SINGLE, PROTECTION_OFF, kControl)))
-      .WillOnce(Return(E_SUCCESS));
+      .WillOnce(DoAll(NotifyTestAsyncWaiter(waiter), Return(E_SUCCESS)));
+  times++;
 
   // Act
   protocol_handler_impl->SendMessageToMobileApp(message, is_final);