Add until_all_done to CallTrigger (#37)

kuznia-rdzeni · Jan 8, 2025 · a976b44 · a976b44
1 parent 8e0435b
commit a976b44
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/transactron/testing/testbenchio.py b/transactron/testing/testbenchio.py
@@ -88,6 +88,12 @@ async def until_done(self) -> Any:
             if any(res is not None for res in results):
                 return results
 
+    async def until_all_done(self) -> Any:
+        """Same as `until_done` but wait for all results instead of any result."""
+        async for results in self:
+            if all(res is not None for res in results):
+                return results
+
     def __await__(self) -> Generator:
         only_calls = [t for t in self.calls_and_values if isinstance(t, tuple)]
         only_values = [t for t in self.calls_and_values if not isinstance(t, tuple)]