Adds requestion options on wait method

Author: nunomaduro

algoliasearch/responses.py

@@ -1,6 +1,6 @@
 import abc
 
-from typing import List, Union
+from typing import List, Union, Optional
 
 from algoliasearch.exceptions import RequestException
 from algoliasearch.helpers import get_items, sleep_for
@@ -21,8 +21,8 @@ class Response(object):
     __metaclass__ = abc.ABCMeta
 
     @abc.abstractmethod
-    def wait(self):
-        # type:() -> Response
+    def wait(self, request_options=None):
+        # type: (Optional[Union[RequestOptions, dict]]) -> Response
 
         pass  # pragma: no cover
 
@@ -36,12 +36,15 @@ def __init__(self, index, raw_responses):
         self.raw_responses = raw_responses
         self.waited = False
 
-    def wait(self):
-        # type: () -> IndexingResponse
+    def wait(self, request_options=None):
+        # type: (Optional[Union[RequestOptions, dict]]) -> IndexingResponse
 
         if not self.waited:
             for raw_response in self.raw_responses:
-                self._index._sync().wait_task(raw_response['taskID'])
+                self._index._sync().wait_task(
+                    raw_response['taskID'],
+                    request_options
+                )
 
         # No longer waits on this responses.
         self.waited = True
@@ -69,11 +72,11 @@ def push(self, response):
 
         self._waitable.append(response)
 
-    def wait(self):
-        # type: () -> MultipleResponse
+    def wait(self, request_options=None):
+        # type: (Optional[Union[RequestOptions, dict]])-> MultipleResponse
 
         for response in self._waitable:
-            response.wait()
+            response.wait(request_options)
 
         # No longer waits on this responses.
         self._waitable = []
@@ -95,14 +98,17 @@ def __init__(self, client, raw_response):
         self._client = client
         self._done = False
 
-    def wait(self):
-        # type: () -> AddApiKeyResponse
+    def wait(self, request_options=None):
+        # type: (Optional[Union[RequestOptions, dict]]) -> AddApiKeyResponse
 
         retries_count = 1
 
         while not self._done:
             try:
-                self._client._sync().get_api_key(self.raw_response['key'])
+                self._client._sync().get_api_key(
+                    self.raw_response['key'],
+                    request_options
+                )
                 self._done = True
             except RequestException as e:
                 if e.status_code != 404:
@@ -131,15 +137,15 @@ def __init__(self, client, raw_response, request_options):
         self._request_options = request_options
         self._done = False
 
-    def wait(self):
-        # type: () -> UpdateApiKeyResponse
+    def wait(self, request_options=None):
+        # type: (Optional[Union[RequestOptions, dict]]) -> UpdateApiKeyResponse
 
         retries_count = 1
 
         while not self._done:
             api_key = self._client._sync().get_api_key(
                 self.raw_response['key'],
-                self._request_options
+                request_options
             )
 
             self._done = self._have_changed(api_key)
@@ -187,14 +193,17 @@ def __init__(self, client, raw_response, key):
         self._key = key
         self._done = False
 
-    def wait(self):
-        # type: () -> DeleteApiKeyResponse
+    def wait(self, request_options=None):
+        # type: (Optional[Union[RequestOptions, dict]]) -> DeleteApiKeyResponse
 
         retries_count = 1
 
         while not self._done:
             try:
-                self._client._sync().get_api_key(self._key)
+                self._client._sync().get_api_key(
+                    self._key,
+                    request_options
+                )
             except RequestException as e:
                 self._done = e.status_code == 404
 
@@ -223,14 +232,17 @@ def __init__(self, client, raw_response, key):
         self._key = key
         self._done = False
 
-    def wait(self):
-        # type: () -> RestoreApiKeyResponse
+    def wait(self, request_options=None):
+        # type: (Optional[Union[RequestOptions, dict]]) -> RestoreApiKeyResponse # noqa: E501
 
         retries_count = 1
 
         while not self._done:
             try:
-                self._client._sync().get_api_key(self._key)
+                self._client._sync().get_api_key(
+                    self._key,
+                    request_options
+                )
                 self._done = True
             except RequestException as e:
                 if e.status_code != 404:
@@ -258,12 +270,16 @@ def __init__(self, client, raw_response):
         self._client = client
         self._done = False
 
-    def wait(self):
-        # type: () -> MultipleIndexBatchIndexingResponse
+    def wait(self, request_options=None):
+        # type: (Optional[Union[RequestOptions, dict]]) -> MultipleIndexBatchIndexingResponse # noqa: E501
 
         while not self._done:
             for index_name, task_id in get_items(self.raw_response['taskID']):
-                self._client._sync().wait_task(index_name, task_id)
+                self._client._sync().wait_task(
+                    index_name,
+                    task_id,
+                    request_options
+                )
             self._done = True
 
         return self

tests/unit/test_responses.py

@@ -2,6 +2,7 @@
 
 import mock
 
+from algoliasearch.configs import SearchConfig
 from algoliasearch.exceptions import RequestException
 from algoliasearch.http.request_options import RequestOptions
 from algoliasearch.responses import (
@@ -80,3 +81,13 @@ def test_update_api_key_response(self):
         self.assertFalse(response._have_changed({}))
         self.assertFalse(response._have_changed({'acl': []}))
         self.assertTrue(response._have_changed({'acl': ['search']}))
+
+    def test_uses_request_options_on_wait(self):
+        index = SearchClient.create('foo', 'bar').init_index('foo')
+        index.wait_task = mock.Mock(name='wait_task')
+        index._sync = mock.Mock(name='_sync')
+        index._sync.return_value = index
+
+        response = IndexingResponse(index, [{'taskID': 1}])
+        response.wait({'bar': 2})
+        index.wait_task.assert_called_once_with(1, {'bar': 2})