1 files changed, 59 insertions, 9 deletions

diff --git a/src/components/connection_handler/test/heart_beat_monitor_test.cc b/src/components/connection_handler/test/heart_beat_monitor_test.cc
index 4c67c97191..e089a07ec4 100644
--- a/src/components/connection_handler/test/heart_beat_monitor_test.cc
+++ b/src/components/connection_handler/test/heart_beat_monitor_test.cc
@@ -51,6 +51,7 @@ namespace connection_handler_test {
 
 using ::testing::DoAll;
 using ::testing::_;
+using ::testing::Return;
 
 class HeartBeatMonitorTest : public testing::Test {
  public:
@@ -64,6 +65,8 @@ class HeartBeatMonitorTest : public testing::Test {
   uint32_t kTimeout;
   static const connection_handler::ConnectionHandle kConnectionHandle =
       0xABCDEF;
+  static const transport_manager::ConnectionUID kDefaultConnectionHandle = 1;
+  static const uint32_t kDefaultSessionId = 1;
 
   virtual void SetUp() {
     conn = new connection_handler::Connection(
@@ -80,25 +83,40 @@ ACTION_P2(RemoveSession, conn, session_id) {
 }
 
 TEST_F(HeartBeatMonitorTest, TimerNotStarted) {
+  EXPECT_CALL(connection_handler_mock, AddSession(_))
+      .WillOnce(Return(kDefaultSessionId));
+  EXPECT_CALL(connection_handler_mock, RemoveSession(kDefaultSessionId))
+      .WillOnce(Return(true));  // called by destructor of Connection
+
   // Whithout StartHeartBeat nothing to be call
   EXPECT_CALL(connection_handler_mock, CloseSession(_, _)).Times(0);
   EXPECT_CALL(connection_handler_mock, CloseConnection(_)).Times(0);
   EXPECT_CALL(connection_handler_mock, SendHeartBeat(_, _)).Times(0);
 
-  conn->AddNewSession();
+  conn->AddNewSession(kDefaultConnectionHandle);
 }
 
 TEST_F(HeartBeatMonitorTest, TimerNotElapsed) {
+  EXPECT_CALL(connection_handler_mock, AddSession(_))
+      .WillOnce(Return(kDefaultSessionId));
+  EXPECT_CALL(connection_handler_mock, RemoveSession(kDefaultSessionId))
+      .WillOnce(Return(true));
+
   EXPECT_CALL(connection_handler_mock, SendHeartBeat(_, _)).Times(0);
   EXPECT_CALL(connection_handler_mock, CloseSession(_, _)).Times(0);
   EXPECT_CALL(connection_handler_mock, CloseConnection(_)).Times(0);
 
-  const uint32_t session = conn->AddNewSession();
+  const uint32_t session = conn->AddNewSession(kDefaultConnectionHandle);
   conn->StartHeartBeat(session);
 }
 
 TEST_F(HeartBeatMonitorTest, TimerElapsed) {
-  const uint32_t session = conn->AddNewSession();
+  EXPECT_CALL(connection_handler_mock, AddSession(_))
+      .WillOnce(Return(kDefaultSessionId));
+  EXPECT_CALL(connection_handler_mock, RemoveSession(kDefaultSessionId))
+      .WillOnce(Return(true));  // invoked by RemoveSession action
+
+  const uint32_t session = conn->AddNewSession(kDefaultConnectionHandle);
 
   TestAsyncWaiter waiter;
   uint32_t times = 0;
@@ -121,11 +139,16 @@ TEST_F(HeartBeatMonitorTest, TimerElapsed) {
 }
 
 TEST_F(HeartBeatMonitorTest, KeptAlive) {
+  EXPECT_CALL(connection_handler_mock, AddSession(_))
+      .WillOnce(Return(kDefaultSessionId));
+  EXPECT_CALL(connection_handler_mock, RemoveSession(kDefaultSessionId))
+      .WillOnce(Return(true));
+
   EXPECT_CALL(connection_handler_mock, CloseSession(_, _)).Times(0);
   EXPECT_CALL(connection_handler_mock, CloseConnection(_)).Times(0);
   EXPECT_CALL(connection_handler_mock, SendHeartBeat(_, _)).Times(0);
 
-  const uint32_t session = conn->AddNewSession();
+  const uint32_t session = conn->AddNewSession(kDefaultConnectionHandle);
   conn->StartHeartBeat(session);
   usleep(kTimeout * MICROSECONDS_IN_MILLISECONDS - MICROSECONDS_IN_SECOND);
   conn->KeepAlive(session);
@@ -137,7 +160,12 @@ TEST_F(HeartBeatMonitorTest, KeptAlive) {
 }
 
 TEST_F(HeartBeatMonitorTest, NotKeptAlive) {
-  const uint32_t session = conn->AddNewSession();
+  EXPECT_CALL(connection_handler_mock, AddSession(_))
+      .WillOnce(Return(kDefaultSessionId));
+  EXPECT_CALL(connection_handler_mock, RemoveSession(kDefaultSessionId))
+      .WillOnce(Return(true));
+
+  const uint32_t session = conn->AddNewSession(kDefaultConnectionHandle);
 
   TestAsyncWaiter waiter;
   uint32_t times = 0;
@@ -167,8 +195,20 @@ TEST_F(HeartBeatMonitorTest, NotKeptAlive) {
 }
 
 TEST_F(HeartBeatMonitorTest, TwoSessionsElapsed) {
-  const uint32_t kSession1 = conn->AddNewSession();
-  const uint32_t kSession2 = conn->AddNewSession();
+  const uint32_t kMockSessionId1 = 1;
+  const uint32_t kMockSessionId2 = 2;
+  EXPECT_CALL(connection_handler_mock, AddSession(_))
+      .WillOnce(Return(kMockSessionId1))
+      .WillOnce(Return(kMockSessionId2));
+  EXPECT_CALL(connection_handler_mock, RemoveSession(kMockSessionId1))
+      .WillOnce(Return(true));
+  EXPECT_CALL(connection_handler_mock, RemoveSession(kMockSessionId2))
+      .WillOnce(Return(true));
+
+  const uint32_t kSession1 = conn->AddNewSession(kDefaultConnectionHandle);
+
+  const transport_manager::ConnectionUID kAnotherConnectionHandle = 2;
+  const uint32_t kSession2 = conn->AddNewSession(kAnotherConnectionHandle);
 
   TestAsyncWaiter waiter;
   uint32_t times = 0;
@@ -199,7 +239,12 @@ TEST_F(HeartBeatMonitorTest, TwoSessionsElapsed) {
 }
 
 TEST_F(HeartBeatMonitorTest, IncreaseHeartBeatTimeout) {
-  const uint32_t kSession = conn->AddNewSession();
+  EXPECT_CALL(connection_handler_mock, AddSession(_))
+      .WillOnce(Return(kDefaultSessionId));
+  EXPECT_CALL(connection_handler_mock, RemoveSession(kDefaultSessionId))
+      .WillOnce(Return(true));
+
+  const uint32_t kSession = conn->AddNewSession(kDefaultConnectionHandle);
 
   EXPECT_CALL(connection_handler_mock, CloseSession(_, _)).Times(0);
   EXPECT_CALL(connection_handler_mock, CloseConnection(_)).Times(0);
@@ -211,7 +256,12 @@ TEST_F(HeartBeatMonitorTest, IncreaseHeartBeatTimeout) {
 }
 
 TEST_F(HeartBeatMonitorTest, DecreaseHeartBeatTimeout) {
-  const uint32_t kSession = conn->AddNewSession();
+  EXPECT_CALL(connection_handler_mock, AddSession(_))
+      .WillOnce(Return(kDefaultSessionId));
+  EXPECT_CALL(connection_handler_mock, RemoveSession(kDefaultSessionId))
+      .WillOnce(Return(true));
+
+  const uint32_t kSession = conn->AddNewSession(kDefaultConnectionHandle);
 
   TestAsyncWaiter waiter;
   uint32_t times = 0;