1 files changed, 69 insertions, 56 deletions
diff --git a/lldb/unittests/Host/SocketTest.cpp b/lldb/unittests/Host/SocketTest.cpp
index deed222c8c1..26bef353209 100644
--- a/lldb/unittests/Host/SocketTest.cpp
+++ b/lldb/unittests/Host/SocketTest.cpp
@@ -13,66 +13,79 @@
 
 #include "lldb/Host/Socket.h"
 
-class SocketTest: public ::testing::Test
-{
-};
-
 using namespace lldb_private;
 
-void AcceptThread (Socket* listen_socket,
-                  const char* listen_remote_address,
-                  bool child_processes_inherit,
-                  Socket** accept_socket,
-                  Error* error)
+class SocketTest : public testing::Test
 {
-    *error = listen_socket->BlockingAccept (listen_remote_address, child_processes_inherit, *accept_socket);
-}
+  public:
+    void
+    SetUp() override
+    {
+#if defined(_MSC_VER)
+        WSADATA data;
+        ::WSAStartup(MAKEWORD(2, 2), &data);
+#endif
+    }
 
-void CreateConnectedSockets (std::unique_ptr<Socket>* a_up, std::unique_ptr<Socket>* b_up)
-{
-    Predicate<uint16_t> port_predicate;
-    // Used when binding to port zero to wait for the thread
-    // that creates the socket, binds and listens to resolve
-    // the port number.
-    
-    port_predicate.SetValue (0, eBroadcastNever);
-    
-    bool child_processes_inherit = false;
-    Socket *socket = nullptr;
-    const char* listen_remote_address = "localhost:0";
-    Error error = Socket::TcpListen (listen_remote_address, child_processes_inherit, socket, &port_predicate);
-    std::unique_ptr<Socket> listen_socket_up (socket);
-    socket = nullptr;
-    EXPECT_FALSE (error.Fail ());
-    EXPECT_NE (nullptr, listen_socket_up.get ());
-    EXPECT_TRUE (listen_socket_up->IsValid ());
-
-    Error accept_error;
-    Socket* accept_socket;
-    std::thread accept_thread (AcceptThread,
-                               listen_socket_up.get (),
-                               listen_remote_address,
-                               child_processes_inherit,
-                               &accept_socket,
-                               &accept_error);
-    
-    char connect_remote_address[64];
-    snprintf (connect_remote_address, sizeof (connect_remote_address), "localhost:%u", port_predicate.GetValue ());
-    error = Socket::TcpConnect (connect_remote_address, child_processes_inherit, socket);
-    a_up->reset (socket);
-    socket = nullptr;
-    EXPECT_TRUE (error.Success ());
-    EXPECT_NE (nullptr, a_up->get ());
-    EXPECT_TRUE ((*a_up)->IsValid ());
-    
-    accept_thread.join ();
-    b_up->reset (accept_socket);
-    EXPECT_TRUE (accept_error.Success ());
-    EXPECT_NE (nullptr, b_up->get ());
-    EXPECT_TRUE ((*b_up)->IsValid ());
-    
-    listen_socket_up.reset ();
-}
+    void
+    TearDown() override
+    {
+#if defined(_MSC_VER)
+        ::WSACleanup();
+#endif
+    }
+
+  protected:
+    static void
+    AcceptThread(Socket *listen_socket, const char *listen_remote_address, bool child_processes_inherit,
+                 Socket **accept_socket, Error *error)
+    {
+        *error = listen_socket->BlockingAccept(listen_remote_address, child_processes_inherit, *accept_socket);
+    }
+
+    void
+    CreateConnectedSockets(std::unique_ptr<Socket> *a_up, std::unique_ptr<Socket> *b_up)
+    {
+        Predicate<uint16_t> port_predicate;
+        // Used when binding to port zero to wait for the thread
+        // that creates the socket, binds and listens to resolve
+        // the port number.
+
+        port_predicate.SetValue(0, eBroadcastNever);
+
+        bool child_processes_inherit = false;
+        Socket *socket = nullptr;
+        const char *listen_remote_address = "localhost:0";
+        Error error = Socket::TcpListen(listen_remote_address, child_processes_inherit, socket, &port_predicate);
+        std::unique_ptr<Socket> listen_socket_up(socket);
+        socket = nullptr;
+        EXPECT_FALSE(error.Fail());
+        EXPECT_NE(nullptr, listen_socket_up.get());
+        EXPECT_TRUE(listen_socket_up->IsValid());
+
+        Error accept_error;
+        Socket *accept_socket;
+        std::thread accept_thread(AcceptThread, listen_socket_up.get(), listen_remote_address, child_processes_inherit,
+                                  &accept_socket, &accept_error);
+
+        char connect_remote_address[64];
+        snprintf(connect_remote_address, sizeof(connect_remote_address), "localhost:%u", port_predicate.GetValue());
+        error = Socket::TcpConnect(connect_remote_address, child_processes_inherit, socket);
+        a_up->reset(socket);
+        socket = nullptr;
+        EXPECT_TRUE(error.Success());
+        EXPECT_NE(nullptr, a_up->get());
+        EXPECT_TRUE((*a_up)->IsValid());
+
+        accept_thread.join();
+        b_up->reset(accept_socket);
+        EXPECT_TRUE(accept_error.Success());
+        EXPECT_NE(nullptr, b_up->get());
+        EXPECT_TRUE((*b_up)->IsValid());
+
+        listen_socket_up.reset();
+    }
+};
 
 TEST_F (SocketTest, DecodeHostAndPort)
 {