From: Kienan Stewart <kstewart@efficios.com>
Date: Thu, 10 Oct 2024 19:36:17 +0000 (+0000)
Subject: Tests: Add wait_until_disconnected() to python test LiveViewer
X-Git-Url: http://git.lttng.org./?a=commitdiff_plain;h=7b874e71ef90c4e0e0d5bc9a5ac697ed0ff09da5;p=lttng-tools.git

Tests: Add wait_until_disconnected() to python test LiveViewer

Change-Id: I70cf5e46a08c658b3fecbcea9317e1b2d9f769fb
Signed-off-by: Kienan Stewart <kstewart@efficios.com>
Signed-off-by: Jérémie Galarneau <jeremie.galarneau@efficios.com>
---

diff --git a/tests/utils/lttngtest/environment.py b/tests/utils/lttngtest/environment.py
index 65d74a41d..ec64c0a53 100644
--- a/tests/utils/lttngtest/environment.py
+++ b/tests/utils/lttngtest/environment.py
@@ -181,10 +181,10 @@ class _LiveViewer:
                 return True
         return False
 
-    def wait_until_connected(self, timeout=0):
-        connected = False
+    def _wait_until(self, desired_state: bool, timeout=0):
+        connected_state = not desired_state
         started = time.time()
-        while not connected:
+        while connected_state != desired_state:
             try:
                 if timeout != 0 and (time.time() - started) > timeout:
                     raise RuntimeError(
@@ -193,11 +193,17 @@ class _LiveViewer:
                         )
                     )
 
-                connected = self.is_connected()
+                connected_state = self.is_connected()
             except bt2._Error:
                 time.sleep(0.01)
                 continue
-        return connected
+        return connected_state
+
+    def wait_until_disconnected(self, timeout=0):
+        return self._wait_until(False, timeout)
+
+    def wait_until_connected(self, timeout=0):
+        return self._wait_until(True, timeout)
 
     def wait(self):
         if self._live_iterator: