Enhance WrapperHandler

- docstrings
- add try_call_all method
- update tests

src/app/utils/wrapper_handler.py

@@ -6,7 +6,24 @@ W = TypeVar("W")
 T = TypeVar("T")
 
 class WrapperHandler(Generic[W]):
+    """
+    A handler for managing multiple wrappers with retry logic.
+    It attempts to call a function on the current wrapper, and if it fails,
+    it retries a specified number of times before switching to the next wrapper.
+    If all wrappers fail, it raises an exception.
+
+    Note: use `build_wrappers` to create an instance of this class for better error handling.
+    """
+
     def __init__(self, wrappers: list[W], try_per_wrapper: int = 3, retry_delay: int = 2):
+        """
+        Initializes the WrapperHandler with a list of wrappers and retry settings.\n
+        Use `build_wrappers` to create an instance of this class for better error handling.
+        Args:
+            wrappers (list[W]): A list of wrapper instances to manage.
+            try_per_wrapper (int): Number of retries per wrapper before switching to the next.
+            retry_delay (int): Delay in seconds between retries.
+        """
         self.wrappers = wrappers
         self.retry_per_wrapper = try_per_wrapper
         self.retry_delay = retry_delay
@@ -14,9 +31,19 @@ class WrapperHandler(Generic[W]):
         self.retry_count = 0
 
     def try_call(self, func: Callable[[W], T]) -> T:
+        """
+        Attempts to call the provided function on the current wrapper.
+        If it fails, it retries a specified number of times before switching to the next wrapper.
+        If all wrappers fail, it raises an exception.
+        Args:
+            func (Callable[[W], T]): A function that takes a wrapper and returns a result.
+        Returns:
+            T: The result of the function call.
+        Raises:
+            Exception: If all wrappers fail after retries.
+        """
         iterations = 0
         while iterations < len(self.wrappers):
-            print(f"Trying wrapper {self.index}")
             try:
                 wrapper = self.wrappers[self.index]
                 result = func(wrapper)
@@ -24,7 +51,6 @@ class WrapperHandler(Generic[W]):
                 return result
             except Exception as e:
                 self.retry_count += 1
-                print(f"Error occurred {self.retry_count}/{self.retry_per_wrapper}: {e}")
                 if self.retry_count >= self.retry_per_wrapper:
                     self.index = (self.index + 1) % len(self.wrappers)
                     self.retry_count = 0
@@ -35,8 +61,44 @@ class WrapperHandler(Generic[W]):
 
         raise Exception(f"All wrappers failed after retries")
 
+    def try_call_all(self, func: Callable[[W], T]) -> list[T]:
+        """
+        Calls the provided function on all wrappers, collecting results.
+        If a wrapper fails, it logs a warning and continues with the next.
+        If all wrappers fail, it raises an exception.
+        Args:
+            func (Callable[[W], T]): A function that takes a wrapper and returns a result.
+        Returns:
+            list[T]: A list of results from the function calls.
+        Raises:
+            Exception: If all wrappers fail.
+        """
+        results = []
+        for wrapper in self.wrappers:
+            try:
+                result = func(wrapper)
+                results.append(result)
+            except Exception as e:
+                log_warning(f"{wrapper} failed: {e}")
+        if not results:
+            raise Exception("All wrappers failed")
+        return results
+
     @staticmethod
     def build_wrappers(constructors: Iterable[Type[W]], try_per_wrapper: int = 3, retry_delay: int = 2) -> 'WrapperHandler[W]':
+        """
+        Builds a WrapperHandler instance with the given wrapper constructors.
+        It attempts to initialize each wrapper and logs a warning if any cannot be initialized.
+        Only successfully initialized wrappers are included in the handler.
+        Args:
+            constructors (Iterable[Type[W]]): An iterable of wrapper classes to instantiate. e.g. [WrapperA, WrapperB]
+            try_per_wrapper (int): Number of retries per wrapper before switching to the next.
+            retry_delay (int): Delay in seconds between retries.
+        Returns:
+            WrapperHandler[W]: An instance of WrapperHandler with the initialized wrappers.
+        Raises:
+            Exception: If no wrappers could be initialized.
+        """
         result = []
         for wrapper_class in constructors:
             try:

tests/utils/test_wrapper_handler.py

@@ -58,3 +58,22 @@ class TestWrapperHandler:
         assert result == "Success"
         assert handler.index == 1  # Should return to the second wrapper after failure
         assert handler.retry_count == 0
+
+    def test_try_call_all(self):
+        wrappers = [FailingWrapper, MockWrapper, FailingWrapper]
+        handler: WrapperHandler[MockWrapper] = WrapperHandler.build_wrappers(wrappers, try_per_wrapper=1, retry_delay=0)
+
+        results = handler.try_call_all(lambda w: w.do_something())
+        assert results == ["Success"]  # Only the second wrapper should succeed
+
+        wrappers = [FailingWrapper, MockWrapper, FailingWrapper, MockWrapper]
+        handler: WrapperHandler[MockWrapper] = WrapperHandler.build_wrappers(wrappers, try_per_wrapper=1, retry_delay=0)
+
+        results = handler.try_call_all(lambda w: w.do_something())
+        assert results == ["Success", "Success"]  # Only the second and fourth wrappers should succeed
+
+        # Test when all wrappers fail
+        handler_all_fail: WrapperHandler[MockWrapper] = WrapperHandler.build_wrappers([FailingWrapper, FailingWrapper], try_per_wrapper=1, retry_delay=0)
+        with pytest.raises(Exception) as exc_info:
+            handler_all_fail.try_call_all(lambda w: w.do_something())
+        assert "All wrappers failed" in str(exc_info.value)