9 files changed, 503 insertions, 273 deletions

diff --git a/unittests/tools/lldb-server/tests/CMakeLists.txt b/unittests/tools/lldb-server/tests/CMakeLists.txt
index 4ef4165b27fef..ed5eb88ba5270 100644
--- a/unittests/tools/lldb-server/tests/CMakeLists.txt
+++ b/unittests/tools/lldb-server/tests/CMakeLists.txt
@@ -1,6 +1,8 @@
 add_lldb_unittest(LLDBServerTests
-  TestClient.cpp
+  LLGSTest.cpp
   MessageObjects.cpp
+  TestBase.cpp
+  TestClient.cpp
   ThreadIdsInJstopinfoTest.cpp
 
   LINK_LIBS
@@ -10,6 +12,10 @@ add_lldb_unittest(LLDBServerTests
     lldbTarget
     lldbPluginPlatformLinux
     lldbPluginProcessGDBRemote
+
+    LLVMTestingSupport
   LINK_COMPONENTS
     Support
   )
+
+add_dependencies(LLDBServerTests lldb-server ${ALL_LLDB_TEST_EXECUTABLES})
diff --git a/unittests/tools/lldb-server/tests/LLGSTest.cpp b/unittests/tools/lldb-server/tests/LLGSTest.cpp
new file mode 100644
index 0000000000000..75d964ed26f44
--- /dev/null
+++ b/unittests/tools/lldb-server/tests/LLGSTest.cpp
@@ -0,0 +1,36 @@
+//===-- LLGSTest.cpp --------------------------------------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestBase.h"
+#include "lldb/Host/Host.h"
+#include "llvm/Testing/Support/Error.h"
+
+using namespace llgs_tests;
+using namespace lldb_private;
+using namespace llvm;
+
+TEST_F(TestBase, LaunchModePreservesEnvironment) {
+  if (TestClient::IsDebugServer()) {
+    GTEST_LOG_(WARNING) << "Test fails with debugserver: llvm.org/pr35671";
+    return;
+  }
+
+  putenv(const_cast<char *>("LLDB_TEST_MAGIC_VARIABLE=LLDB_TEST_MAGIC_VALUE"));
+
+  auto ClientOr = TestClient::launch(getLogFileName(),
+                                     {getInferiorPath("environment_check")});
+  ASSERT_THAT_EXPECTED(ClientOr, Succeeded());
+  auto &Client = **ClientOr;
+
+  ASSERT_THAT_ERROR(Client.ContinueAll(), Succeeded());
+  ASSERT_THAT_EXPECTED(
+      Client.GetLatestStopReplyAs<StopReplyExit>(),
+      HasValue(testing::Property(&StopReply::getKind,
+                                 WaitStatus{WaitStatus::Exit, 0})));
+}
diff --git a/unittests/tools/lldb-server/tests/MessageObjects.cpp b/unittests/tools/lldb-server/tests/MessageObjects.cpp
index 6ac067bce0621..e5927d037f985 100644
--- a/unittests/tools/lldb-server/tests/MessageObjects.cpp
+++ b/unittests/tools/lldb-server/tests/MessageObjects.cpp
@@ -19,7 +19,7 @@ namespace llgs_tests {
 
 Expected<ProcessInfo> ProcessInfo::Create(StringRef response) {
   ProcessInfo process_info;
-  auto elements_or_error = SplitPairList("ProcessInfo", response);
+  auto elements_or_error = SplitUniquePairList("ProcessInfo", response);
   if (!elements_or_error)
     return elements_or_error.takeError();
 
@@ -65,19 +65,16 @@ StringRef ThreadInfo::ReadRegister(unsigned int register_id) const {
   return m_registers.lookup(register_id);
 }
 
-bool ThreadInfo::ReadRegisterAsUint64(unsigned int register_id,
-                                      uint64_t &value) const {
+Expected<uint64_t>
+ThreadInfo::ReadRegisterAsUint64(unsigned int register_id) const {
+  uint64_t value;
   std::string value_str(m_registers.lookup(register_id));
-  if (!llvm::to_integer(value_str, value, 16)) {
-    GTEST_LOG_(ERROR)
-        << formatv("ThreadInfo: Unable to parse register value at {0}.",
-                   register_id)
-               .str();
-    return false;
-  }
+  if (!llvm::to_integer(value_str, value, 16))
+    return make_parsing_error("ThreadInfo value for register {0}: {1}",
+                              register_id, value_str);
 
   sys::swapByteOrder(value);
-  return true;
+  return value;
 }
 
 //====== JThreadsInfo ==========================================================
@@ -136,60 +133,96 @@ const ThreadInfoMap &JThreadsInfo::GetThreadInfos() const {
 }
 
 //====== StopReply =============================================================
-const U64Map &StopReply::GetThreadPcs() const { return m_thread_pcs; }
+Expected<std::unique_ptr<StopReply>>
+StopReply::create(StringRef Response, llvm::support::endianness Endian) {
+  if (Response.size() < 3)
+    return make_parsing_error("StopReply: Invalid packet");
+  if (Response.consume_front("T"))
+    return StopReplyStop::create(Response, Endian);
+  if (Response.consume_front("W"))
+    return StopReplyExit::create(Response);
+  return make_parsing_error("StopReply: Invalid packet");
+}
 
-Expected<StopReply> StopReply::Create(StringRef response,
-                                      llvm::support::endianness endian) {
-  StopReply stop_reply;
+Expected<std::unique_ptr<StopReplyStop>>
+StopReplyStop::create(StringRef Response, llvm::support::endianness Endian) {
+  unsigned int Signal;
+  StringRef SignalStr = Response.take_front(2);
+  Response = Response.drop_front(2);
+  if (!to_integer(SignalStr, Signal, 16))
+    return make_parsing_error("StopReply: stop signal");
+
+  auto Elements = SplitPairList(Response);
+  for (StringRef Field :
+       {"name", "reason", "thread", "threads", "thread-pcs"}) {
+    // This will insert an empty field if there is none. In the future, we
+    // should probably differentiate between these fields not being present and
+    // them being empty, but right now no tests depends on this.
+    if (Elements.insert({Field, {""}}).first->second.size() != 1)
+      return make_parsing_error(
+          "StopReply: got multiple responses for the {0} field", Field);
+  }
+  StringRef Name = Elements["name"][0];
+  StringRef Reason = Elements["reason"][0];
+
+  lldb::tid_t Thread;
+  if (!to_integer(Elements["thread"][0], Thread, 16))
+    return make_parsing_error("StopReply: thread");
+
+  SmallVector<StringRef, 20> Threads;
+  SmallVector<StringRef, 20> Pcs;
+  Elements["threads"][0].split(Threads, ',');
+  Elements["thread-pcs"][0].split(Pcs, ',');
+  if (Threads.size() != Pcs.size())
+    return make_parsing_error("StopReply: thread/PC count mismatch");
 
-  auto elements_or_error = SplitPairList("StopReply", response);
-  if (auto split_error = elements_or_error.takeError()) {
-    return std::move(split_error);
+  U64Map ThreadPcs;
+  for (auto ThreadPc : zip(Threads, Pcs)) {
+    lldb::tid_t Id;
+    uint64_t Pc;
+    if (!to_integer(std::get<0>(ThreadPc), Id, 16))
+      return make_parsing_error("StopReply: Thread id '{0}'",
+                                std::get<0>(ThreadPc));
+    if (!to_integer(std::get<1>(ThreadPc), Pc, 16))
+      return make_parsing_error("StopReply Thread Pc '{0}'",
+                                std::get<1>(ThreadPc));
+
+    ThreadPcs[Id] = Pc;
   }
 
-  auto elements = *elements_or_error;
-  stop_reply.m_name = elements["name"];
-  stop_reply.m_reason = elements["reason"];
+  RegisterMap Registers;
+  for (const auto &E : Elements) {
+    StringRef Key = E.getKey();
+    const auto &Val = E.getValue();
+    if (Key.size() != 2)
+      continue;
 
-  SmallVector<StringRef, 20> threads;
-  SmallVector<StringRef, 20> pcs;
-  elements["threads"].split(threads, ',');
-  elements["thread-pcs"].split(pcs, ',');
-  if (threads.size() != pcs.size())
-    return make_parsing_error("StopReply: thread/PC count mismatch");
+    unsigned int Reg;
+    if (!to_integer(Key, Reg, 16))
+      continue;
 
-  for (size_t i = 0; i < threads.size(); i++) {
-    lldb::tid_t thread_id;
-    uint64_t pc;
-    if (threads[i].getAsInteger(16, thread_id))
-      return make_parsing_error("StopReply: thread ID at [{0}].", i);
-    if (pcs[i].getAsInteger(16, pc))
-      return make_parsing_error("StopReply: thread PC at [{0}].", i);
+    if (Val.size() != 1)
+      return make_parsing_error(
+          "StopReply: multiple entries for register field [{0:x}]", Reg);
 
-    stop_reply.m_thread_pcs[thread_id] = pc;
+    Registers[Reg] = Val[0].str();
   }
 
-  for (auto i = elements.begin(); i != elements.end(); i++) {
-    StringRef key = i->getKey();
-    StringRef val = i->getValue();
-    if (key.size() >= 9 && key[0] == 'T' && key.substr(3, 6) == "thread") {
-      if (val.getAsInteger(16, stop_reply.m_thread))
-        return make_parsing_error("StopReply: thread id");
-      if (key.substr(1, 2).getAsInteger(16, stop_reply.m_signal))
-        return make_parsing_error("StopReply: stop signal");
-    } else if (key.size() == 2) {
-      unsigned int reg;
-      if (!key.getAsInteger(16, reg)) {
-        stop_reply.m_registers[reg] = val.str();
-      }
-    }
-  }
+  return llvm::make_unique<StopReplyStop>(Signal, Thread, Name, ThreadPcs,
+                                          Registers, Reason);
+}
 
-  return stop_reply;
+Expected<std::unique_ptr<StopReplyExit>>
+StopReplyExit::create(StringRef Response) {
+  uint8_t Status;
+  if (!to_integer(Response, Status, 16))
+    return make_parsing_error("StopReply: exit status");
+  return llvm::make_unique<StopReplyExit>(Status);
 }
 
 //====== Globals ===============================================================
-Expected<StringMap<StringRef>> SplitPairList(StringRef caller, StringRef str) {
+Expected<StringMap<StringRef>> SplitUniquePairList(StringRef caller,
+                                                   StringRef str) {
   SmallVector<StringRef, 20> elements;
   str.split(elements, ';');
 
@@ -199,7 +232,20 @@ Expected<StringMap<StringRef>> SplitPairList(StringRef caller, StringRef str) {
     if (pairs.count(pair.first))
       return make_parsing_error("{0}: Duplicate Key: {1}", caller, pair.first);
 
-    pairs.insert(s.split(':'));
+    pairs.insert(pair);
+  }
+
+  return pairs;
+}
+
+StringMap<SmallVector<StringRef, 2>> SplitPairList(StringRef str) {
+  SmallVector<StringRef, 20> elements;
+  str.split(elements, ';');
+
+  StringMap<SmallVector<StringRef, 2>> pairs;
+  for (StringRef s : elements) {
+    std::pair<StringRef, StringRef> pair = s.split(':');
+    pairs[pair.first].push_back(pair.second);
   }
 
   return pairs;
diff --git a/unittests/tools/lldb-server/tests/MessageObjects.h b/unittests/tools/lldb-server/tests/MessageObjects.h
index 82551e2bb5490..bbb7f75bafd5c 100644
--- a/unittests/tools/lldb-server/tests/MessageObjects.h
+++ b/unittests/tools/lldb-server/tests/MessageObjects.h
@@ -7,6 +7,10 @@
 //
 //===----------------------------------------------------------------------===//
 
+#ifndef LLDB_SERVER_TESTS_MESSAGEOBJECTS_H
+#define LLDB_SERVER_TESTS_MESSAGEOBJECTS_H
+
+#include "lldb/Host/Host.h"
 #include "lldb/lldb-types.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallString.h"
@@ -48,7 +52,7 @@ public:
              const RegisterMap &registers, unsigned int signal);
 
   llvm::StringRef ReadRegister(unsigned int register_id) const;
-  bool ReadRegisterAsUint64(unsigned int register_id, uint64_t &value) const;
+  llvm::Expected<uint64_t> ReadRegisterAsUint64(unsigned int register_id) const;
 
 private:
   std::string m_name;
@@ -71,25 +75,75 @@ private:
 
 class StopReply {
 public:
-  static llvm::Expected<StopReply> Create(llvm::StringRef response,
-                                          llvm::support::endianness endian);
-  const U64Map &GetThreadPcs() const;
+  StopReply() = default;
+  virtual ~StopReply() = default;
+
+  static llvm::Expected<std::unique_ptr<StopReply>>
+  create(llvm::StringRef response, llvm::support::endianness endian);
+
+  // for llvm::cast<>
+  virtual lldb_private::WaitStatus getKind() const = 0;
+
+  StopReply(const StopReply &) = delete;
+  void operator=(const StopReply &) = delete;
+};
+
+class StopReplyStop : public StopReply {
+public:
+  StopReplyStop(uint8_t Signal, lldb::tid_t ThreadId, llvm::StringRef Name,
+                U64Map ThreadPcs, RegisterMap Registers, llvm::StringRef Reason)
+      : Signal(Signal), ThreadId(ThreadId), Name(Name),
+        ThreadPcs(std::move(ThreadPcs)), Registers(std::move(Registers)),
+        Reason(Reason) {}
+
+  static llvm::Expected<std::unique_ptr<StopReplyStop>>
+  create(llvm::StringRef response, llvm::support::endianness endian);
+
+  const U64Map &getThreadPcs() const { return ThreadPcs; }
+  lldb::tid_t getThreadId() const { return ThreadId; }
+
+  // for llvm::cast<>
+  lldb_private::WaitStatus getKind() const override {
+    return lldb_private::WaitStatus{lldb_private::WaitStatus::Stop, Signal};
+  }
+  static bool classof(const StopReply *R) {
+    return R->getKind().type == lldb_private::WaitStatus::Stop;
+  }
 
 private:
-  StopReply() = default;
-  void ParseResponse(llvm::StringRef response,
-                     llvm::support::endianness endian);
-  unsigned int m_signal;
-  lldb::tid_t m_thread;
-  std::string m_name;
-  U64Map m_thread_pcs;
-  RegisterMap m_registers;
-  std::string m_reason;
+  uint8_t Signal;
+  lldb::tid_t ThreadId;
+  std::string Name;
+  U64Map ThreadPcs;
+  RegisterMap Registers;
+  std::string Reason;
+};
+
+class StopReplyExit : public StopReply {
+public:
+  explicit StopReplyExit(uint8_t Status) : Status(Status) {}
+
+  static llvm::Expected<std::unique_ptr<StopReplyExit>>
+  create(llvm::StringRef response);
+
+  // for llvm::cast<>
+  lldb_private::WaitStatus getKind() const override {
+    return lldb_private::WaitStatus{lldb_private::WaitStatus::Exit, Status};
+  }
+  static bool classof(const StopReply *R) {
+    return R->getKind().type == lldb_private::WaitStatus::Exit;
+  }
+
+private:
+  uint8_t Status;
 };
 
 // Common functions for parsing packet data.
 llvm::Expected<llvm::StringMap<llvm::StringRef>>
-SplitPairList(llvm::StringRef caller, llvm::StringRef s);
+SplitUniquePairList(llvm::StringRef caller, llvm::StringRef s);
+
+llvm::StringMap<llvm::SmallVector<llvm::StringRef, 2>>
+SplitPairList(llvm::StringRef s);
 
 template <typename... Args>
 llvm::Error make_parsing_error(llvm::StringRef format, Args &&... args) {
@@ -99,4 +153,7 @@ llvm::Error make_parsing_error(llvm::StringRef format, Args &&... args) {
   return llvm::make_error<llvm::StringError>(error,
                                              llvm::inconvertibleErrorCode());
 }
+
 } // namespace llgs_tests
+
+#endif // LLDB_SERVER_TESTS_MESSAGEOBJECTS_H
diff --git a/unittests/tools/lldb-server/tests/TestBase.cpp b/unittests/tools/lldb-server/tests/TestBase.cpp
new file mode 100644
index 0000000000000..b901249dbf4c6
--- /dev/null
+++ b/unittests/tools/lldb-server/tests/TestBase.cpp
@@ -0,0 +1,36 @@
+//===-- TestBase.cpp --------------------------------------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "TestBase.h"
+#include <cstdlib>
+
+using namespace llgs_tests;
+using namespace llvm;
+
+std::string TestBase::getLogFileName() {
+  const auto *test_info =
+      ::testing::UnitTest::GetInstance()->current_test_info();
+  assert(test_info);
+
+  const char *Dir = getenv("LOG_FILE_DIRECTORY");
+  if (!Dir)
+    return "";
+
+  if (!llvm::sys::fs::is_directory(Dir)) {
+    GTEST_LOG_(WARNING) << "Cannot access log directory: " << Dir;
+    return "";
+  }
+
+  SmallString<64> DirStr(Dir);
+  sys::path::append(DirStr, std::string("server-") +
+                                test_info->test_case_name() + "-" +
+                                test_info->name() + ".log");
+  return DirStr.str();
+}
+
diff --git a/unittests/tools/lldb-server/tests/TestBase.h b/unittests/tools/lldb-server/tests/TestBase.h
new file mode 100644
index 0000000000000..8956490e945c1
--- /dev/null
+++ b/unittests/tools/lldb-server/tests/TestBase.h
@@ -0,0 +1,48 @@
+//===-- TestBase.h ----------------------------------------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLDB_SERVER_TESTS_TESTBASE_H
+#define LLDB_SERVER_TESTS_TESTBASE_H
+
+#include "TestClient.h"
+#include "lldb/Host/HostInfo.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Testing/Support/Error.h"
+#include "gtest/gtest.h"
+
+namespace llgs_tests {
+
+class TestBase: public ::testing::Test {
+public:
+  static void SetUpTestCase() { lldb_private::HostInfo::Initialize(); }
+
+  static std::string getInferiorPath(llvm::StringRef Name) {
+    llvm::SmallString<64> Path(LLDB_TEST_INFERIOR_PATH);
+    llvm::sys::path::append(Path, Name + LLDB_TEST_INFERIOR_SUFFIX);
+    return Path.str();
+  }
+
+  static std::string getLogFileName();
+};
+
+class StandardStartupTest: public TestBase {
+public:
+  void SetUp() override {
+    auto ClientOr = TestClient::launch(getLogFileName());
+    ASSERT_THAT_EXPECTED(ClientOr, llvm::Succeeded());
+    Client = std::move(*ClientOr);
+  }
+
+protected:
+  std::unique_ptr<TestClient> Client;
+};
+
+} // namespace llgs_tests
+
+#endif // LLDB_SERVER_TESTS_TESTBASE_H
diff --git a/unittests/tools/lldb-server/tests/TestClient.cpp b/unittests/tools/lldb-server/tests/TestClient.cpp
index 9553ddd65fac3..773466abe8553 100644
--- a/unittests/tools/lldb-server/tests/TestClient.cpp
+++ b/unittests/tools/lldb-server/tests/TestClient.cpp
@@ -8,13 +8,14 @@
 //===----------------------------------------------------------------------===//
 
 #include "TestClient.h"
-#include "lldb/Core/ArchSpec.h"
 #include "lldb/Host/HostInfo.h"
 #include "lldb/Host/common/TCPSocket.h"
 #include "lldb/Host/posix/ConnectionFileDescriptorPosix.h"
 #include "lldb/Interpreter/Args.h"
 #include "lldb/Target/ProcessLaunchInfo.h"
 #include "llvm/ADT/StringExtras.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Testing/Support/Error.h"
 #include "gtest/gtest.h"
 #include <cstdlib>
 #include <future>
@@ -26,67 +27,100 @@ using namespace lldb_private;
 using namespace llvm;
 
 namespace llgs_tests {
-void TestClient::Initialize() { HostInfo::Initialize(); }
+bool TestClient::IsDebugServer() {
+  return sys::path::filename(LLDB_SERVER).contains("debugserver");
+}
+
+bool TestClient::IsLldbServer() { return !IsDebugServer(); }
+
+TestClient::TestClient(std::unique_ptr<Connection> Conn) {
+  SetConnection(Conn.release());
+
+  SendAck(); // Send this as a handshake.
+}
 
-TestClient::TestClient(const std::string &test_name,
-                       const std::string &test_case_name)
-    : m_test_name(test_name), m_test_case_name(test_case_name),
-      m_pc_register(UINT_MAX) {}
+TestClient::~TestClient() {
+  if (!IsConnected())
+    return;
 
-TestClient::~TestClient() {}
+  std::string response;
+  // Debugserver (non-conformingly?) sends a reply to the k packet instead of
+  // simply closing the connection.
+  PacketResult result =
+      IsDebugServer() ? PacketResult::Success : PacketResult::ErrorDisconnected;
+  EXPECT_THAT_ERROR(SendMessage("k", response, result), Succeeded());
+}
+
+Expected<std::unique_ptr<TestClient>> TestClient::launch(StringRef Log) {
+  return launch(Log, {});
+}
 
-bool TestClient::StartDebugger() {
+Expected<std::unique_ptr<TestClient>> TestClient::launch(StringRef Log, ArrayRef<StringRef> InferiorArgs) {
   const ArchSpec &arch_spec = HostInfo::GetArchitecture();
   Args args;
   args.AppendArgument(LLDB_SERVER);
-  args.AppendArgument("gdbserver");
-  args.AppendArgument("--log-channels=gdb-remote packets");
+  if (IsLldbServer())
+    args.AppendArgument("gdbserver");
   args.AppendArgument("--reverse-connect");
-  std::string log_file_name = GenerateLogFileName(arch_spec);
-  if (log_file_name.size()) {
-    args.AppendArgument("--log-file=" + log_file_name);
+
+  if (!Log.empty()) {
+    args.AppendArgument(("--log-file=" + Log).str());
+    if (IsLldbServer())
+      args.AppendArgument("--log-channels=gdb-remote packets");
+    else
+      args.AppendArgument("--log-flags=0x800000");
   }
 
-  Status error;
+  Status status;
   TCPSocket listen_socket(true, false);
-  error = listen_socket.Listen("127.0.0.1:0", 5);
-  if (error.Fail()) {
-    GTEST_LOG_(ERROR) << "Unable to open listen socket.";
-    return false;
+  status = listen_socket.Listen("127.0.0.1:0", 5);
+  if (status.Fail())
+    return status.ToError();
+
+  args.AppendArgument(
+      ("localhost:" + Twine(listen_socket.GetLocalPortNumber())).str());
+
+  if (!InferiorArgs.empty()) {
+    args.AppendArgument("--");
+    for (StringRef arg : InferiorArgs)
+      args.AppendArgument(arg);
   }
 
-  char connect_remote_address[64];
-  snprintf(connect_remote_address, sizeof(connect_remote_address),
-           "localhost:%u", listen_socket.GetLocalPortNumber());
+  ProcessLaunchInfo Info;
+  Info.SetArchitecture(arch_spec);
+  Info.SetArguments(args, true);
 
-  args.AppendArgument(connect_remote_address);
+  StringList Env;
+  Host::GetEnvironment(Env);
+  Info.GetEnvironmentEntries() = Args(Env);
 
-  m_server_process_info.SetArchitecture(arch_spec);
-  m_server_process_info.SetArguments(args, true);
-  Status status = Host::LaunchProcess(m_server_process_info);
-  if (status.Fail()) {
-    GTEST_LOG_(ERROR)
-        << formatv("Failure to launch lldb server: {0}.", status).str();
-    return false;
-  }
+  status = Host::LaunchProcess(Info);
+  if (status.Fail())
+    return status.ToError();
 
-  char connect_remote_uri[64];
-  snprintf(connect_remote_uri, sizeof(connect_remote_uri), "connect://%s",
-           connect_remote_address);
   Socket *accept_socket;
   listen_socket.Accept(accept_socket);
-  SetConnection(new ConnectionFileDescriptor(accept_socket));
+  auto Conn = llvm::make_unique<ConnectionFileDescriptor>(accept_socket);
+  auto Client = std::unique_ptr<TestClient>(new TestClient(std::move(Conn)));
 
-  SendAck(); // Send this as a handshake.
-  return true;
-}
+  if (!InferiorArgs.empty()) {
+    if (Error E = Client->QueryProcessInfo())
+      return std::move(E);
+  }
 
-bool TestClient::StopDebugger() {
-  std::string response;
-  return SendMessage("k", response, PacketResult::ErrorDisconnected);
+  return std::move(Client);
 }
 
-bool TestClient::SetInferior(llvm::ArrayRef<std::string> inferior_args) {
+Error TestClient::SetInferior(llvm::ArrayRef<std::string> inferior_args) {
+  StringList env;
+  Host::GetEnvironment(env);
+  for (size_t i = 0; i < env.GetSize(); ++i) {
+    if (SendEnvironmentPacket(env[i].c_str()) != 0) {
+      return make_error<StringError>(
+          formatv("Failed to set environment variable: {0}", env[i]).str(),
+          inconvertibleErrorCode());
+    }
+  }
   std::stringstream command;
   command << "A";
   for (size_t i = 0; i < inferior_args.size(); i++) {
@@ -96,36 +130,26 @@ bool TestClient::SetInferior(llvm::ArrayRef<std::string> inferior_args) {
     command << hex_encoded.size() << ',' << i << ',' << hex_encoded;
   }
 
-  if (!SendMessage(command.str()))
-    return false;
-  if (!SendMessage("qLaunchSuccess"))
-    return false;
-  std::string response;
-  if (!SendMessage("qProcessInfo", response))
-    return false;
-  auto create_or_error = ProcessInfo::Create(response);
-  if (auto create_error = create_or_error.takeError()) {
-    GTEST_LOG_(ERROR) << toString(std::move(create_error));
-    return false;
-  }
-
-  m_process_info = *create_or_error;
-  return true;
+  if (Error E = SendMessage(command.str()))
+    return E;
+  if (Error E = SendMessage("qLaunchSuccess"))
+    return E;
+  if (Error E = QueryProcessInfo())
+    return E;
+  return Error::success();
 }
 
-bool TestClient::ListThreadsInStopReply() {
+Error TestClient::ListThreadsInStopReply() {
   return SendMessage("QListThreadsInStopReply");
 }
 
-bool TestClient::SetBreakpoint(unsigned long address) {
-  std::stringstream command;
-  command << "Z0," << std::hex << address << ",1";
-  return SendMessage(command.str());
+Error TestClient::SetBreakpoint(unsigned long address) {
+  return SendMessage(formatv("Z0,{0:x-},1", address).str());
 }
 
-bool TestClient::ContinueAll() { return Continue("vCont;c"); }
+Error TestClient::ContinueAll() { return Continue("vCont;c"); }
 
-bool TestClient::ContinueThread(unsigned long thread_id) {
+Error TestClient::ContinueThread(unsigned long thread_id) {
   return Continue(formatv("vCont;c:{0:x-}", thread_id).str());
 }
 
@@ -133,7 +157,7 @@ const ProcessInfo &TestClient::GetProcessInfo() { return *m_process_info; }
 
 Optional<JThreadsInfo> TestClient::GetJThreadsInfo() {
   std::string response;
-  if (!SendMessage("jThreadsInfo", response))
+  if (SendMessage("jThreadsInfo", response))
     return llvm::None;
   auto creation = JThreadsInfo::Create(response, m_process_info->GetEndian());
   if (auto create_error = creation.takeError()) {
@@ -145,39 +169,41 @@ Optional<JThreadsInfo> TestClient::GetJThreadsInfo() {
 }
 
 const StopReply &TestClient::GetLatestStopReply() {
-  return m_stop_reply.getValue();
+  assert(m_stop_reply);
+  return *m_stop_reply;
 }
 
-bool TestClient::SendMessage(StringRef message) {
+Error TestClient::SendMessage(StringRef message) {
   std::string dummy_string;
   return SendMessage(message, dummy_string);
 }
 
-bool TestClient::SendMessage(StringRef message, std::string &response_string) {
-  if (!SendMessage(message, response_string, PacketResult::Success))
-    return false;
-  else if (response_string[0] == 'E') {
-    GTEST_LOG_(ERROR) << "Error " << response_string
-                      << " while sending message: " << message.str();
-    return false;
+Error TestClient::SendMessage(StringRef message, std::string &response_string) {
+  if (Error E = SendMessage(message, response_string, PacketResult::Success))
+    return E;
+  if (response_string[0] == 'E') {
+    return make_error<StringError>(
+        formatv("Error `{0}` while sending message: {1}", response_string,
+                message)
+            .str(),
+        inconvertibleErrorCode());
   }
-
-  return true;
+  return Error::success();
 }
 
-bool TestClient::SendMessage(StringRef message, std::string &response_string,
-                             PacketResult expected_result) {
+Error TestClient::SendMessage(StringRef message, std::string &response_string,
+                              PacketResult expected_result) {
   StringExtractorGDBRemote response;
   GTEST_LOG_(INFO) << "Send Packet: " << message.str();
   PacketResult result = SendPacketAndWaitForResponse(message, response, false);
   response.GetEscapedBinaryData(response_string);
   GTEST_LOG_(INFO) << "Read Packet: " << response_string;
-  if (result != expected_result) {
-    GTEST_LOG_(ERROR) << FormatFailedResult(message, result);
-    return false;
-  }
+  if (result != expected_result)
+    return make_error<StringError>(
+        formatv("Error sending message `{0}`: {1}", message, result).str(),
+        inconvertibleErrorCode());
 
-  return true;
+  return Error::success();
 }
 
 unsigned int TestClient::GetPcRegisterId() {
@@ -187,12 +213,12 @@ unsigned int TestClient::GetPcRegisterId() {
   for (unsigned int register_id = 0;; register_id++) {
     std::string message = formatv("qRegisterInfo{0:x-}", register_id).str();
     std::string response;
-    if (!SendMessage(message, response)) {
+    if (SendMessage(message, response)) {
       GTEST_LOG_(ERROR) << "Unable to query register ID for PC register.";
       return UINT_MAX;
     }
 
-    auto elements_or_error = SplitPairList("GetPcRegisterId", response);
+    auto elements_or_error = SplitUniquePairList("GetPcRegisterId", response);
     if (auto split_error = elements_or_error.takeError()) {
       GTEST_LOG_(ERROR) << "GetPcRegisterId: Error splitting response: "
                         << response;
@@ -209,78 +235,41 @@ unsigned int TestClient::GetPcRegisterId() {
   return m_pc_register;
 }
 
-bool TestClient::Continue(StringRef message) {
-  if (!m_process_info.hasValue()) {
-    GTEST_LOG_(ERROR) << "Continue() called before m_process_info initialized.";
-    return false;
-  }
-
+llvm::Error TestClient::QueryProcessInfo() {
   std::string response;
-  if (!SendMessage(message, response))
-    return false;
-  auto creation = StopReply::Create(response, m_process_info->GetEndian());
-  if (auto create_error = creation.takeError()) {
-    GTEST_LOG_(ERROR) << toString(std::move(create_error));
-    return false;
-  }
-
-  m_stop_reply = std::move(*creation);
-  return true;
+  if (Error E = SendMessage("qProcessInfo", response))
+    return E;
+  auto create_or_error = ProcessInfo::Create(response);
+  if (!create_or_error)
+    return create_or_error.takeError();
+  m_process_info = *create_or_error;
+  return Error::success();
 }
 
-std::string TestClient::GenerateLogFileName(const ArchSpec &arch) const {
-  char *log_directory = getenv("LOG_FILE_DIRECTORY");
-  if (!log_directory)
-    return "";
+Error TestClient::Continue(StringRef message) {
+  assert(m_process_info.hasValue());
 
-  if (!llvm::sys::fs::is_directory(log_directory)) {
-    GTEST_LOG_(WARNING) << "Cannot access log directory: " << log_directory;
-    return "";
-  }
-
-  std::string log_file_name;
-  raw_string_ostream log_file(log_file_name);
-  log_file << log_directory << "/lldb-" << m_test_case_name << '-'
-           << m_test_name << '-' << arch.GetArchitectureName() << ".log";
-  return log_file.str();
-}
+  std::string response;
+  if (Error E = SendMessage(message, response))
+    return E;
+  auto creation = StopReply::create(response, m_process_info->GetEndian());
+  if (Error E = creation.takeError())
+    return E;
 
-std::string TestClient::FormatFailedResult(const std::string &message,
-                                           PacketResult result) {
-  std::string formatted_error;
-  raw_string_ostream error_stream(formatted_error);
-  error_stream << "Failure sending message: " << message << " Result: ";
-
-  switch (result) {
-  case PacketResult::ErrorSendFailed:
-    error_stream << "ErrorSendFailed";
-    break;
-  case PacketResult::ErrorSendAck:
-    error_stream << "ErrorSendAck";
-    break;
-  case PacketResult::ErrorReplyFailed:
-    error_stream << "ErrorReplyFailed";
-    break;
-  case PacketResult::ErrorReplyTimeout:
-    error_stream << "ErrorReplyTimeout";
-    break;
-  case PacketResult::ErrorReplyInvalid:
-    error_stream << "ErrorReplyInvalid";
-    break;
-  case PacketResult::ErrorReplyAck:
-    error_stream << "ErrorReplyAck";
-    break;
-  case PacketResult::ErrorDisconnected:
-    error_stream << "ErrorDisconnected";
-    break;
-  case PacketResult::ErrorNoSequenceLock:
-    error_stream << "ErrorNoSequenceLock";
-    break;
-  default:
-    error_stream << "Unknown Error";
+  m_stop_reply = std::move(*creation);
+  if (!isa<StopReplyStop>(m_stop_reply)) {
+    StringExtractorGDBRemote R;
+    PacketResult result = ReadPacket(R, GetPacketTimeout(), false);
+    if (result != PacketResult::ErrorDisconnected) {
+      return make_error<StringError>(
+          formatv("Expected connection close after receiving {0}. Got {1}/{2} "
+                  "instead.",
+                  response, result, R.GetStringRef())
+              .str(),
+          inconvertibleErrorCode());
+    }
   }
-
-  error_stream.str();
-  return formatted_error;
+  return Error::success();
 }
+
 } // namespace llgs_tests
diff --git a/unittests/tools/lldb-server/tests/TestClient.h b/unittests/tools/lldb-server/tests/TestClient.h
index ae620cc972078..4fa1bbb8ac935 100644
--- a/unittests/tools/lldb-server/tests/TestClient.h
+++ b/unittests/tools/lldb-server/tests/TestClient.h
@@ -7,55 +7,75 @@
 //
 //===----------------------------------------------------------------------===//
 
+#ifndef LLDB_SERVER_TESTS_TESTCLIENT_H
+#define LLDB_SERVER_TESTS_TESTCLIENT_H
+
 #include "MessageObjects.h"
 #include "Plugins/Process/gdb-remote/GDBRemoteCommunicationClient.h"
-#include "lldb/Core/ArchSpec.h"
 #include "lldb/Target/ProcessLaunchInfo.h"
+#include "lldb/Utility/ArchSpec.h"
+#include "lldb/Utility/Connection.h"
 #include "llvm/ADT/Optional.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/FormatVariadic.h"
 #include <memory>
 #include <string>
 
 namespace llgs_tests {
-// TODO: Make the test client an abstract base class, with different children
-// for different types of connections: llgs v. debugserver
 class TestClient
     : public lldb_private::process_gdb_remote::GDBRemoteCommunicationClient {
 public:
-  static void Initialize();
-  TestClient(const std::string &test_name, const std::string &test_case_name);
-  virtual ~TestClient();
-  LLVM_NODISCARD bool StartDebugger();
-  LLVM_NODISCARD bool StopDebugger();
-  LLVM_NODISCARD bool SetInferior(llvm::ArrayRef<std::string> inferior_args);
-  LLVM_NODISCARD bool ListThreadsInStopReply();
-  LLVM_NODISCARD bool SetBreakpoint(unsigned long address);
-  LLVM_NODISCARD bool ContinueAll();
-  LLVM_NODISCARD bool ContinueThread(unsigned long thread_id);
+  static bool IsDebugServer();
+  static bool IsLldbServer();
+
+  /// Launches the server, connects it to the client and returns the client. If
+  /// Log is non-empty, the server will write it's log to this file.
+  static llvm::Expected<std::unique_ptr<TestClient>> launch(llvm::StringRef Log);
+
+  /// Launches the server, while specifying the inferior on its command line.
+  /// When the client connects, it already has a process ready.
+  static llvm::Expected<std::unique_ptr<TestClient>>
+  launch(llvm::StringRef Log, llvm::ArrayRef<llvm::StringRef> InferiorArgs);
+
+  ~TestClient() override;
+  llvm::Error SetInferior(llvm::ArrayRef<std::string> inferior_args);
+  llvm::Error ListThreadsInStopReply();
+  llvm::Error SetBreakpoint(unsigned long address);
+  llvm::Error ContinueAll();
+  llvm::Error ContinueThread(unsigned long thread_id);
   const ProcessInfo &GetProcessInfo();
   llvm::Optional<JThreadsInfo> GetJThreadsInfo();
   const StopReply &GetLatestStopReply();
-  LLVM_NODISCARD bool SendMessage(llvm::StringRef message);
-  LLVM_NODISCARD bool SendMessage(llvm::StringRef message,
-                                  std::string &response_string);
-  LLVM_NODISCARD bool SendMessage(llvm::StringRef message,
-                                  std::string &response_string,
-                                  PacketResult expected_result);
+  template <typename T> llvm::Expected<const T &> GetLatestStopReplyAs() {
+    assert(m_stop_reply);
+    if (const auto *Reply = llvm::dyn_cast<T>(m_stop_reply.get()))
+      return *Reply;
+    return llvm::make_error<llvm::StringError>(
+        llvm::formatv("Unexpected Stop Reply {0}", m_stop_reply->getKind()),
+        llvm::inconvertibleErrorCode());
+  }
+  llvm::Error SendMessage(llvm::StringRef message);
+  llvm::Error SendMessage(llvm::StringRef message,
+                          std::string &response_string);
+  llvm::Error SendMessage(llvm::StringRef message, std::string &response_string,
+                          PacketResult expected_result);
   unsigned int GetPcRegisterId();
 
 private:
-  LLVM_NODISCARD bool Continue(llvm::StringRef message);
-  LLVM_NODISCARD bool GenerateConnectionAddress(std::string &address);
-  std::string GenerateLogFileName(const lldb_private::ArchSpec &arch) const;
+  TestClient(std::unique_ptr<lldb_private::Connection> Conn);
+
+  llvm::Error QueryProcessInfo();
+  llvm::Error Continue(llvm::StringRef message);
   std::string FormatFailedResult(
       const std::string &message,
       lldb_private::process_gdb_remote::GDBRemoteCommunication::PacketResult
           result);
 
   llvm::Optional<ProcessInfo> m_process_info;
-  llvm::Optional<StopReply> m_stop_reply;
-  lldb_private::ProcessLaunchInfo m_server_process_info;
-  std::string m_test_name;
-  std::string m_test_case_name;
-  unsigned int m_pc_register;
+  std::unique_ptr<StopReply> m_stop_reply;
+  unsigned int m_pc_register = UINT_MAX;
 };
+
 } // namespace llgs_tests
+
+#endif // LLDB_SERVER_TESTS_TESTCLIENT_H
diff --git a/unittests/tools/lldb-server/tests/ThreadIdsInJstopinfoTest.cpp b/unittests/tools/lldb-server/tests/ThreadIdsInJstopinfoTest.cpp
index 961b0a3b3167c..6ff777bd17abf 100644
--- a/unittests/tools/lldb-server/tests/ThreadIdsInJstopinfoTest.cpp
+++ b/unittests/tools/lldb-server/tests/ThreadIdsInJstopinfoTest.cpp
@@ -7,38 +7,33 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "TestBase.h"
 #include "TestClient.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Testing/Support/Error.h"
 #include "gtest/gtest.h"
 #include <string>
 
 using namespace llgs_tests;
+using namespace llvm;
 
-class ThreadsInJstopinfoTest : public ::testing::Test {
-protected:
-  virtual void SetUp() { TestClient::Initialize(); }
-};
+TEST_F(StandardStartupTest, TestStopReplyContainsThreadPcs) {
+  // This inferior spawns 4 threads, then forces a break.
+  ASSERT_THAT_ERROR(
+      Client->SetInferior({getInferiorPath("thread_inferior"), "4"}),
+      Succeeded());
 
-TEST_F(ThreadsInJstopinfoTest, TestStopReplyContainsThreadPcsLlgs) {
-  std::vector<std::string> inferior_args;
-  // This inferior spawns N threads, then forces a break.
-  inferior_args.push_back(THREAD_INFERIOR);
-  inferior_args.push_back("4");
-
-  auto test_info = ::testing::UnitTest::GetInstance()->current_test_info();
-
-  TestClient client(test_info->name(), test_info->test_case_name());
-  ASSERT_TRUE(client.StartDebugger());
-  ASSERT_TRUE(client.SetInferior(inferior_args));
-  ASSERT_TRUE(client.ListThreadsInStopReply());
-  ASSERT_TRUE(client.ContinueAll());
-  unsigned int pc_reg = client.GetPcRegisterId();
+  ASSERT_THAT_ERROR(Client->ListThreadsInStopReply(), Succeeded());
+  ASSERT_THAT_ERROR(Client->ContinueAll(), Succeeded());
+  unsigned int pc_reg = Client->GetPcRegisterId();
   ASSERT_NE(pc_reg, UINT_MAX);
 
-  auto jthreads_info = client.GetJThreadsInfo();
+  auto jthreads_info = Client->GetJThreadsInfo();
   ASSERT_TRUE(jthreads_info);
 
-  auto stop_reply = client.GetLatestStopReply();
-  auto stop_reply_pcs = stop_reply.GetThreadPcs();
+  auto stop_reply = Client->GetLatestStopReplyAs<StopReplyStop>();
+  ASSERT_THAT_EXPECTED(stop_reply, Succeeded());
+  auto stop_reply_pcs = stop_reply->getThreadPcs();
   auto thread_infos = jthreads_info->GetThreadInfos();
   ASSERT_EQ(stop_reply_pcs.size(), thread_infos.size())
       << "Thread count mismatch.";
@@ -47,12 +42,9 @@ TEST_F(ThreadsInJstopinfoTest, TestStopReplyContainsThreadPcsLlgs) {
     unsigned long tid = stop_reply_pc.first;
     ASSERT_TRUE(thread_infos.find(tid) != thread_infos.end())
         << "Thread ID: " << tid << " not in JThreadsInfo.";
-    uint64_t pc_value;
-    ASSERT_TRUE(thread_infos[tid].ReadRegisterAsUint64(pc_reg, pc_value))
-        << "Failure reading ThreadInfo register " << pc_reg;
-    ASSERT_EQ(stop_reply_pcs[tid], pc_value)
+    auto pc_value = thread_infos[tid].ReadRegisterAsUint64(pc_reg);
+    ASSERT_THAT_EXPECTED(pc_value, Succeeded());
+    ASSERT_EQ(stop_reply_pcs[tid], *pc_value)
         << "Mismatched PC for thread: " << tid;
   }
-
-  ASSERT_TRUE(client.StopDebugger());
 }