1 files changed, 19 insertions, 3 deletions
diff --git a/unittests/tools/lldb-server/tests/TestClient.h b/unittests/tools/lldb-server/tests/TestClient.h
index b5195307876d..ebcb4a549b39 100644
--- a/unittests/tools/lldb-server/tests/TestClient.h
+++ b/unittests/tools/lldb-server/tests/TestClient.h
@@ -59,7 +59,7 @@ public:
   llvm::Error ContinueAll();
   llvm::Error ContinueThread(unsigned long thread_id);
   const ProcessInfo &GetProcessInfo();
-  llvm::Optional<JThreadsInfo> GetJThreadsInfo();
+  llvm::Expected<JThreadsInfo> GetJThreadsInfo();
   const StopReply &GetLatestStopReply();
   template <typename T> llvm::Expected<const T &> GetLatestStopReplyAs() {
     assert(m_stop_reply);
@@ -74,12 +74,18 @@ public:
                           std::string &response_string);
   llvm::Error SendMessage(llvm::StringRef message, std::string &response_string,
                           PacketResult expected_result);
+
+  template <typename P, typename... CreateArgs>
+  llvm::Expected<typename P::result_type> SendMessage(llvm::StringRef Message,
+                                                      CreateArgs &&... Args);
   unsigned int GetPcRegisterId();
 
 private:
   TestClient(std::unique_ptr<lldb_private::Connection> Conn);
 
-  llvm::Error QueryProcessInfo();
+  llvm::Error qProcessInfo();
+  llvm::Error qRegisterInfos();
+  llvm::Error queryProcess();
   llvm::Error Continue(llvm::StringRef message);
   std::string FormatFailedResult(
       const std::string &message,
@@ -88,9 +94,19 @@ private:
 
   llvm::Optional<ProcessInfo> m_process_info;
   std::unique_ptr<StopReply> m_stop_reply;
-  unsigned int m_pc_register = UINT_MAX;
+  std::vector<lldb_private::RegisterInfo> m_register_infos;
+  unsigned int m_pc_register = LLDB_INVALID_REGNUM;
 };
 
+template <typename P, typename... CreateArgs>
+llvm::Expected<typename P::result_type>
+TestClient::SendMessage(llvm::StringRef Message, CreateArgs &&... Args) {
+  std::string ResponseText;
+  if (llvm::Error E = SendMessage(Message, ResponseText))
+    return std::move(E);
+  return P::create(ResponseText, std::forward<CreateArgs>(Args)...);
+}
+
 } // namespace llgs_tests
 
 #endif // LLDB_SERVER_TESTS_TESTCLIENT_H