def test_rest_assistant_call_with_pre_and_post_processing(self, mocked_api):
url = "https://www.test.com/url"
resp = {"one": 1}
pre_processor_ran = False
post_processor_ran = False
mocked_api.get(url, body=json.dumps(resp).encode())
class PreProcessor(RESTPreProcessorBase):
async def pre_process(self, request: RESTRequest) -> RESTRequest:
nonlocal pre_processor_ran
pre_processor_ran = True
return request
class PostProcessor(RESTPostProcessorBase):
async def post_process(self, response: RESTResponse) -> RESTResponse:
nonlocal post_processor_ran
post_processor_ran = True
return response
pre_processors = [PreProcessor()]
post_processors = [PostProcessor()]
connection = RESTConnection(aiohttp.ClientSession())
assistant = RESTAssistant(
connection=connection,
throttler=AsyncThrottler(rate_limits=[]),
rest_pre_processors=pre_processors,
rest_post_processors=post_processors)
req = RESTRequest(method=RESTMethod.GET, url=url)
ret = self.async_run_with_timeout(assistant.call(req))
ret_json = self.async_run_with_timeout(ret.json())
self.assertEqual(resp, ret_json)
self.assertTrue(pre_processor_ran)
self.assertTrue(post_processor_ran)
def test_rest_assistant_authenticates(self, mocked_call):
url = "https://www.test.com/url"
resp = {"one": 1}
call_request: Optional[RESTRequest] = None
auth_header = {"authenticated": True}
async def register_request_and_return(request: RESTRequest):
nonlocal call_request
call_request = request
return resp
mocked_call.side_effect = register_request_and_return
class AuthDummy(AuthBase):
async def rest_authenticate(self,
request: RESTRequest) -> RESTRequest:
request.headers = auth_header
return request
async def ws_authenticate(self, request: WSRequest) -> WSRequest:
pass
connection = RESTConnection(aiohttp.ClientSession())
assistant = RESTAssistant(connection, auth=AuthDummy())
req = RESTRequest(method=RESTMethod.GET, url=url)
auth_req = RESTRequest(method=RESTMethod.GET,
url=url,
is_auth_required=True)
self.async_run_with_timeout(assistant.call(req))
self.assertIsNotNone(call_request)
self.assertIsNone(call_request.headers)
self.async_run_with_timeout(assistant.call(auth_req))
self.assertIsNotNone(call_request)
self.assertIsNotNone(call_request.headers)
self.assertEqual(call_request.headers, auth_header)
async def api_call_with_retries(request: GateIORESTRequest,
rest_assistant: RESTAssistant,
throttler: AsyncThrottler,
logger: logging.Logger,
gate_io_auth: Optional[GateIoAuth] = None,
try_count: int = 0) -> Dict[str, Any]:
headers = {"Content-Type": "application/json"}
async with throttler.execute_task(limit_id=request.throttler_limit_id):
if request.is_auth_required:
if gate_io_auth is None:
raise RuntimeError(
f"Authentication required for request, but no GateIoAuth object supplied."
f" Request: {request}.")
auth_params = request.data if request.method == RESTMethod.POST else request.params
request.data = auth_params
headers: dict = gate_io_auth.get_headers(str(request.method),
request.auth_url,
auth_params)
request.headers = headers
response_coro = asyncio.wait_for(rest_assistant.call(request),
CONSTANTS.API_CALL_TIMEOUT)
http_status, parsed_response, request_errors = await rest_response_with_errors(
response_coro)
if request_errors or parsed_response is None:
if try_count < CONSTANTS.API_MAX_RETRIES:
try_count += 1
time_sleep = retry_sleep_time(try_count)
logger.info(
f"Error fetching data from {request.url}. HTTP status is {http_status}."
f" Retrying in {time_sleep:.0f}s.")
await _sleep(time_sleep)
return await api_call_with_retries(request, rest_assistant,
throttler, logger, gate_io_auth,
try_count)
else:
raise GateIoAPIError({
"label": "HTTP_ERROR",
"message": parsed_response,
"status": http_status
})
if "message" in parsed_response:
raise GateIoAPIError(parsed_response)
return parsed_response