Merge pull request #22 from 29next/20-improve-rate-limit-hanling

20 improve rate limit hanling

ntk/gateway.py

@@ -14,7 +14,10 @@ def _request(self, request_type, url, apikey=None, payload={}, files={}):
         if apikey:
             headers = {'Authorization': f'Bearer {apikey}'}
 
-        return requests.request(request_type, url, headers=headers, data=payload, files=files)
+        response = requests.request(request_type, url, headers=headers, data=payload, files=files)
+        if response.status_code == 429 and "throttled" in response.content.decode():
+            return self._request(request_type, url, apikey, payload, files)
+        return response
 
     @check_error(error_format='Missing Themes in {store}')
     def get_themes(self):

setup.py

@@ -1,6 +1,6 @@
 from setuptools import find_packages, setup
 
-__version__ = '1.0.3'
+__version__ = '1.0.4'
 
 tests_require = [
     "flake8==3.9.2",

tests/test_command.py

@@ -345,11 +345,10 @@ def test_watch_command_should_call_gateway_with_correct_arguments_belong_to_file
                 theme_id=1234, template_name='layout/base.html')
             self.assertIn(expected_call_deleted, self.mock_gateway.mock_calls)
 
-    @patch("ntk.command.time", autospec=True)
     @patch("ntk.command.Command._get_accept_files", autospec=True)
     @patch("builtins.open", autospec=True)
     def test_watch_command_with_create_image_file_should_call_gateway_with_correct_arguments(
-        self, mock_open_file, mock_get_accept_file, mock_time
+        self, mock_open_file, mock_get_accept_file
     ):
         mock_get_accept_file.return_value = [
             f'{os.getcwd()}/assets/image.jpg',
@@ -372,8 +371,6 @@ def test_watch_command_with_create_image_file_should_call_gateway_with_correct_a
         )
         self.assertIn(expected_call_added, self.mock_gateway.mock_calls)
 
-        mock_time.sleep.assert_called_once_with(0.07)
-
     @patch("ntk.command.Command._get_accept_files", autospec=True)
     @patch("ntk.command.Command._compile_sass", autospec=True)
     def test_watch_command_with_sass_directory_should_call_compile_sass(

tests/test_gateway.py

@@ -17,6 +17,43 @@ def setUp(self):
     #####
     @patch('ntk.gateway.requests.request', autospec=True)
     def test_request(self, mock_request):
+        mock_response_200 = MagicMock()
+        mock_response_200.status_code = 200
+
+        mock_request.return_value = mock_response_200
+
+        request_type = 'POST'
+        url = 'http://simple.com/api/admin/themes/5/templates/'
+        payload = {
+            'name': 'assets/base.html',
+            'content': '{% load i18n %}\n\n<div class="mt-2">My home page</div>'
+        }
+        files = {'file': ('assets/image.jpg', self.mock_img_file)}
+
+        self.gateway._request(request_type, url, apikey=self.apikey, payload=payload, files=files)
+
+        expected_calls = [
+            call(
+                'POST', 'http://simple.com/api/admin/themes/5/templates/',
+                headers={'Authorization': 'Bearer apikey'},
+                data={
+                    'name': 'assets/base.html', 'content': '{% load i18n %}\n\n<div class="mt-2">My home page</div>'},
+                files=files)
+        ]
+        assert mock_request.mock_calls == expected_calls
+
+    @patch('ntk.gateway.requests.request', autospec=True)
+    def test_request_with_rate_limit_should_retry(self, mock_request):
+        mock_response_429 = MagicMock()
+        mock_response_429.status_code = 429
+        mock_response_429.content.decode.return_value = "throttled"
+
+        # Mock the response for the second call with status code 200
+        mock_response_200 = MagicMock()
+        mock_response_200.status_code = 200
+
+        mock_request.side_effect = [mock_response_429, mock_response_200]
+
         request_type = 'POST'
         url = 'http://simple.com/api/admin/themes/5/templates/'
         payload = {
@@ -27,13 +64,23 @@ def test_request(self, mock_request):
 
         self.gateway._request(request_type, url, apikey=self.apikey, payload=payload, files=files)
 
-        expected_calls = [call(
-            'POST', 'http://simple.com/api/admin/themes/5/templates/',
-            headers={'Authorization': 'Bearer apikey'},
-            data={'name': 'assets/base.html', 'content': '{% load i18n %}\n\n<div class="mt-2">My home page</div>'},
-            files=files)
+        assert mock_request.call_count == 2
+
+        expected_calls = [
+            call(
+                'POST', 'http://simple.com/api/admin/themes/5/templates/',
+                headers={'Authorization': 'Bearer apikey'},
+                data={
+                    'name': 'assets/base.html', 'content': '{% load i18n %}\n\n<div class="mt-2">My home page</div>'
+                }, files=files),
+            call(
+                'POST', 'http://simple.com/api/admin/themes/5/templates/',
+                headers={'Authorization': 'Bearer apikey'},
+                data={
+                    'name': 'assets/base.html', 'content': '{% load i18n %}\n\n<div class="mt-2">My home page</div>'
+                }, files=files)
         ]
-        self.assertEqual(mock_request.mock_calls, expected_calls)
+        assert mock_request.mock_calls == expected_calls
 
     #####
     # get_themes