Coverage for /usr/local/lib/python3.12/site-packages/prefect/_internal/testing.py: 0%

1"""Testing utilities for internal use.""" 

2 

3import asyncio 

4from typing import Any, AsyncIterator 

5 

6from typing_extensions import Self 

7 

8 

9class AssertionRetryAttempt: 

10 """Context manager for capturing exceptions during retry attempts.""" 

11 

12 def __init__(self, attempt_number: int): 

13 self.attempt_number = attempt_number 

14 self.exception: Exception | None = None 

15 

16 def __enter__(self) -> Self: 

17 return self 

18 

19 def __exit__( 

20 self, 

21 exc_type: type[BaseException] | None, 

22 exc_val: BaseException | None, 

23 exc_tb: Any, 

24 ) -> bool: 

25 if exc_val is not None: 

26 self.exception = exc_val # type: ignore 

27 return exc_type is AssertionError 

28 

29 

30async def retry_asserts( 

31 max_attempts: int = 3, 

32 delay: float = 1.0, 

33) -> AsyncIterator[AssertionRetryAttempt]: 

34 """ 

35 Async generator that retries a block of assertions until it succeeds or max attempts is reached. 

36 

37 Useful for testing eventual consistency scenarios where changes may not 

38 propagate immediately. 

39 

40 Args: 

41 max_attempts: Maximum number of attempts before raising the exception. 

42 delay: Time in seconds to wait between retry attempts. 

43 

44 Yields: 

45 A context manager that captures exceptions during each attempt. 

46 

47 Raises: 

48 The last exception raised within the block if all attempts fail. 

49 

50 Example: 

51 ```python 

52 async for attempt in retry_asserts(max_attempts=3): 

53 with attempt: 

54 for deployment in deployments: 

55 await session.refresh(deployment) 

56 assert deployment.status == DeploymentStatus.READY 

57 ``` 

58 """ 

59 for attempt_number in range(1, max_attempts + 1): 

60 attempt = AssertionRetryAttempt(attempt_number) 

61 yield attempt 

62 

63 if attempt.exception is None: 

64 return # Success, exit early 

65 

66 if attempt_number == max_attempts: 

67 raise attempt.exception 

68 

69 await asyncio.sleep(delay)