| 126 | asyncio.run(_run()) |
| 127 | |
| 128 | def test_on_status_callback(self) -> None: |
| 129 | async def _run() -> None: |
| 130 | statuses: list[RetryStatusMessage] = [] |
| 131 | |
| 132 | async def operation(attempt: int, ctx: RetryContext) -> str: |
| 133 | if attempt < 2: |
| 134 | err = RateLimitError() |
| 135 | err.headers = {"retry-after": "0.01"} # type: ignore[attr-defined] |
| 136 | raise err |
| 137 | return "ok" |
| 138 | |
| 139 | result = await with_retry( |
| 140 | operation, |
| 141 | RetryOptions(max_retries=5), |
| 142 | on_status=lambda msg: statuses.append(msg), |
| 143 | ) |
| 144 | self.assertEqual(result, "ok") |
| 145 | self.assertEqual(len(statuses), 1) |
| 146 | self.assertEqual(statuses[0].attempt, 1) |
| 147 | |
| 148 | asyncio.run(_run()) |
| 149 | |
| 150 | def test_exhaust_retries(self) -> None: |
| 151 | async def _run() -> None: |