diff --git a/aioetherscan/client.py b/aioetherscan/client.py index c5ac3f2..66ce64f 100644 --- a/aioetherscan/client.py +++ b/aioetherscan/client.py @@ -1,5 +1,5 @@ from asyncio import AbstractEventLoop -from typing import AsyncContextManager +from typing import AsyncContextManager, Union from aiohttp import ClientTimeout from aiohttp_retry import RetryOptionsBase @@ -20,7 +20,7 @@ class Client: def __init__( self, - api_key: str, + api_key: Union[str, list[str]], api_kind: str = 'eth', network: str = 'main', loop: AbstractEventLoop = None, @@ -29,7 +29,8 @@ def __init__( throttler: AsyncContextManager = None, retry_options: RetryOptionsBase = None, ) -> None: - self._url_builder = UrlBuilder(api_key, api_kind, network) + api_keys = [api_key] if isinstance(api_key, str) else api_key + self._url_builder = UrlBuilder(api_keys, api_kind, network) self._http = Network(self._url_builder, loop, timeout, proxy, throttler, retry_options) self.account = Account(self) diff --git a/aioetherscan/exceptions.py b/aioetherscan/exceptions.py index c764fd3..a015578 100644 --- a/aioetherscan/exceptions.py +++ b/aioetherscan/exceptions.py @@ -20,6 +20,10 @@ def __str__(self): return f'[{self.message}] {self.result}' +class EtherscanClientApiRateLimitError(EtherscanClientApiError): + pass + + class EtherscanClientProxyError(EtherscanClientError): """JSON-RPC 2.0 Specification diff --git a/aioetherscan/modules/extra/generators/generator_utils.py b/aioetherscan/modules/extra/generators/generator_utils.py index f6056ae..0f764f2 100644 --- a/aioetherscan/modules/extra/generators/generator_utils.py +++ b/aioetherscan/modules/extra/generators/generator_utils.py @@ -84,7 +84,7 @@ async def _parse_by_pages( api_method: Callable, request_params: dict[str, Any] ) -> AsyncIterator[Transfer]: page = count(1) - while True: + while True: # pragma: no cover request_params['page'] = next(page) try: result = await api_method(**request_params) diff --git a/aioetherscan/network.py b/aioetherscan/network.py index dffae21..a0f27d4 100644 --- a/aioetherscan/network.py +++ b/aioetherscan/network.py @@ -1,6 +1,7 @@ import asyncio import logging from asyncio import AbstractEventLoop +from functools import wraps from typing import Union, AsyncContextManager, Optional import aiohttp @@ -15,10 +16,30 @@ EtherscanClientError, EtherscanClientApiError, EtherscanClientProxyError, + EtherscanClientApiRateLimitError, ) from aioetherscan.url_builder import UrlBuilder +def retry_limit_attempt(f): + @wraps(f) + async def inner(self, *args, **kwargs): + attempt = 1 + max_attempts = self._url_builder.keys_count + while True: # pragma: no cover + try: + return await f(self, *args, **kwargs) + except EtherscanClientApiRateLimitError as e: + self._logger.warning(f'Key daily limit exceeded, {attempt=}: {e}') + attempt += 1 + if attempt > max_attempts: + raise e + await asyncio.sleep(0.01) + self._url_builder.rotate_api_key() + + return inner + + class Network: def __init__( self, @@ -48,9 +69,11 @@ async def close(self): if self._retry_client is not None: await self._retry_client.close() + @retry_limit_attempt async def get(self, params: dict = None) -> Union[dict, list, str]: return await self._request(METH_GET, params=self._url_builder.filter_and_sign(params)) + @retry_limit_attempt async def post(self, data: dict = None) -> Union[dict, list, str]: return await self._request(METH_POST, data=self._url_builder.filter_and_sign(data)) @@ -68,6 +91,7 @@ async def _request( if self._retry_client is None: self._retry_client = self._get_retry_client() session_method = getattr(self._retry_client, method.lower()) + async with self._throttler: async with session_method( self._url_builder.API_URL, params=params, data=data, proxy=self._proxy @@ -93,6 +117,10 @@ async def _handle_response(self, response: aiohttp.ClientResponse) -> Union[dict def _raise_if_error(response_json: dict): if 'status' in response_json and response_json['status'] != '1': message, result = response_json.get('message'), response_json.get('result') + + if 'max daily rate limit reached' in result.lower(): + raise EtherscanClientApiRateLimitError(message, result) + raise EtherscanClientApiError(message, result) if 'error' in response_json: diff --git a/aioetherscan/url_builder.py b/aioetherscan/url_builder.py index 6e29db7..9054cff 100644 --- a/aioetherscan/url_builder.py +++ b/aioetherscan/url_builder.py @@ -1,3 +1,5 @@ +import logging +from itertools import cycle from typing import Optional from urllib.parse import urlunsplit, urljoin @@ -19,8 +21,10 @@ class UrlBuilder: BASE_URL: str = None API_URL: str = None - def __init__(self, api_key: str, api_kind: str, network: str) -> None: - self._API_KEY = api_key + def __init__(self, api_keys: list[str], api_kind: str, network: str) -> None: + self._api_keys = api_keys + self._api_keys_cycle = cycle(self._api_keys) + self._api_key = self._get_next_api_key() self._set_api_kind(api_kind) self._network = network.lower().strip() @@ -28,6 +32,8 @@ def __init__(self, api_key: str, api_kind: str, network: str) -> None: self.API_URL = self._get_api_url() self.BASE_URL = self._get_base_url() + self._logger = logging.getLogger(__name__) + def _set_api_kind(self, api_kind: str) -> None: api_kind = api_kind.lower().strip() if api_kind not in self._API_KINDS: @@ -87,9 +93,35 @@ def filter_and_sign(self, params: dict): def _sign(self, params: dict) -> dict: if not params: params = {} - params['apikey'] = self._API_KEY + params['apikey'] = self._api_key return params @staticmethod def _filter_params(params: dict) -> dict: return {k: v for k, v in params.items() if v is not None} + + def _get_next_api_key(self) -> str: + return next(self._api_keys_cycle) + + def rotate_api_key(self) -> None: + prev_api_key = self._api_key + next_api_key = self._get_next_api_key() + + self._logger.info( + f'Rotating API key from {self._mask_api_key(prev_api_key)!r} to {self._mask_api_key(next_api_key)!r}' + ) + + self._api_key = next_api_key + + @staticmethod + def _mask_api_key(api_key: str, masked_chars_count: int = 30, symbol: str = '*') -> str: + api_key_len = len(api_key) + if masked_chars_count >= api_key_len or masked_chars_count <= 0: + return symbol * len(api_key) + + right_part = api_key[-(api_key_len - masked_chars_count) :] + return right_part.rjust(len(api_key), symbol) + + @property + def keys_count(self) -> int: + return len(self._api_keys) diff --git a/tests/test_network.py b/tests/test_network.py index 1aef1e8..43ff21c 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -1,7 +1,7 @@ import asyncio import json import logging -from unittest.mock import patch, AsyncMock, MagicMock, Mock +from unittest.mock import patch, AsyncMock, MagicMock, Mock, call import aiohttp import aiohttp_retry @@ -17,8 +17,9 @@ EtherscanClientError, EtherscanClientApiError, EtherscanClientProxyError, + EtherscanClientApiRateLimitError, ) -from aioetherscan.network import Network +from aioetherscan.network import Network, retry_limit_attempt from aioetherscan.url_builder import UrlBuilder @@ -45,7 +46,7 @@ def get_loop(): @pytest_asyncio.fixture async def ub(): - ub = UrlBuilder('test_api_key', 'eth', 'main') + ub = UrlBuilder(['test_api_key'], 'eth', 'main') yield ub @@ -86,25 +87,25 @@ def test_no_loop(ub): async def test_get(nw): with patch('aioetherscan.network.Network._request', new=AsyncMock()) as mock: await nw.get() - mock.assert_called_once_with(METH_GET, params={'apikey': nw._url_builder._API_KEY}) + mock.assert_called_once_with(METH_GET, params={'apikey': nw._url_builder._api_key}) @pytest.mark.asyncio async def test_post(nw): with patch('aioetherscan.network.Network._request', new=AsyncMock()) as mock: await nw.post() - mock.assert_called_once_with(METH_POST, data={'apikey': nw._url_builder._API_KEY}) + mock.assert_called_once_with(METH_POST, data={'apikey': nw._url_builder._api_key}) with patch('aioetherscan.network.Network._request', new=AsyncMock()) as mock: await nw.post({'some': 'data'}) mock.assert_called_once_with( - METH_POST, data={'apikey': nw._url_builder._API_KEY, 'some': 'data'} + METH_POST, data={'apikey': nw._url_builder._api_key, 'some': 'data'} ) with patch('aioetherscan.network.Network._request', new=AsyncMock()) as mock: await nw.post({'some': 'data', 'null': None}) mock.assert_called_once_with( - METH_POST, data={'apikey': nw._url_builder._API_KEY, 'some': 'data'} + METH_POST, data={'apikey': nw._url_builder._api_key, 'some': 'data'} ) @@ -238,3 +239,77 @@ def test_get_retry_client(nw): retry_options=nw._retry_options, ) assert result is m.return_value + + +def test_raise_if_error_daily_limit_reached(nw): + data = dict( + status='0', + message='NOTOK', + result='Max daily rate limit reached. 110000 (100%) of 100000 day/limit', + ) + with pytest.raises(EtherscanClientApiRateLimitError) as e: + nw._raise_if_error(data) + + assert e.value.message == data['message'] + assert e.value.result == data['result'] + + +class TestRetryClass: + def __init__(self, limit: int, keys_count: int): + self._url_builder = Mock() + self._url_builder.keys_count = keys_count + self._url_builder.rotate_api_key = Mock() + + self._logger = Mock() + self._logger.warning = Mock() + + self._count = 1 + self._limit = limit + + @retry_limit_attempt + async def some_method(self): + self._count += 1 + + if self._count > self._limit: + raise EtherscanClientApiRateLimitError( + 'NOTOK', + 'Max daily rate limit reached. 110000 (100%) of 100000 day/limit', + ) + + +@pytest.mark.asyncio +async def test_retry_limit_attempt_error_limit_exceeded(nw): + c = TestRetryClass(1, 1) + + with pytest.raises(EtherscanClientApiRateLimitError): + await c.some_method() + c._url_builder.rotate_api_key.assert_not_called() + c._logger.warning.assert_called_once_with( + 'Key daily limit exceeded, attempt=1: [NOTOK] Max daily rate limit reached. 110000 (100%) of 100000 day/limit' + ) + + +@pytest.mark.asyncio +async def test_retry_limit_attempt_error_limit_rotate(nw): + c = TestRetryClass(1, 2) + + with pytest.raises(EtherscanClientApiRateLimitError): + await c.some_method() + c._url_builder.rotate_api_key.assert_called_once() + c._logger.warning.assert_has_calls( + [ + call( + 'Key daily limit exceeded, attempt=1: [NOTOK] Max daily rate limit reached. 110000 (100%) of 100000 day/limit' + ), + call( + 'Key daily limit exceeded, attempt=2: [NOTOK] Max daily rate limit reached. 110000 (100%) of 100000 day/limit' + ), + ] + ) + + +@pytest.mark.asyncio +async def test_retry_limit_attempt_ok(nw): + c = TestRetryClass(2, 1) + + await c.some_method() diff --git a/tests/test_url_builder.py b/tests/test_url_builder.py index 5323016..d4dcd6a 100644 --- a/tests/test_url_builder.py +++ b/tests/test_url_builder.py @@ -1,4 +1,5 @@ -from unittest.mock import patch +from itertools import cycle +from unittest.mock import patch, Mock import pytest import pytest_asyncio @@ -7,7 +8,7 @@ def apikey(): - return 'test_api_key' + return ['test_api_key'] @pytest_asyncio.fixture @@ -17,8 +18,8 @@ async def ub(): def test_sign(ub): - assert ub._sign({}) == {'apikey': ub._API_KEY} - assert ub._sign({'something': 'something'}) == {'something': 'something', 'apikey': ub._API_KEY} + assert ub._sign({}) == {'apikey': ub._api_key} + assert ub._sign({'something': 'something'}) == {'something': 'something', 'apikey': ub._api_key} def test_filter_params(ub): @@ -130,3 +131,42 @@ def test_get_link(ub): path = 'some_path' ub.get_link(path) join_mock.assert_called_once_with(ub.BASE_URL, path) + + +@pytest.mark.parametrize( + 'api_key,masked_chars_count,expected', + [ + ('abc', 4, '***'), + ('abc', 2, '**c'), + ('abcd', 4, '****'), + ('abcdef', 4, '****ef'), + ], +) +def test_mask_api_key(ub, api_key, masked_chars_count, expected): + assert ub._mask_api_key(api_key, masked_chars_count) == expected + + assert ub._mask_api_key('qwe', 2, '#') == '##e' + + +def test_keys_count(ub): + ub._api_keys = [1, 2] + assert ub.keys_count == 2 + + +def test_rotate_api_key(ub): + api_keys = ['one', 'two', 'three'] + first, second, _ = api_keys + + ub._api_keys = api_keys + ub._api_keys_cycle = cycle(ub._api_keys) + ub._api_key = ub._get_next_api_key() + + assert ub._api_key == first + + ub._logger = Mock() + ub._logger.info = Mock() + + ub.rotate_api_key() + + ub._logger.info.assert_called_once() + assert ub._api_key == second