|
| 1 | +"""Tests for execute_rest() throttle retry logic in rest_api.py. |
| 2 | +
|
| 3 | +Verifies: |
| 4 | +- Normal (non-throttled) requests are unaffected |
| 5 | +- Throttled requests retry up to 3 times with exponential backoff |
| 6 | +- KeeperApiError raised after max retries |
| 7 | +- --fail-on-throttle skips retries entirely |
| 8 | +- Server's "try again in X" message is parsed correctly (seconds + minutes) |
| 9 | +- Server wait capped at 300s |
| 10 | +- Backoff takes the larger of server hint vs exponential schedule |
| 11 | +""" |
| 12 | + |
| 13 | +import json |
| 14 | +import os |
| 15 | +import sys |
| 16 | +import unittest |
| 17 | +from unittest.mock import patch, MagicMock |
| 18 | + |
| 19 | +# Add parent dir so imports work from unit-tests/ |
| 20 | +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) |
| 21 | + |
| 22 | +from keepercommander.rest_api import execute_rest |
| 23 | +from keepercommander.error import KeeperApiError |
| 24 | +from keepercommander.params import RestApiContext |
| 25 | +from keepercommander.proto import APIRequest_pb2 as proto |
| 26 | + |
| 27 | + |
| 28 | +def make_context(fail_on_throttle=False): |
| 29 | + """Create a real RestApiContext with QRC disabled to simplify mocking.""" |
| 30 | + ctx = RestApiContext(server='https://keepersecurity.com', locale='en_US') |
| 31 | + ctx.transmission_key = os.urandom(32) |
| 32 | + ctx.server_key_id = 7 |
| 33 | + ctx.fail_on_throttle = fail_on_throttle |
| 34 | + ctx.disable_qrc() # Skip QRC key negotiation |
| 35 | + return ctx |
| 36 | + |
| 37 | + |
| 38 | +def make_throttle_response(message="Please try again in 1 minute"): |
| 39 | + """Build a fake HTTP 403 throttle response.""" |
| 40 | + body = {"error": "throttled", "message": message} |
| 41 | + resp = MagicMock() |
| 42 | + resp.status_code = 403 |
| 43 | + resp.headers = {'Content-Type': 'application/json'} |
| 44 | + resp.json.return_value = body |
| 45 | + resp.content = json.dumps(body).encode() |
| 46 | + return resp |
| 47 | + |
| 48 | + |
| 49 | +def make_success_response(): |
| 50 | + """Build a fake HTTP 200 response with empty body.""" |
| 51 | + resp = MagicMock() |
| 52 | + resp.status_code = 200 |
| 53 | + resp.headers = {'Content-Type': 'application/octet-stream'} |
| 54 | + resp.content = b'' |
| 55 | + return resp |
| 56 | + |
| 57 | + |
| 58 | +def make_payload(): |
| 59 | + """Create a minimal ApiRequestPayload for execute_rest.""" |
| 60 | + return proto.ApiRequestPayload() |
| 61 | + |
| 62 | + |
| 63 | +class TestThrottleRetry(unittest.TestCase): |
| 64 | + """Tests for the bounded retry with exponential backoff on 403 throttle.""" |
| 65 | + |
| 66 | + @patch('keepercommander.rest_api.time.sleep') |
| 67 | + @patch('keepercommander.rest_api.requests.post') |
| 68 | + def test_normal_request_unaffected(self, mock_post, mock_sleep): |
| 69 | + """Non-throttled 200 response should pass through with no retries.""" |
| 70 | + mock_post.return_value = make_success_response() |
| 71 | + |
| 72 | + execute_rest(make_context(), 'test/endpoint', make_payload()) |
| 73 | + |
| 74 | + self.assertEqual(mock_post.call_count, 1) |
| 75 | + mock_sleep.assert_not_called() |
| 76 | + |
| 77 | + @patch('keepercommander.rest_api.time.sleep') |
| 78 | + @patch('keepercommander.rest_api.requests.post') |
| 79 | + def test_retries_then_succeeds(self, mock_post, mock_sleep): |
| 80 | + """Throttle twice, succeed on 3rd attempt.""" |
| 81 | + mock_post.side_effect = [ |
| 82 | + make_throttle_response("try again in 30 seconds"), |
| 83 | + make_throttle_response("try again in 30 seconds"), |
| 84 | + make_success_response(), |
| 85 | + ] |
| 86 | + |
| 87 | + execute_rest(make_context(), 'test/endpoint', make_payload()) |
| 88 | + |
| 89 | + self.assertEqual(mock_post.call_count, 3) |
| 90 | + self.assertEqual(mock_sleep.call_count, 2) |
| 91 | + # 1st: max(30, 30*2^0=30) = 30 |
| 92 | + # 2nd: max(30, 30*2^1=60) = 60 |
| 93 | + calls = [c.args[0] for c in mock_sleep.call_args_list] |
| 94 | + self.assertEqual(calls, [30, 60]) |
| 95 | + |
| 96 | + @patch('keepercommander.rest_api.time.sleep') |
| 97 | + @patch('keepercommander.rest_api.requests.post') |
| 98 | + def test_raises_after_max_retries(self, mock_post, mock_sleep): |
| 99 | + """Always throttled — should raise KeeperApiError after 3 retries.""" |
| 100 | + mock_post.return_value = make_throttle_response("try again in 1 minute") |
| 101 | + |
| 102 | + with self.assertRaises(KeeperApiError): |
| 103 | + execute_rest(make_context(), 'test/endpoint', make_payload()) |
| 104 | + |
| 105 | + # 1 initial + 3 retries = 4 posts, error raised when retry 4 > max 3 |
| 106 | + self.assertEqual(mock_post.call_count, 4) |
| 107 | + self.assertEqual(mock_sleep.call_count, 3) |
| 108 | + |
| 109 | + @patch('keepercommander.rest_api.time.sleep') |
| 110 | + @patch('keepercommander.rest_api.requests.post') |
| 111 | + def test_fail_on_throttle_skips_retry(self, mock_post, mock_sleep): |
| 112 | + """--fail-on-throttle should return error immediately with no retries.""" |
| 113 | + mock_post.return_value = make_throttle_response() |
| 114 | + |
| 115 | + result = execute_rest(make_context(fail_on_throttle=True), 'test/endpoint', make_payload()) |
| 116 | + |
| 117 | + # When fail_on_throttle=True, the throttle block is skipped and the |
| 118 | + # failure dict is returned directly (no retry, no exception) |
| 119 | + self.assertEqual(result.get('error'), 'throttled') |
| 120 | + self.assertEqual(mock_post.call_count, 1) |
| 121 | + mock_sleep.assert_not_called() |
| 122 | + |
| 123 | + @patch('keepercommander.rest_api.time.sleep') |
| 124 | + @patch('keepercommander.rest_api.requests.post') |
| 125 | + def test_parses_seconds_hint(self, mock_post, mock_sleep): |
| 126 | + """Server says 'try again in 45 seconds' — wait should be 45s.""" |
| 127 | + mock_post.side_effect = [ |
| 128 | + make_throttle_response("try again in 45 seconds"), |
| 129 | + make_success_response(), |
| 130 | + ] |
| 131 | + |
| 132 | + execute_rest(make_context(), 'test/endpoint', make_payload()) |
| 133 | + |
| 134 | + # max(45, 30*2^0=30) = 45 |
| 135 | + mock_sleep.assert_called_once_with(45) |
| 136 | + |
| 137 | + @patch('keepercommander.rest_api.time.sleep') |
| 138 | + @patch('keepercommander.rest_api.requests.post') |
| 139 | + def test_parses_minutes_hint(self, mock_post, mock_sleep): |
| 140 | + """Server says 'try again in 2 minutes' — wait should be 120s.""" |
| 141 | + mock_post.side_effect = [ |
| 142 | + make_throttle_response("try again in 2 minutes"), |
| 143 | + make_success_response(), |
| 144 | + ] |
| 145 | + |
| 146 | + execute_rest(make_context(), 'test/endpoint', make_payload()) |
| 147 | + |
| 148 | + # max(120, 30*2^0=30) = 120 |
| 149 | + mock_sleep.assert_called_once_with(120) |
| 150 | + |
| 151 | + @patch('keepercommander.rest_api.time.sleep') |
| 152 | + @patch('keepercommander.rest_api.requests.post') |
| 153 | + def test_caps_server_wait_at_300s(self, mock_post, mock_sleep): |
| 154 | + """Server says 'try again in 49 minutes' — capped to 300s.""" |
| 155 | + mock_post.side_effect = [ |
| 156 | + make_throttle_response("try again in 49 minutes"), |
| 157 | + make_success_response(), |
| 158 | + ] |
| 159 | + |
| 160 | + execute_rest(make_context(), 'test/endpoint', make_payload()) |
| 161 | + |
| 162 | + # min(2940, 300)=300; max(300, 30*2^0=30) = 300 |
| 163 | + mock_sleep.assert_called_once_with(300) |
| 164 | + |
| 165 | + @patch('keepercommander.rest_api.time.sleep') |
| 166 | + @patch('keepercommander.rest_api.requests.post') |
| 167 | + def test_exponential_backoff_progression(self, mock_post, mock_sleep): |
| 168 | + """Verify backoff doubles: 30s, 60s, 120s when server hint is small.""" |
| 169 | + mock_post.return_value = make_throttle_response("try again in 10 seconds") |
| 170 | + |
| 171 | + with self.assertRaises(KeeperApiError): |
| 172 | + execute_rest(make_context(), 'test/endpoint', make_payload()) |
| 173 | + |
| 174 | + # Server says 10s, but backoff wins: max(10, 30*2^0)=30, max(10, 30*2^1)=60, max(10, 30*2^2)=120 |
| 175 | + calls = [c.args[0] for c in mock_sleep.call_args_list] |
| 176 | + self.assertEqual(calls, [30, 60, 120]) |
| 177 | + |
| 178 | + @patch('keepercommander.rest_api.time.sleep') |
| 179 | + @patch('keepercommander.rest_api.requests.post') |
| 180 | + def test_no_message_defaults_to_60s(self, mock_post, mock_sleep): |
| 181 | + """Missing 'try again' text defaults to 60s server hint.""" |
| 182 | + mock_post.side_effect = [ |
| 183 | + make_throttle_response("Rate limit exceeded"), # no "try again in X" |
| 184 | + make_success_response(), |
| 185 | + ] |
| 186 | + |
| 187 | + execute_rest(make_context(), 'test/endpoint', make_payload()) |
| 188 | + |
| 189 | + # Default 60s; max(60, 30*2^0=30) = 60 |
| 190 | + mock_sleep.assert_called_once_with(60) |
| 191 | + |
| 192 | + |
| 193 | +if __name__ == '__main__': |
| 194 | + unittest.main() |
0 commit comments