diff --git a/zulip/README.md b/zulip/README.md index 75845bfee..59edc5b66 100644 --- a/zulip/README.md +++ b/zulip/README.md @@ -94,6 +94,15 @@ keys: msg, result. For successful calls, result will be "success" and msg will be the empty string. On error, result will be "error" and msg will describe what went wrong. +#### Rate Limiting + +The Zulip API client automatically handles rate limiting errors (`RATE_LIMIT_HIT`). When a rate limit error is encountered: + +1. If the server provides a `Retry-After` header, the client will pause for the specified number of seconds and then retry the request. +2. If no `Retry-After` header is provided, the client will use an exponential backoff strategy to retry the request. + +This automatic handling ensures that your application doesn't need to implement its own rate limit handling logic. + #### Examples The API bindings package comes with several nice example scripts that diff --git a/zulip/tests/test_rate_limit_handling.py b/zulip/tests/test_rate_limit_handling.py new file mode 100644 index 000000000..d28e5514e --- /dev/null +++ b/zulip/tests/test_rate_limit_handling.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 + +import unittest +import time +import responses +from unittest.mock import patch, MagicMock +from zulip import Client + +class TestRateLimitHandling(unittest.TestCase): + """Test the automatic handling of RATE_LIMIT_HIT errors.""" + + def setUp(self): + # Create a test client with a mocked get_server_settings method + with patch.object(Client, 'get_server_settings', return_value={"zulip_version": "1.0", "zulip_feature_level": 1}): + self.client = Client( + email="test@example.com", + api_key="test_api_key", + site="https://example.com", + ) + # Make sure we have a session + self.client.ensure_session() + + @responses.activate + def test_rate_limit_retry_with_header(self): + """Test that the client retries after a rate limit error with Retry-After header.""" + + # Add a mocked response for the first request that returns a rate limit error + responses.add( + responses.POST, + "https://example.com/api/v1/test_endpoint", + json={"result": "error", "code": "RATE_LIMIT_HIT", "msg": "Rate limit hit"}, + status=429, + headers={"Retry-After": "1"} # 1 second retry + ) + + # Add a mocked response for the second request (after retry) that succeeds + responses.add( + responses.POST, + "https://example.com/api/v1/test_endpoint", + json={"result": "success", "msg": ""}, + status=200 + ) + + # Mock time.sleep to avoid actually waiting during the test + with patch('time.sleep') as mock_sleep: + result = self.client.call_endpoint(url="test_endpoint") + + # Verify that sleep was called with the correct retry value + mock_sleep.assert_called_once_with(1) + + # Verify that we got the success response + self.assertEqual(result["result"], "success") + + # Verify that both responses were requested + self.assertEqual(len(responses.calls), 2) + + @responses.activate + def test_rate_limit_retry_without_header(self): + """Test that the client retries after a rate limit error without Retry-After header.""" + + # Add a mocked response for the first request that returns a rate limit error + responses.add( + responses.POST, + "https://example.com/api/v1/test_endpoint", + json={"result": "error", "code": "RATE_LIMIT_HIT", "msg": "Rate limit hit"}, + status=429 + # No Retry-After header + ) + + # Add a mocked response for the second request (after retry) that succeeds + responses.add( + responses.POST, + "https://example.com/api/v1/test_endpoint", + json={"result": "success", "msg": ""}, + status=200 + ) + + # Mock time.sleep to avoid actually waiting during the test + with patch('time.sleep') as mock_sleep: + result = self.client.call_endpoint(url="test_endpoint") + + # Verify that sleep was called (with any value) + mock_sleep.assert_called_once() + + # Verify that we got the success response + self.assertEqual(result["result"], "success") + + # Verify that both responses were requested + self.assertEqual(len(responses.calls), 2) + +if __name__ == "__main__": + unittest.main() diff --git a/zulip/zulip/__init__.py b/zulip/zulip/__init__.py index 8d930e7e7..12c11a8a8 100644 --- a/zulip/zulip/__init__.py +++ b/zulip/zulip/__init__.py @@ -687,6 +687,26 @@ def end_error_retry(succeeded: bool) -> None: "status_code": res.status_code, } + # Handle rate limiting automatically + if json_result.get("result") == "error" and json_result.get("code") == "RATE_LIMIT_HIT": + retry_after = None + # Check for Retry-After header (in seconds) + if "Retry-After" in res.headers: + try: + retry_after = int(res.headers["Retry-After"]) + except (ValueError, TypeError): + pass + + # If we have a valid retry_after value, sleep and retry + if retry_after and retry_after > 0: + if self.verbose: + print(f"Rate limit hit. Retrying after {retry_after} seconds...") + time.sleep(retry_after) + continue + # If no valid retry_after header, use a default backoff + elif error_retry(" (rate limited)"): + continue + end_error_retry(True) return json_result