Merge pull request #684 from ag2ai/polish-suppress-in-tests

Polish suppress decorator in tests
ag2ai · Jan 28, 2025 · 267ce37 · 267ce37
2 parents 101cc53 + 55c6559
commit 267ce37
Showing 1 changed file with 14 additions and 28 deletions.
diff --git a/test/conftest.py b/test/conftest.py
@@ -418,7 +418,7 @@ def credentials_from_test_param(request: pytest.FixtureRequest) -> Credentials:
 T = TypeVar("T", bound=Callable[..., Any])
 
 
-def suppress(exception: type[BaseException], *, retries: Optional[int] = None, timeout: int = 60) -> Callable[[T], T]:
+def suppress(exception: type[BaseException], *, retries: int = 0, timeout: int = 60) -> Callable[[T], T]:
     """Suppresses the specified exception and retries the function a specified number of times.
 
     Args:
@@ -429,58 +429,44 @@ def suppress(exception: type[BaseException], *, retries: Optional[int] = None, t
     """
 
     def decorator(
-        func: T, exception: type[BaseException] = exception, retries: Optional[int] = retries, timeout: int = timeout
+        func: T, exception: type[BaseException] = exception, retries: int = retries, timeout: int = timeout
     ) -> T:
         if inspect.iscoroutinefunction(func):
 
             @functools.wraps(func)
             async def wrapper(
                 *args: Any,
                 exception: type[BaseException] = exception,
-                retries: Optional[int] = retries,
+                retries: int = retries,
                 timeout: int = timeout,
                 **kwargs: Any,
             ) -> Any:
-                if retries is None:
+                for i in range(retries + 1):
                     try:
                         return await func(*args, **kwargs)
                     except exception:
-                        pytest.xfail(f"Suppressed '{exception}' raised")
-                        raise
-                else:
-                    for i in range(retries):
-                        try:
-                            return await func(*args, **kwargs)
-                        except exception:
-                            if i >= retries - 1:
-                                pytest.xfail(f"Suppressed '{exception}' raised {i + 1} times")
-                                raise
-                            await asyncio.sleep(timeout)
+                        if i >= retries - 1:
+                            pytest.xfail(f"Suppressed '{exception}' raised {i + 1} times")
+                            raise
+                        await asyncio.sleep(timeout)
         else:
 
             @functools.wraps(func)
             def wrapper(
                 *args: Any,
                 exception: type[BaseException] = exception,
-                retries: Optional[int] = retries,
+                retries: int = retries,
                 timeout: int = timeout,
                 **kwargs: Any,
             ) -> Any:
-                if retries is None:
+                for i in range(retries + 1):
                     try:
                         return func(*args, **kwargs)
                     except exception:
-                        pytest.xfail(f"Suppressed '{exception}' raised")
-                        raise
-                else:
-                    for i in range(retries):
-                        try:
-                            return func(*args, **kwargs)
-                        except exception:
-                            if i >= retries - 1:
-                                pytest.xfail(f"Suppressed '{exception}' raised {i + 1} times")
-                                raise
-                            time.sleep(timeout)
+                        if i >= retries - 1:
+                            pytest.xfail(f"Suppressed '{exception}' raised {i + 1} times")
+                            raise
+                        time.sleep(timeout)
 
         return wrapper  # type: ignore[return-value]