|
| 1 | +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. |
| 2 | +# SPDX-License-Identifier: Apache-2.0 |
| 3 | + |
| 4 | +from asyncio import gather, sleep |
| 5 | + |
| 6 | +import pytest |
| 7 | +from smithy_core.exceptions import CallError, ClientTimeoutError, RetryError |
| 8 | +from smithy_core.interfaces import retries as retries_interface |
| 9 | +from smithy_core.retries import ( |
| 10 | + ExponentialBackoffJitterType, |
| 11 | + ExponentialRetryBackoffStrategy, |
| 12 | + StandardRetryQuota, |
| 13 | + StandardRetryStrategy, |
| 14 | +) |
| 15 | + |
| 16 | + |
| 17 | +# TODO: Refactor this to use a smithy-testing generated client |
| 18 | +async def retry_operation( |
| 19 | + strategy: retries_interface.RetryStrategy, |
| 20 | + responses: list[int | Exception], |
| 21 | +) -> tuple[str, int]: |
| 22 | + token = strategy.acquire_initial_retry_token() |
| 23 | + response_iter = iter(responses) |
| 24 | + |
| 25 | + while True: |
| 26 | + if token.retry_delay: |
| 27 | + await sleep(token.retry_delay) |
| 28 | + |
| 29 | + response = next(response_iter) |
| 30 | + attempt = token.retry_count + 1 |
| 31 | + |
| 32 | + # Success case |
| 33 | + if response == 200: |
| 34 | + strategy.record_success(token=token) |
| 35 | + return "success", attempt |
| 36 | + |
| 37 | + # Error case - either status code or exception |
| 38 | + if isinstance(response, Exception): |
| 39 | + error = response |
| 40 | + else: |
| 41 | + error = CallError( |
| 42 | + fault="server" if response >= 500 else "client", |
| 43 | + message=f"HTTP {response}", |
| 44 | + is_retry_safe=response >= 500, |
| 45 | + ) |
| 46 | + |
| 47 | + try: |
| 48 | + token = strategy.refresh_retry_token_for_retry( |
| 49 | + token_to_renew=token, error=error |
| 50 | + ) |
| 51 | + except RetryError: |
| 52 | + raise error |
| 53 | + |
| 54 | + |
| 55 | +async def test_standard_retry_eventually_succeeds(): |
| 56 | + quota = StandardRetryQuota(initial_capacity=500) |
| 57 | + strategy = StandardRetryStrategy(max_attempts=3, retry_quota=quota) |
| 58 | + |
| 59 | + result, attempts = await retry_operation(strategy, [500, 500, 200]) |
| 60 | + |
| 61 | + assert result == "success" |
| 62 | + assert attempts == 3 |
| 63 | + assert quota.available_capacity == 495 |
| 64 | + |
| 65 | + |
| 66 | +async def test_standard_retry_fails_due_to_max_attempts(): |
| 67 | + quota = StandardRetryQuota(initial_capacity=500) |
| 68 | + strategy = StandardRetryStrategy(max_attempts=3, retry_quota=quota) |
| 69 | + |
| 70 | + with pytest.raises(CallError, match="502"): |
| 71 | + await retry_operation(strategy, [502, 502, 502]) |
| 72 | + |
| 73 | + assert quota.available_capacity == 490 |
| 74 | + |
| 75 | + |
| 76 | +async def test_retry_quota_exhausted_after_single_retry(): |
| 77 | + quota = StandardRetryQuota(initial_capacity=5) |
| 78 | + strategy = StandardRetryStrategy(max_attempts=3, retry_quota=quota) |
| 79 | + |
| 80 | + with pytest.raises(CallError, match="502"): |
| 81 | + await retry_operation(strategy, [500, 502]) |
| 82 | + |
| 83 | + assert quota.available_capacity == 0 |
| 84 | + |
| 85 | + |
| 86 | +async def test_retry_quota_prevents_retries_when_quota_zero(): |
| 87 | + quota = StandardRetryQuota(initial_capacity=0) |
| 88 | + strategy = StandardRetryStrategy(max_attempts=3, retry_quota=quota) |
| 89 | + |
| 90 | + with pytest.raises(CallError, match="500"): |
| 91 | + await retry_operation(strategy, [500]) |
| 92 | + |
| 93 | + assert quota.available_capacity == 0 |
| 94 | + |
| 95 | + |
| 96 | +async def test_retry_quota_stops_retries_when_exhausted(): |
| 97 | + quota = StandardRetryQuota(initial_capacity=10) |
| 98 | + strategy = StandardRetryStrategy(max_attempts=5, retry_quota=quota) |
| 99 | + |
| 100 | + with pytest.raises(CallError, match="503"): |
| 101 | + await retry_operation(strategy, [500, 502, 503]) |
| 102 | + |
| 103 | + assert quota.available_capacity == 0 |
| 104 | + |
| 105 | + |
| 106 | +async def test_retry_quota_recovers_after_successful_responses(): |
| 107 | + quota = StandardRetryQuota(initial_capacity=15) |
| 108 | + strategy = StandardRetryStrategy(max_attempts=5, retry_quota=quota) |
| 109 | + |
| 110 | + # First operation: 2 retries then success |
| 111 | + await retry_operation(strategy, [500, 502, 200]) |
| 112 | + assert quota.available_capacity == 10 |
| 113 | + |
| 114 | + # Second operation: 1 retry then success |
| 115 | + await retry_operation(strategy, [500, 200]) |
| 116 | + assert quota.available_capacity == 10 |
| 117 | + |
| 118 | + |
| 119 | +async def test_retry_quota_shared_across_concurrent_operations(): |
| 120 | + quota = StandardRetryQuota(initial_capacity=500) |
| 121 | + backoff = ExponentialRetryBackoffStrategy( |
| 122 | + backoff_scale_value=1, |
| 123 | + max_backoff=10, |
| 124 | + jitter_type=ExponentialBackoffJitterType.FULL, |
| 125 | + ) |
| 126 | + strategy = StandardRetryStrategy( |
| 127 | + max_attempts=5, |
| 128 | + retry_quota=quota, |
| 129 | + backoff_strategy=backoff, |
| 130 | + ) |
| 131 | + |
| 132 | + result1, result2 = await gather( |
| 133 | + retry_operation(strategy, [500, 500, 200]), |
| 134 | + retry_operation(strategy, [500, 200]), |
| 135 | + ) |
| 136 | + |
| 137 | + assert result1 == ("success", 3) |
| 138 | + assert result2 == ("success", 2) |
| 139 | + assert quota.available_capacity == 495 |
| 140 | + |
| 141 | + |
| 142 | +async def test_retry_quota_handles_timeout_errors(): |
| 143 | + quota = StandardRetryQuota(initial_capacity=500) |
| 144 | + strategy = StandardRetryStrategy(max_attempts=3, retry_quota=quota) |
| 145 | + |
| 146 | + timeout1 = ClientTimeoutError() |
| 147 | + timeout2 = ClientTimeoutError() |
| 148 | + |
| 149 | + result, attempts = await retry_operation(strategy, [timeout1, timeout2, 200]) |
| 150 | + |
| 151 | + assert result == "success" |
| 152 | + assert attempts == 3 |
| 153 | + assert quota.available_capacity == 490 |
0 commit comments