Addressed review comments

3 files changed, 30 insertions, 35 deletions

diff --git a/src/components/transport_manager/test/include/transport_manager/cloud/sample_websocket_server.h b/src/components/transport_manager/test/include/transport_manager/cloud/sample_websocket_server.h
index 141fa1ff89..a60aefc7b0 100644
--- a/src/components/transport_manager/test/include/transport_manager/cloud/sample_websocket_server.h
+++ b/src/components/transport_manager/test/include/transport_manager/cloud/sample_websocket_server.h
@@ -65,9 +65,9 @@ namespace ssl = boost::asio::ssl;        // from <boost/asio/ssl.hpp>
 using tcp = boost::asio::ip::tcp;        // from <boost/asio/ip/tcp.hpp>
 
 // Accepts incoming connections and launches the WSServer
-class WSSession : public std::enable_shared_from_this<WSSession> {
+class WSSession {
  private:
-  class WSServer : public std::enable_shared_from_this<WSServer> {
+  class WSServer {
    public:
     explicit WSServer(tcp::socket&& socket);
     void Run();
@@ -97,9 +97,9 @@ class WSSession : public std::enable_shared_from_this<WSSession> {
 };
 
 // Accepts incoming connections and launches the sessions
-class WSSSession : public std::enable_shared_from_this<WSSSession> {
+class WSSSession {
  private:
-  class WSSServer : public std::enable_shared_from_this<WSSServer> {
+  class WSSServer {
    public:
     // Take ownership of the socket
     WSSServer(tcp::socket&& socket, ssl::context& ctx);
diff --git a/src/components/transport_manager/test/sample_websocket_server.cc b/src/components/transport_manager/test/sample_websocket_server.cc
index ae15d46b62..56899e72e2 100644
--- a/src/components/transport_manager/test/sample_websocket_server.cc
+++ b/src/components/transport_manager/test/sample_websocket_server.cc
@@ -16,9 +16,7 @@ WSSession::WSServer::WSServer(tcp::socket&& socket)
 void WSSession::WSServer::Run() {
   // Accept the websocket handshake
   ws_.async_accept(boost::asio::bind_executor(
-      strand_,
-      std::bind(
-          &WSServer::OnAccept, shared_from_this(), std::placeholders::_1)));
+      strand_, std::bind(&WSServer::OnAccept, this, std::placeholders::_1)));
 }
 
 void WSSession::WSServer::OnAccept(beast::error_code ec) {
@@ -95,10 +93,9 @@ WSSSession::WSSServer::WSSServer(tcp::socket&& socket, ssl::context& ctx)
 
 void WSSSession::WSSServer::Run() {
   // Perform the SSL handshake
-  wss_.next_layer().async_handshake(ssl::stream_base::server,
-                                    std::bind(&WSSServer::OnSSLHandshake,
-                                              shared_from_this(),
-                                              std::placeholders::_1));
+  wss_.next_layer().async_handshake(
+      ssl::stream_base::server,
+      std::bind(&WSSServer::OnSSLHandshake, this, std::placeholders::_1));
 }
 
 void WSSSession::WSSServer::OnSSLHandshake(beast::error_code ec) {
@@ -107,8 +104,8 @@ void WSSSession::WSSServer::OnSSLHandshake(beast::error_code ec) {
   }
 
   // Accept the websocket handshake
-  wss_.async_accept(std::bind(
-      &WSSServer::OnAccept, shared_from_this(), std::placeholders::_1));
+  wss_.async_accept(
+      std::bind(&WSSServer::OnAccept, this, std::placeholders::_1));
 }
 
 void WSSSession::WSSServer::OnAccept(beast::error_code ec) {
@@ -192,8 +189,7 @@ void WSSSession::do_accept() {
   if (acceptor_.is_open()) {
     acceptor_.async_accept(
         socket_,
-        std::bind(
-            &WSSSession::on_accept, shared_from_this(), std::placeholders::_1));
+        std::bind(&WSSSession::on_accept, this, std::placeholders::_1));
     ioc_.run();
   }
 }
diff --git a/src/components/transport_manager/test/websocket_connection_test.cc b/src/components/transport_manager/test/websocket_connection_test.cc
index 2b458a7427..f1f8b0eac1 100644
--- a/src/components/transport_manager/test/websocket_connection_test.cc
+++ b/src/components/transport_manager/test/websocket_connection_test.cc
@@ -58,30 +58,30 @@ class WebsocketConnectionTest : public ::testing::Test {
     std::shared_ptr<WebsocketClientConnection> connection;
   };
 
-  WebsocketClient CreateWebsocketClient(
+  void InitWebsocketClient(
       const transport_manager::transport_adapter::CloudAppProperties&
-          properties) {
+          properties,
+      WebsocketClient& client_out) {
     uniq_id = dev_id = properties.endpoint;
 
-    WebsocketClient client{std::make_shared<CloudWebsocketTransportAdapter>(
-                               last_state_, transport_manager_settings),
-                           nullptr};
-    client.adapter->SetAppCloudTransportConfig(uniq_id, properties);
+    client_out =
+        WebsocketClient{std::make_shared<CloudWebsocketTransportAdapter>(
+                            last_state_, transport_manager_settings),
+                        nullptr};
+    client_out.adapter->SetAppCloudTransportConfig(uniq_id, properties);
 
     TransportAdapterImpl* ta_cloud =
-        dynamic_cast<TransportAdapterImpl*>(client.adapter.get());
+        dynamic_cast<TransportAdapterImpl*>(client_out.adapter.get());
     ta_cloud->CreateDevice(uniq_id);
 
-    auto connection = client.adapter->FindPendingConnection(dev_id, app_handle);
+    auto connection = client_out.adapter->FindPendingConnection(dev_id, 0);
 
     EXPECT_NE(connection, nullptr);
 
-    client.connection =
+    client_out.connection =
         std::dynamic_pointer_cast<WebsocketClientConnection>(connection);
 
-    EXPECT_NE(client.connection.use_count(), 0);
-
-    return client;
+    EXPECT_NE(client_out.connection.use_count(), 0);
   }
 
   void StartWSServer() {
@@ -102,7 +102,6 @@ class WebsocketConnectionTest : public ::testing::Test {
   ~WebsocketConnectionTest() {}
 
   void SetUp() OVERRIDE {
-    app_handle = 0;
     ON_CALL(transport_manager_settings, use_last_state())
         .WillByDefault(Return(true));
   }
@@ -111,9 +110,9 @@ class WebsocketConnectionTest : public ::testing::Test {
   resumption::LastStateImpl last_state_;
   std::string dev_id;
   std::string uniq_id;
-  int app_handle;
   std::shared_ptr<websocket::WSSession> ws_session;
   std::shared_ptr<websocket::WSSSession> wss_session;
+  WebsocketClient ws_client;
   std::string kHost = "127.0.0.1";
   uint16_t kPort = 8080;
   // Sample certificate for localhost
@@ -275,9 +274,9 @@ TEST_F(WebsocketConnectionTest, WSConnection_SUCCESS) {
   sleep(1);
 
   // Start client
-  auto client_connection = CreateWebsocketClient(properties);
+  InitWebsocketClient(properties, ws_client);
   std::shared_ptr<WebsocketClientConnection> ws_connection =
-      client_connection.connection;
+      ws_client.connection;
 
   // Check websocket connection
   TransportAdapter::Error ret_code = ws_connection->Start();
@@ -306,9 +305,9 @@ TEST_F(WebsocketConnectionTest, WSSConnection_SUCCESS) {
   sleep(1);
 
   // Start client
-  auto client_connection = CreateWebsocketClient(properties);
+  InitWebsocketClient(properties, ws_client);
   std::shared_ptr<WebsocketClientConnection> wss_connection =
-      client_connection.connection;
+      ws_client.connection;
 
   // Check websocket connection
   TransportAdapter::Error ret_code = wss_connection->Start();
@@ -337,9 +336,9 @@ TEST_F(WebsocketConnectionTest, WSSConnection_FAILURE_IncorrectCert) {
   sleep(1);
 
   // Start client
-  auto client_connection = CreateWebsocketClient(properties);
+  InitWebsocketClient(properties, ws_client);
   std::shared_ptr<WebsocketClientConnection> wss_connection =
-      client_connection.connection;
+      ws_client.connection;
 
   // Check websocket connection
   TransportAdapter::Error ret_code = wss_connection->Start();