|
| 1 | +#!/usr/bin/env python3 |
| 2 | + |
| 3 | +import unittest |
| 4 | +import time |
| 5 | +import responses |
| 6 | +from unittest.mock import patch, MagicMock |
| 7 | +from zulip import Client |
| 8 | + |
| 9 | +class TestRateLimitHandling(unittest.TestCase): |
| 10 | + """Test the automatic handling of RATE_LIMIT_HIT errors.""" |
| 11 | + |
| 12 | + def setUp(self): |
| 13 | + # Create a test client with a mocked get_server_settings method |
| 14 | + with patch.object(Client, 'get_server_settings', return_value={"zulip_version": "1.0", "zulip_feature_level": 1}): |
| 15 | + self.client = Client( |
| 16 | + email="test@example.com", |
| 17 | + api_key="test_api_key", |
| 18 | + site="https://example.com", |
| 19 | + ) |
| 20 | + # Make sure we have a session |
| 21 | + self.client.ensure_session() |
| 22 | + |
| 23 | + @responses.activate |
| 24 | + def test_rate_limit_retry_with_header(self): |
| 25 | + """Test that the client retries after a rate limit error with Retry-After header.""" |
| 26 | + |
| 27 | + # Add a mocked response for the first request that returns a rate limit error |
| 28 | + responses.add( |
| 29 | + responses.POST, |
| 30 | + "https://example.com/api/v1/test_endpoint", |
| 31 | + json={"result": "error", "code": "RATE_LIMIT_HIT", "msg": "Rate limit hit"}, |
| 32 | + status=429, |
| 33 | + headers={"Retry-After": "1"} # 1 second retry |
| 34 | + ) |
| 35 | + |
| 36 | + # Add a mocked response for the second request (after retry) that succeeds |
| 37 | + responses.add( |
| 38 | + responses.POST, |
| 39 | + "https://example.com/api/v1/test_endpoint", |
| 40 | + json={"result": "success", "msg": ""}, |
| 41 | + status=200 |
| 42 | + ) |
| 43 | + |
| 44 | + # Mock time.sleep to avoid actually waiting during the test |
| 45 | + with patch('time.sleep') as mock_sleep: |
| 46 | + result = self.client.call_endpoint(url="test_endpoint") |
| 47 | + |
| 48 | + # Verify that sleep was called with the correct retry value |
| 49 | + mock_sleep.assert_called_once_with(1) |
| 50 | + |
| 51 | + # Verify that we got the success response |
| 52 | + self.assertEqual(result["result"], "success") |
| 53 | + |
| 54 | + # Verify that both responses were requested |
| 55 | + self.assertEqual(len(responses.calls), 2) |
| 56 | + |
| 57 | + @responses.activate |
| 58 | + def test_rate_limit_retry_without_header(self): |
| 59 | + """Test that the client retries after a rate limit error without Retry-After header.""" |
| 60 | + |
| 61 | + # Add a mocked response for the first request that returns a rate limit error |
| 62 | + responses.add( |
| 63 | + responses.POST, |
| 64 | + "https://example.com/api/v1/test_endpoint", |
| 65 | + json={"result": "error", "code": "RATE_LIMIT_HIT", "msg": "Rate limit hit"}, |
| 66 | + status=429 |
| 67 | + # No Retry-After header |
| 68 | + ) |
| 69 | + |
| 70 | + # Add a mocked response for the second request (after retry) that succeeds |
| 71 | + responses.add( |
| 72 | + responses.POST, |
| 73 | + "https://example.com/api/v1/test_endpoint", |
| 74 | + json={"result": "success", "msg": ""}, |
| 75 | + status=200 |
| 76 | + ) |
| 77 | + |
| 78 | + # Mock time.sleep to avoid actually waiting during the test |
| 79 | + with patch('time.sleep') as mock_sleep: |
| 80 | + result = self.client.call_endpoint(url="test_endpoint") |
| 81 | + |
| 82 | + # Verify that sleep was called (with any value) |
| 83 | + mock_sleep.assert_called_once() |
| 84 | + |
| 85 | + # Verify that we got the success response |
| 86 | + self.assertEqual(result["result"], "success") |
| 87 | + |
| 88 | + # Verify that both responses were requested |
| 89 | + self.assertEqual(len(responses.calls), 2) |
| 90 | + |
| 91 | +if __name__ == "__main__": |
| 92 | + unittest.main() |
0 commit comments