1 files changed, 47 insertions, 27 deletions

diff --git a/unittests/tools/lldb-server/tests/TestClient.h b/unittests/tools/lldb-server/tests/TestClient.h
index ae620cc97207..4fa1bbb8ac93 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