Fix flaky TestDAP_console test. (#94494)

Test Plan:
llvm-lit
llvm-project/lldb/test/API/tools/lldb-dap/console/TestDAP_console.py

2 files changed, 8 insertions, 4 deletions

diff --git a/lldb/packages/Python/lldbsuite/test/tools/lldb-dap/dap_server.py b/lldb/packages/Python/lldbsuite/test/tools/lldb-dap/dap_server.py
index e2126d67a5fe..a9eeec184099 100644
--- a/lldb/packages/Python/lldbsuite/test/tools/lldb-dap/dap_server.py
+++ b/lldb/packages/Python/lldbsuite/test/tools/lldb-dap/dap_server.py
@@ -171,13 +171,15 @@ class DebugCommunication(object):
         self.output_condition.release()
         return output
 
-    def collect_output(self, category, duration, clear=True):
-        end_time = time.time() + duration
+    def collect_output(self, category, timeout_secs, pattern, clear=True):
+        end_time = time.time() + timeout_secs
         collected_output = ""
         while end_time > time.time():
             output = self.get_output(category, timeout=0.25, clear=clear)
             if output:
                 collected_output += output
+                if pattern is not None and pattern in output:
+                    break
         return collected_output if collected_output else None
 
     def enqueue_recv_packet(self, packet):
diff --git a/lldb/packages/Python/lldbsuite/test/tools/lldb-dap/lldbdap_testcase.py b/lldb/packages/Python/lldbsuite/test/tools/lldb-dap/lldbdap_testcase.py
index d56ea5dca14b..600dc2b5fe9b 100644
--- a/lldb/packages/Python/lldbsuite/test/tools/lldb-dap/lldbdap_testcase.py
+++ b/lldb/packages/Python/lldbsuite/test/tools/lldb-dap/lldbdap_testcase.py
@@ -195,8 +195,10 @@ class DAPTestCaseBase(TestBase):
     def get_console(self, timeout=0.0):
         return self.dap_server.get_output("console", timeout=timeout)
 
-    def collect_console(self, duration):
-        return self.dap_server.collect_output("console", duration=duration)
+    def collect_console(self, timeout_secs, pattern=None):
+        return self.dap_server.collect_output(
+            "console", timeout_secs=timeout_secs, pattern=pattern
+        )
 
     def get_local_as_int(self, name, threadId=None):
         value = self.dap_server.get_local_variable_value(name, threadId=threadId)