From f70efb604de1241e4e7d795ae5ef1d65ee1f6b4c Mon Sep 17 00:00:00 2001
From: mulhern <amulhern@redhat.com>
Date: Wed, 18 Sep 2024 14:16:41 -0400
Subject: [PATCH] Use tenacity for wait_for_pools

Signed-off-by: mulhern <amulhern@redhat.com>
---
 tests/client-dbus/tests/udev/_utils.py | 64 +++++++++++++++-----------
 1 file changed, 38 insertions(+), 26 deletions(-)

diff --git a/tests/client-dbus/tests/udev/_utils.py b/tests/client-dbus/tests/udev/_utils.py
index 1fc9aa589e..aaf2911769 100644
--- a/tests/client-dbus/tests/udev/_utils.py
+++ b/tests/client-dbus/tests/udev/_utils.py
@@ -31,7 +31,15 @@
 import dbus
 import psutil
 import pyudev
-from tenacity import Retrying, retry_if_exception_type, stop_after_delay, wait_fixed
+from tenacity import (
+    RetryError,
+    Retrying,
+    retry_if_exception_type,
+    retry_if_not_result,
+    stop_after_attempt,
+    stop_after_delay,
+    wait_fixed,
+)
 
 # isort: LOCAL
 from stratisd_client_dbus import (
@@ -535,32 +543,36 @@ def wait_for_pools(self, expected_num, *, name=None):
         :return: list of pool information found
         :rtype: list of (str * MOPool)
         """
-        (count, limit, dbus_err, found_num, known_pools, start_time) = (
-            0,
-            expected_num + 1,
-            None,
-            None,
-            None,
-            time.time(),
-        )
-        while count < limit and not expected_num == found_num:
-            try:
-                known_pools = get_pools(name=name)
-            except dbus.exceptions.DBusException as err:
-                dbus_err = err
-
-            if known_pools is not None:
-                found_num = len(known_pools)
-
-            time.sleep(3)
-            count += 1
-
-        if found_num is None and dbus_err is not None:
+        try:
+            for attempt in Retrying(
+                retry=(
+                    retry_if_exception_type(dbus.exceptions.DBusException)
+                    | retry_if_not_result(
+                        lambda known_pools: len(known_pools) == expected_num
+                    )
+                ),
+                wait=wait_fixed(3),
+                stop=stop_after_attempt(expected_num + 1),
+                reraise=True,
+            ):
+                with attempt:
+                    known_pools = get_pools(name=name)
+                attempt.retry_state.set_result(known_pools)
+        except dbus.exceptions.DBusException as err:
             raise RuntimeError(
-                f"After {time.time() - start_time:.2f} seconds, the only "
-                "response is a D-Bus exception"
-            ) from dbus_err
+                "Failed to obtain any information about pools from the D-Bus"
+            ) from err
+        except RetryError:
+            pass
 
-        self.assertEqual(found_num, expected_num)
+        # At this point, it is not known whether the number of pools is more,
+        # fewer, or equal to the number of pools expected. It is only known
+        # that the D-Bus is returning a result. If the number of pools is
+        # equal to the number of pools expected, that may be due to timing,
+        # and be about to change. So, regardless, try to get the pools one
+        # more time and verify the correct count.
+        time.sleep(3)
+        known_pools = get_pools(name=name)
+        self.assertEqual(len(known_pools), expected_num)
 
         return known_pools