Merge pull request #1965 from AByzhynar/fix/Wrong_connection_closure

Fix/wrong connection closure

4 files changed, 138 insertions, 145 deletions

diff --git a/src/components/connection_handler/src/connection.cc b/src/components/connection_handler/src/connection.cc
index 5a35919c6c..3f5b0e1901 100644
--- a/src/components/connection_handler/src/connection.cc
+++ b/src/components/connection_handler/src/connection.cc
@@ -430,7 +430,6 @@ const SessionMap Connection::session_map() const {
 }
 
 void Connection::CloseSession(uint8_t session_id) {
-  size_t size;
   {
     sync_primitives::AutoLock lock(session_map_lock_);
 
@@ -438,16 +437,10 @@ void Connection::CloseSession(uint8_t session_id) {
     if (session_it == session_map_.end()) {
       return;
     }
-    size = session_map_.size();
   }
 
   connection_handler_->CloseSession(
       connection_handle_, session_id, connection_handler::kCommon);
-
-  // Close connection if it is last session
-  if (1 == size) {
-    connection_handler_->CloseConnection(connection_handle_);
-  }
 }
 
 void Connection::UpdateProtocolVersionSession(uint8_t session_id,
diff --git a/src/components/connection_handler/src/heartbeat_monitor.cc b/src/components/connection_handler/src/heartbeat_monitor.cc
index 12994fd23a..4c2d5899af 100644
--- a/src/components/connection_handler/src/heartbeat_monitor.cc
+++ b/src/components/connection_handler/src/heartbeat_monitor.cc
@@ -89,18 +89,20 @@ void HeartBeatMonitor::threadMain() {
 }
 
 void HeartBeatMonitor::AddSession(uint8_t session_id) {
-  LOG4CXX_DEBUG(logger_,
-                "Add session with id " << static_cast<int32_t>(session_id));
+  const uint32_t converted_session_id = static_cast<int32_t>(session_id);
+  UNUSED(converted_session_id);
+  LOG4CXX_DEBUG(logger_, "Add session with id " << converted_session_id);
   AutoLock auto_lock(sessions_list_lock_);
   if (sessions_.end() != sessions_.find(session_id)) {
     LOG4CXX_WARN(logger_,
-                 "Session with id " << static_cast<int32_t>(session_id)
-                                    << " already exists");
+                 "Session with id: " << converted_session_id
+                                     << " already exists");
     return;
   }
   sessions_.insert(
       std::make_pair(session_id, SessionState(default_heartbeat_timeout_)));
-  LOG4CXX_INFO(logger_, "Start heartbeat for session " << session_id);
+  LOG4CXX_INFO(logger_,
+               "Start heartbeat for session: " << converted_session_id);
 }
 
 void HeartBeatMonitor::RemoveSession(uint8_t session_id) {
diff --git a/src/components/connection_handler/test/connection_handler_impl_test.cc b/src/components/connection_handler/test/connection_handler_impl_test.cc
index f66d6b350c..64de24b103 100644
--- a/src/components/connection_handler/test/connection_handler_impl_test.cc
+++ b/src/components/connection_handler/test/connection_handler_impl_test.cc
@@ -88,7 +88,7 @@ class ConnectionHandlerTest : public ::testing::Test {
     mac_address_ = "test_address";
 
     connection_handler_ = new ConnectionHandlerImpl(
-        mock_connection_handler_settings, mock_transport_manager_);
+        mock_connection_handler_settings_, mock_transport_manager_);
     uid_ = 1u;
     connection_key_ = connection_handler_->KeyFromPair(0, 0u);
     protected_services_.clear();
@@ -101,23 +101,29 @@ class ConnectionHandlerTest : public ::testing::Test {
 
   void SetSpecificServices() {
 #ifdef ENABLE_SECURITY
-    ON_CALL(mock_connection_handler_settings, force_protected_service())
+    ON_CALL(mock_connection_handler_settings_, force_protected_service())
         .WillByDefault(ReturnRefOfCopy(protected_services_));
-    ON_CALL(mock_connection_handler_settings, force_unprotected_service())
+    ON_CALL(mock_connection_handler_settings_, force_unprotected_service())
         .WillByDefault(ReturnRefOfCopy(unprotected_services_));
 #endif  // ENABLE_SECURITY
   }
+
   // Additional SetUp
   void AddTestDeviceConnection() {
     const transport_manager::DeviceInfo device_info(
         device_handle_, mac_address_, device_name_, connection_type_);
     // Add Device and connection
-    ON_CALL(mock_connection_handler_settings, heart_beat_timeout())
+    ON_CALL(mock_connection_handler_settings_, heart_beat_timeout())
         .WillByDefault(Return(1000u));
     connection_handler_->addDeviceConnection(device_info, uid_);
     connection_key_ = connection_handler_->KeyFromPair(uid_, 0u);
     // Remove all specific services
   }
+
+  /**
+   * @brief Adds test session
+   * @return session_id
+   */
   void AddTestSession() {
     protocol_handler_test::MockProtocolHandler temp_protocol_handler;
     connection_handler_->set_protocol_handler(&temp_protocol_handler);
@@ -134,9 +140,11 @@ class ConnectionHandlerTest : public ::testing::Test {
         connection_handler_->KeyFromPair(uid_, out_context_.new_session_id_);
     CheckSessionExists(uid_, out_context_.new_session_id_);
   }
+
   uint32_t SessionHash(const uint32_t connection, const uint32_t session) {
     return connection_handler_->KeyFromPair(connection, session);
   }
+
   void AddTestService(ServiceType service_type) {
     EXPECT_NE(0u, out_context_.new_session_id_);
     EXPECT_EQ(SessionHash(uid_, out_context_.new_session_id_),
@@ -339,7 +347,7 @@ class ConnectionHandlerTest : public ::testing::Test {
   testing::NiceMock<transport_manager_test::MockTransportManager>
       mock_transport_manager_;
   testing::NiceMock<MockConnectionHandlerSettings>
-      mock_connection_handler_settings;
+      mock_connection_handler_settings_;
   protocol_handler_test::MockProtocolHandler mock_protocol_handler_;
   transport_manager::DeviceHandle device_handle_;
   transport_manager::ConnectionUID uid_;
@@ -404,23 +412,23 @@ TEST_F(ConnectionHandlerTest, GetAppIdOnSessionKey) {
   AddTestSession();
 
   uint32_t app_id = 0;
-  const uint32_t testid = SessionHash(uid_, out_context_.new_session_id_);
+  const uint32_t test_id = SessionHash(uid_, out_context_.new_session_id_);
 
-  connection_handler::DeviceHandle null_handle = 0;
+  connection_handler::DeviceHandle* device_id = NULL;
   EXPECT_EQ(0,
             connection_handler_->GetDataOnSessionKey(
-                connection_key_, &app_id, NULL, &null_handle));
-  EXPECT_EQ(testid, app_id);
+                connection_key_, &app_id, NULL, device_id));
+  EXPECT_EQ(test_id, app_id);
 }
 
 TEST_F(ConnectionHandlerTest, GetAppIdOnSessionKey_SessionNotStarted) {
   AddTestDeviceConnection();
 
   uint32_t app_id = 0;
-  connection_handler::DeviceHandle null_handle = 0;
+  connection_handler::DeviceHandle* device_id = NULL;
   EXPECT_EQ(-1,
             connection_handler_->GetDataOnSessionKey(
-                connection_key_, &app_id, NULL, &null_handle));
+                connection_key_, &app_id, NULL, device_id));
 }
 
 TEST_F(ConnectionHandlerTest, GetDeviceID) {
@@ -528,6 +536,13 @@ TEST_F(ConnectionHandlerTest, GetPairFromKey) {
   EXPECT_EQ(out_context_.new_session_id_, session_id);
 }
 
+MATCHER_P(SameDevice, device, "") {
+  return arg.device_handle() == device.device_handle() &&
+         arg.user_friendly_name() == device.user_friendly_name() &&
+         arg.mac_address() == device.mac_address() &&
+         arg.connection_type() == device.connection_type();
+}
+
 TEST_F(ConnectionHandlerTest, IsHeartBeatSupported) {
   AddTestDeviceConnection();
   AddTestSession();
@@ -536,13 +551,6 @@ TEST_F(ConnectionHandlerTest, IsHeartBeatSupported) {
       uid_, out_context_.new_session_id_));
 }
 
-MATCHER_P(SameDevice, device, "") {
-  return arg.device_handle() == device.device_handle() &&
-         arg.user_friendly_name() == device.user_friendly_name() &&
-         arg.mac_address() == device.mac_address() &&
-         arg.connection_type() == device.connection_type();
-}
-
 TEST_F(ConnectionHandlerTest, SendEndServiceWithoutSetProtocolHandler) {
   AddTestDeviceConnection();
   AddTestSession();
@@ -1265,8 +1273,8 @@ TEST_F(ConnectionHandlerTest, SessionStop_CheckSpecificHash) {
   for (uint32_t session = 0; session < 0xFF; ++session) {
     AddTestSession();
 
-    uint32_t wrong_hash = protocol_handler::HASH_ID_WRONG;
-    uint32_t hash = protocol_handler::HASH_ID_NOT_SUPPORTED;
+    uint32_t hash = connection_key_;
+    uint32_t wrong_hash = hash + 1;
 
     const uint32_t end_audio_wrong_hash =
         connection_handler_->OnSessionEndedCallback(
@@ -1494,7 +1502,7 @@ TEST_F(ConnectionHandlerTest, ServiceStarted_Video_Multiple) {
 
 TEST_F(ConnectionHandlerTest,
        SessionStarted_StartSession_SecureSpecific_Unprotect) {
-  EXPECT_CALL(mock_connection_handler_settings, heart_beat_timeout())
+  EXPECT_CALL(mock_connection_handler_settings_, heart_beat_timeout())
       .WillOnce(Return(heartbeat_timeout));
   // Add virtual device and connection
   AddTestDeviceConnection();
@@ -2049,7 +2057,7 @@ TEST_F(ConnectionHandlerTest, StartStopSecondarySession) {
       device_handle_, mac_address_, device_name_, std::string("WIFI"));
   const transport_manager::ConnectionUID secondary_uid = 2u;
   // Add Device and connection
-  ON_CALL(mock_connection_handler_settings, heart_beat_timeout())
+  ON_CALL(mock_connection_handler_settings_, heart_beat_timeout())
       .WillByDefault(Return(1000u));
   connection_handler_->addDeviceConnection(secondary_device_info,
                                            secondary_uid);
@@ -2106,7 +2114,7 @@ TEST_F(ConnectionHandlerTest, StopSecondarySession_NoService) {
   const transport_manager::DeviceInfo secondary_device_info(
       device_handle_, mac_address_, device_name_, std::string("WIFI"));
   const transport_manager::ConnectionUID secondary_uid = 123u;
-  ON_CALL(mock_connection_handler_settings, heart_beat_timeout())
+  ON_CALL(mock_connection_handler_settings_, heart_beat_timeout())
       .WillByDefault(Return(1000u));
   connection_handler_->addDeviceConnection(secondary_device_info,
                                            secondary_uid);
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 e089a07ec4..ccfff34116 100644
--- a/src/components/connection_handler/test/heart_beat_monitor_test.cc
+++ b/src/components/connection_handler/test/heart_beat_monitor_test.cc
@@ -43,6 +43,7 @@ namespace {
 const int32_t MILLISECONDS_IN_SECOND = 1000;
 const int32_t MICROSECONDS_IN_MILLISECONDS = 1000;
 const int32_t MICROSECONDS_IN_SECOND = 1000 * 1000;
+const uint32_t kTime_offset = 20u;
 }
 
 namespace test {
@@ -55,26 +56,30 @@ using ::testing::Return;
 
 class HeartBeatMonitorTest : public testing::Test {
  public:
-  HeartBeatMonitorTest() : conn(NULL) {
-    kTimeout = 5000u;
-  }
+  HeartBeatMonitorTest()
+      : connection_(NULL)
+      , timeout_(100u)
+      , time_allowance_multiplier_(1.5)
+      , expiration_timeout_(timeout_ * 2u) {}
 
  protected:
-  testing::NiceMock<MockConnectionHandler> connection_handler_mock;
-  connection_handler::Connection* conn;
-  uint32_t kTimeout;
-  static const connection_handler::ConnectionHandle kConnectionHandle =
+  testing::NiceMock<MockConnectionHandler> connection_handler_mock_;
+  connection_handler::Connection* connection_;
+  const uint32_t timeout_;
+  const float_t time_allowance_multiplier_;
+  const uint32_t expiration_timeout_;
+  static const connection_handler::ConnectionHandle connection_handle_ =
       0xABCDEF;
   static const transport_manager::ConnectionUID kDefaultConnectionHandle = 1;
   static const uint32_t kDefaultSessionId = 1;
 
-  virtual void SetUp() {
-    conn = new connection_handler::Connection(
-        kConnectionHandle, 0, &connection_handler_mock, kTimeout);
+  void SetUp() OVERRIDE {
+    connection_ = new connection_handler::Connection(
+        connection_handle_, 0, &connection_handler_mock_, timeout_);
   }
 
   void TearDown() OVERRIDE {
-    delete conn;
+    delete connection_;
   }
 };
 
@@ -83,206 +88,191 @@ ACTION_P2(RemoveSession, conn, session_id) {
 }
 
 TEST_F(HeartBeatMonitorTest, TimerNotStarted) {
-  EXPECT_CALL(connection_handler_mock, AddSession(_))
+  EXPECT_CALL(connection_handler_mock_, AddSession(_))
       .WillOnce(Return(kDefaultSessionId));
-  EXPECT_CALL(connection_handler_mock, RemoveSession(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);
+  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(kDefaultConnectionHandle);
+  connection_->AddNewSession(kDefaultConnectionHandle);
 }
 
 TEST_F(HeartBeatMonitorTest, TimerNotElapsed) {
-  EXPECT_CALL(connection_handler_mock, AddSession(_))
+  EXPECT_CALL(connection_handler_mock_, AddSession(_))
       .WillOnce(Return(kDefaultSessionId));
-  EXPECT_CALL(connection_handler_mock, RemoveSession(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);
+  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(kDefaultConnectionHandle);
-  conn->StartHeartBeat(session);
+  const uint32_t session = connection_->AddNewSession(kDefaultConnectionHandle);
+  connection_->StartHeartBeat(session);
 }
 
 TEST_F(HeartBeatMonitorTest, TimerElapsed) {
-  EXPECT_CALL(connection_handler_mock, AddSession(_))
+  EXPECT_CALL(connection_handler_mock_, AddSession(_))
       .WillOnce(Return(kDefaultSessionId));
-  EXPECT_CALL(connection_handler_mock, RemoveSession(kDefaultSessionId))
+  EXPECT_CALL(connection_handler_mock_, RemoveSession(kDefaultSessionId))
       .WillOnce(Return(true));  // invoked by RemoveSession action
 
-  const uint32_t session = conn->AddNewSession(kDefaultConnectionHandle);
+  const uint32_t session = connection_->AddNewSession(kDefaultConnectionHandle);
 
   TestAsyncWaiter waiter;
   uint32_t times = 0;
-  EXPECT_CALL(connection_handler_mock, CloseSession(_, session, _))
-      .WillOnce(
-          DoAll(NotifyTestAsyncWaiter(&waiter), RemoveSession(conn, session)));
-  times++;
-  EXPECT_CALL(connection_handler_mock, CloseConnection(_))
-      .WillOnce(NotifyTestAsyncWaiter(&waiter));
+  EXPECT_CALL(connection_handler_mock_, CloseSession(_, session, _))
+      .WillOnce(DoAll(NotifyTestAsyncWaiter(&waiter),
+                      RemoveSession(connection_, session)));
   times++;
-  EXPECT_CALL(connection_handler_mock, SendHeartBeat(_, session))
+  EXPECT_CALL(connection_handler_mock_, SendHeartBeat(_, session))
       .WillOnce(NotifyTestAsyncWaiter(&waiter));
   times++;
 
-  conn->StartHeartBeat(session);
+  connection_->StartHeartBeat(session);
 
   EXPECT_TRUE(waiter.WaitFor(
       times,
-      2 * kTimeout * MICROSECONDS_IN_MILLISECONDS + MICROSECONDS_IN_SECOND));
+      2 * timeout_ * MICROSECONDS_IN_MILLISECONDS + MICROSECONDS_IN_SECOND));
 }
 
 TEST_F(HeartBeatMonitorTest, KeptAlive) {
-  EXPECT_CALL(connection_handler_mock, AddSession(_))
+  EXPECT_CALL(connection_handler_mock_, AddSession(_))
       .WillOnce(Return(kDefaultSessionId));
-  EXPECT_CALL(connection_handler_mock, RemoveSession(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(kDefaultConnectionHandle);
-  conn->StartHeartBeat(session);
-  usleep(kTimeout * MICROSECONDS_IN_MILLISECONDS - MICROSECONDS_IN_SECOND);
-  conn->KeepAlive(session);
-  usleep(kTimeout * MICROSECONDS_IN_MILLISECONDS - MICROSECONDS_IN_SECOND);
-  conn->KeepAlive(session);
-  usleep(kTimeout * MICROSECONDS_IN_MILLISECONDS - MICROSECONDS_IN_SECOND);
-  conn->KeepAlive(session);
-  usleep(kTimeout * MICROSECONDS_IN_MILLISECONDS - MICROSECONDS_IN_SECOND);
+  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 = connection_->AddNewSession(kDefaultConnectionHandle);
+  connection_->StartHeartBeat(session);
+  usleep(timeout_);
+  connection_->KeepAlive(session);
+  usleep(timeout_);
+  connection_->KeepAlive(session);
 }
 
 TEST_F(HeartBeatMonitorTest, NotKeptAlive) {
-  EXPECT_CALL(connection_handler_mock, AddSession(_))
+  EXPECT_CALL(connection_handler_mock_, AddSession(_))
       .WillOnce(Return(kDefaultSessionId));
-  EXPECT_CALL(connection_handler_mock, RemoveSession(kDefaultSessionId))
+  EXPECT_CALL(connection_handler_mock_, RemoveSession(kDefaultSessionId))
       .WillOnce(Return(true));
 
-  const uint32_t session = conn->AddNewSession(kDefaultConnectionHandle);
+  const uint32_t session = connection_->AddNewSession(kDefaultConnectionHandle);
 
   TestAsyncWaiter waiter;
   uint32_t times = 0;
-  EXPECT_CALL(connection_handler_mock, SendHeartBeat(_, session))
+  EXPECT_CALL(connection_handler_mock_, SendHeartBeat(_, session))
       .WillOnce(NotifyTestAsyncWaiter(&waiter));
   times++;
-  EXPECT_CALL(connection_handler_mock, CloseSession(_, session, _))
-      .WillOnce(
-          DoAll(NotifyTestAsyncWaiter(&waiter), RemoveSession(conn, session)));
-  times++;
-  EXPECT_CALL(connection_handler_mock, CloseConnection(_))
-      .WillOnce(NotifyTestAsyncWaiter(&waiter));
+  EXPECT_CALL(connection_handler_mock_, CloseSession(_, session, _))
+      .WillOnce(DoAll(NotifyTestAsyncWaiter(&waiter),
+                      RemoveSession(connection_, session)));
   times++;
 
-  conn->StartHeartBeat(session);
-  usleep(kTimeout * MICROSECONDS_IN_MILLISECONDS - MICROSECONDS_IN_SECOND);
-  conn->KeepAlive(session);
-  usleep(kTimeout * MICROSECONDS_IN_MILLISECONDS - MICROSECONDS_IN_SECOND);
-  conn->KeepAlive(session);
-  usleep(kTimeout * MICROSECONDS_IN_MILLISECONDS - MICROSECONDS_IN_SECOND);
-  conn->KeepAlive(session);
-  usleep(2 * kTimeout * MICROSECONDS_IN_MILLISECONDS + MICROSECONDS_IN_SECOND);
+  connection_->StartHeartBeat(session);
+  usleep(timeout_);
+  connection_->KeepAlive(session);
 
   EXPECT_TRUE(waiter.WaitFor(
       times,
-      2 * kTimeout * MICROSECONDS_IN_MILLISECONDS + MICROSECONDS_IN_SECOND));
+      2 * timeout_ * MICROSECONDS_IN_MILLISECONDS + MICROSECONDS_IN_SECOND));
 }
 
 TEST_F(HeartBeatMonitorTest, TwoSessionsElapsed) {
   const uint32_t kMockSessionId1 = 1;
   const uint32_t kMockSessionId2 = 2;
-  EXPECT_CALL(connection_handler_mock, AddSession(_))
+  EXPECT_CALL(connection_handler_mock_, AddSession(_))
       .WillOnce(Return(kMockSessionId1))
       .WillOnce(Return(kMockSessionId2));
-  EXPECT_CALL(connection_handler_mock, RemoveSession(kMockSessionId1))
+  EXPECT_CALL(connection_handler_mock_, RemoveSession(kMockSessionId1))
       .WillOnce(Return(true));
-  EXPECT_CALL(connection_handler_mock, RemoveSession(kMockSessionId2))
+  EXPECT_CALL(connection_handler_mock_, RemoveSession(kMockSessionId2))
       .WillOnce(Return(true));
 
-  const uint32_t kSession1 = conn->AddNewSession(kDefaultConnectionHandle);
+  const uint32_t kSession1 =
+      connection_->AddNewSession(kDefaultConnectionHandle);
 
   const transport_manager::ConnectionUID kAnotherConnectionHandle = 2;
-  const uint32_t kSession2 = conn->AddNewSession(kAnotherConnectionHandle);
+  const uint32_t kSession2 =
+      connection_->AddNewSession(kAnotherConnectionHandle);
 
   TestAsyncWaiter waiter;
   uint32_t times = 0;
-  EXPECT_CALL(connection_handler_mock, CloseSession(_, kSession1, _))
+  EXPECT_CALL(connection_handler_mock_, CloseSession(_, kSession1, _))
       .WillOnce(DoAll(NotifyTestAsyncWaiter(&waiter),
-                      RemoveSession(conn, kSession1)));
+                      RemoveSession(connection_, kSession1)));
   times++;
-  EXPECT_CALL(connection_handler_mock, CloseSession(_, kSession2, _))
+  EXPECT_CALL(connection_handler_mock_, CloseSession(_, kSession2, _))
       .WillOnce(DoAll(NotifyTestAsyncWaiter(&waiter),
-                      RemoveSession(conn, kSession2)));
-  times++;
-  EXPECT_CALL(connection_handler_mock, CloseConnection(_))
-      .WillOnce(NotifyTestAsyncWaiter(&waiter));
+                      RemoveSession(connection_, kSession2)));
   times++;
-  EXPECT_CALL(connection_handler_mock, SendHeartBeat(_, kSession1))
+
+  EXPECT_CALL(connection_handler_mock_, SendHeartBeat(_, kSession1))
       .WillOnce(NotifyTestAsyncWaiter(&waiter));
   times++;
-  EXPECT_CALL(connection_handler_mock, SendHeartBeat(_, kSession2))
+  EXPECT_CALL(connection_handler_mock_, SendHeartBeat(_, kSession2))
       .WillOnce(NotifyTestAsyncWaiter(&waiter));
   times++;
 
-  conn->StartHeartBeat(kSession1);
-  conn->StartHeartBeat(kSession2);
+  connection_->StartHeartBeat(kSession1);
+  connection_->StartHeartBeat(kSession2);
 
   EXPECT_TRUE(waiter.WaitFor(
       times,
-      2 * kTimeout * MICROSECONDS_IN_MILLISECONDS + MICROSECONDS_IN_SECOND));
+      2 * timeout_ * MICROSECONDS_IN_MILLISECONDS + MICROSECONDS_IN_SECOND));
 }
 
 TEST_F(HeartBeatMonitorTest, IncreaseHeartBeatTimeout) {
-  EXPECT_CALL(connection_handler_mock, AddSession(_))
+  EXPECT_CALL(connection_handler_mock_, AddSession(_))
       .WillOnce(Return(kDefaultSessionId));
-  EXPECT_CALL(connection_handler_mock, RemoveSession(kDefaultSessionId))
+  EXPECT_CALL(connection_handler_mock_, RemoveSession(kDefaultSessionId))
       .WillOnce(Return(true));
 
-  const uint32_t kSession = conn->AddNewSession(kDefaultConnectionHandle);
+  const uint32_t kSession =
+      connection_->AddNewSession(kDefaultConnectionHandle);
 
-  EXPECT_CALL(connection_handler_mock, CloseSession(_, _)).Times(0);
-  EXPECT_CALL(connection_handler_mock, CloseConnection(_)).Times(0);
-  EXPECT_CALL(connection_handler_mock, SendHeartBeat(_, _)).Times(0);
+  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 kNewTimeout = kTimeout + MICROSECONDS_IN_MILLISECONDS;
-  conn->StartHeartBeat(kSession);
-  conn->SetHeartBeatTimeout(kNewTimeout, kSession);
+  const uint32_t kNewTimeout = timeout_ + MICROSECONDS_IN_MILLISECONDS;
+  connection_->StartHeartBeat(kSession);
+  connection_->SetHeartBeatTimeout(kNewTimeout, kSession);
 }
 
 TEST_F(HeartBeatMonitorTest, DecreaseHeartBeatTimeout) {
-  EXPECT_CALL(connection_handler_mock, AddSession(_))
+  EXPECT_CALL(connection_handler_mock_, AddSession(_))
       .WillOnce(Return(kDefaultSessionId));
-  EXPECT_CALL(connection_handler_mock, RemoveSession(kDefaultSessionId))
+  EXPECT_CALL(connection_handler_mock_, RemoveSession(kDefaultSessionId))
       .WillOnce(Return(true));
 
-  const uint32_t kSession = conn->AddNewSession(kDefaultConnectionHandle);
+  const uint32_t kSession =
+      connection_->AddNewSession(kDefaultConnectionHandle);
 
   TestAsyncWaiter waiter;
   uint32_t times = 0;
-  EXPECT_CALL(connection_handler_mock, CloseSession(_, kSession, _))
-      .WillOnce(
-          DoAll(NotifyTestAsyncWaiter(&waiter), RemoveSession(conn, kSession)));
-  times++;
-  EXPECT_CALL(connection_handler_mock, CloseConnection(_))
-      .WillOnce(NotifyTestAsyncWaiter(&waiter));
+  EXPECT_CALL(connection_handler_mock_, CloseSession(_, kSession, _))
+      .WillOnce(DoAll(NotifyTestAsyncWaiter(&waiter),
+                      RemoveSession(connection_, kSession)));
   times++;
-  EXPECT_CALL(connection_handler_mock, SendHeartBeat(_, kSession))
+  EXPECT_CALL(connection_handler_mock_, SendHeartBeat(_, kSession))
       .WillOnce(NotifyTestAsyncWaiter(&waiter));
   times++;
 
-  const uint32_t kNewTimeout = kTimeout - MICROSECONDS_IN_MILLISECONDS;
-  conn->StartHeartBeat(kSession);
-  conn->SetHeartBeatTimeout(kNewTimeout, kSession);
+  const uint32_t new_timeout = timeout_ - kTime_offset;
+  connection_->StartHeartBeat(kSession);
+  connection_->SetHeartBeatTimeout(new_timeout, kSession);
 
   EXPECT_TRUE(waiter.WaitFor(
       times,
-      2 * kTimeout * MICROSECONDS_IN_MILLISECONDS + MICROSECONDS_IN_SECOND));
+      2 * timeout_ * MICROSECONDS_IN_MILLISECONDS + MICROSECONDS_IN_SECOND));
 }
 
 }  // namespace connection_handler_test