New Rate Limiting Mechanism (#3148)

This commit is contained in:
Bibo-Joshi 2022-08-26 06:50:03 +02:00 committed by GitHub
parent cf6c298b82
commit 741a50ab97
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
30 changed files with 3643 additions and 60 deletions

View file

@ -26,7 +26,7 @@ Setting things up
.. code-block:: bash
$ pip install -r requirements.txt -r requirements-dev.txt
$ pip install -r requirements-all.txt
5. Install pre-commit hooks:

View file

@ -22,8 +22,6 @@ jobs:
- name: Install dependencies
run: |
python -W ignore -m pip install --upgrade pip
python -W ignore -m pip install -r requirements.txt
python -W ignore -m pip install -r requirements-dev.txt
python -W ignore -m pip install -r docs/requirements-docs.txt
python -W ignore -m pip install -r requirements-all.txt
- name: Check Links
run: sphinx-build docs/source docs/build/html -W --keep-going -j auto -b linkcheck

View file

@ -27,8 +27,6 @@ jobs:
- name: Install dependencies
run: |
python -W ignore -m pip install --upgrade pip
python -W ignore -m pip install -r requirements.txt
python -W ignore -m pip install -r requirements-dev.txt
python -W ignore -m pip install -r docs/requirements-docs.txt
python -W ignore -m pip install -r requirements-all.txt
- name: Build docs
run: sphinx-build docs/source docs/build/html -W --keep-going -j auto

View file

@ -3,6 +3,7 @@ on:
pull_request_target:
paths:
- requirements.txt
- requirements-opts.txt
- .pre-commit-config.yaml
permissions:
pull-requests: write
@ -14,5 +15,5 @@ jobs:
- name: running the check
uses: Poolitzer/notifier-action@master
with:
notify-message: Hey! Looks like you edited the requirements or the pre-commit hooks. I'm just a friendly reminder to keep the additional dependencies for the hooks in sync with the requirements :)
notify-message: Hey! Looks like you edited the (optional) requirements or the pre-commit hooks. I'm just a friendly reminder to keep the additional dependencies for the hooks in sync with the requirements :)
repo-token: ${{ secrets.GITHUB_TOKEN }}

View file

@ -42,32 +42,41 @@ jobs:
python -W ignore -m pip install --upgrade pip
python -W ignore -m pip install -U codecov pytest-cov
python -W ignore -m pip install -r requirements.txt
python -W ignore -m pip install -r requirements-opts.txt
python -W ignore -m pip install -r requirements-dev.txt
- name: Test with pytest
# We run 3 different suites here
# We run 4 different suites here
# 1. Test just utils.datetime.py without pytz being installed
# 2. Test just test_no_passport.py without passport dependencies being installed
# 3. Test everything else
# 3. Test just test_rate_limiter.py without passport dependencies being installed
# 4. Test everything else
# The first & second one are achieved by mocking the corresponding import
# See test_helpers.py & test_no_passport.py for details
run: |
pytest -v --cov -k test_no_passport.py
no_passport_exit=$?
export TEST_NO_PASSPORT='false'
export TEST_PASSPORT='true'
pytest -v --cov --cov-append -k test_helpers.py
no_pytz_exit=$?
export TEST_NO_PYTZ='false'
export TEST_PYTZ='true'
pip uninstall aiolimiter -y
pytest -v --cov --cov-append -k test_ratelimiter.py
no_rate_limiter_exit=$?
export TEST_RATE_LIMITER='true'
pip install -r requirements-opts.txt
pytest -v --cov --cov-append
full_exit=$?
special_exit=$(( no_pytz_exit > no_passport_exit ? no_pytz_exit : no_passport_exit ))
special_exit=$(( special_exit > no_rate_limiter_exit ? special_exit : no_rate_limiter_exit ))
global_exit=$(( special_exit > full_exit ? special_exit : full_exit ))
exit ${global_exit}
env:
JOB_INDEX: ${{ strategy.job-index }}
BOTS: W3sidG9rZW4iOiAiNjk2MTg4NzMyOkFBR1Z3RUtmSEhsTmpzY3hFRE5LQXdraEdzdFpfa28xbUMwIiwgInBheW1lbnRfcHJvdmlkZXJfdG9rZW4iOiAiMjg0Njg1MDYzOlRFU1Q6WldGaU1UUmxNbVF5TnpNeSIsICJib3RfbmFtZSI6ICJQVEIgdGVzdHMgb24gVHJhdmlzIHVzaW5nIENQeXRob24gMi43IiwgInN1cGVyX2dyb3VwX2lkIjogIi0xMDAxMzkwOTgzOTk3IiwgImJvdF91c2VybmFtZSI6ICJAcHRiX3RyYXZpc19jcHl0aG9uXzI3X2JvdCJ9LCB7InRva2VuIjogIjY3MTQ2ODg4NjpBQUdQR2ZjaVJJQlVORmU4MjR1SVZkcTdKZTNfWW5BVE5HdyIsICJwYXltZW50X3Byb3ZpZGVyX3Rva2VuIjogIjI4NDY4NTA2MzpURVNUOlpHWXdPVGxrTXpNeE4yWTIiLCAiYm90X25hbWUiOiAiUFRCIHRlc3RzIG9uIFRyYXZpcyB1c2luZyBDUHl0aG9uIDMuNCIsICJzdXBlcl9ncm91cF9pZCI6ICItMTAwMTQ0NjAyMjUyMiIsICJib3RfdXNlcm5hbWUiOiAiQHB0Yl90cmF2aXNfY3B5dGhvbl8zNF9ib3QifSwgeyJ0b2tlbiI6ICI2MjkzMjY1Mzg6QUFGUnJaSnJCN29CM211ekdzR0pYVXZHRTVDUXpNNUNVNG8iLCAicGF5bWVudF9wcm92aWRlcl90b2tlbiI6ICIyODQ2ODUwNjM6VEVTVDpNbU01WVdKaFl6a3hNMlUxIiwgImJvdF9uYW1lIjogIlBUQiB0ZXN0cyBvbiBUcmF2aXMgdXNpbmcgQ1B5dGhvbiAzLjUiLCAic3VwZXJfZ3JvdXBfaWQiOiAiLTEwMDE0OTY5MTc3NTAiLCAiYm90X3VzZXJuYW1lIjogIkBwdGJfdHJhdmlzX2NweXRob25fMzVfYm90In0sIHsidG9rZW4iOiAiNjQwMjA4OTQzOkFBRmhCalFwOXFtM1JUeFN6VXBZekJRakNsZS1Kano1aGNrIiwgInBheW1lbnRfcHJvdmlkZXJfdG9rZW4iOiAiMjg0Njg1MDYzOlRFU1Q6WXpoa1pUZzFOamMxWXpWbCIsICJib3RfbmFtZSI6ICJQVEIgdGVzdHMgb24gVHJhdmlzIHVzaW5nIENQeXRob24gMy42IiwgInN1cGVyX2dyb3VwX2lkIjogIi0xMDAxMzMzODcxNDYxIiwgImJvdF91c2VybmFtZSI6ICJAcHRiX3RyYXZpc19jcHl0aG9uXzM2X2JvdCJ9LCB7InRva2VuIjogIjY5NTEwNDA4ODpBQUhmenlsSU9qU0lJUy1lT25JMjB5MkUyMEhvZEhzZnotMCIsICJwYXltZW50X3Byb3ZpZGVyX3Rva2VuIjogIjI4NDY4NTA2MzpURVNUOk9HUTFNRGd3WmpJd1pqRmwiLCAiYm90X25hbWUiOiAiUFRCIHRlc3RzIG9uIFRyYXZpcyB1c2luZyBDUHl0aG9uIDMuNyIsICJzdXBlcl9ncm91cF9pZCI6ICItMTAwMTQ3ODI5MzcxNCIsICJib3RfdXNlcm5hbWUiOiAiQHB0Yl90cmF2aXNfY3B5dGhvbl8zN19ib3QifSwgeyJ0b2tlbiI6ICI2OTE0MjM1NTQ6QUFGOFdrakNaYm5IcVBfaTZHaFRZaXJGRWxackdhWU9oWDAiLCAicGF5bWVudF9wcm92aWRlcl90b2tlbiI6ICIyODQ2ODUwNjM6VEVTVDpZamM1TlRoaU1tUXlNV1ZoIiwgImJvdF9uYW1lIjogIlBUQiB0ZXN0cyBvbiBUcmF2aXMgdXNpbmcgUHlQeSAyLjciLCAic3VwZXJfZ3JvdXBfaWQiOiAiLTEwMDEzNjM5MzI1NzMiLCAiYm90X3VzZXJuYW1lIjogIkBwdGJfdHJhdmlzX3B5cHlfMjdfYm90In0sIHsidG9rZW4iOiAiNjg0MzM5OTg0OkFBRk1nRUVqcDAxcjVyQjAwN3lDZFZOc2c4QWxOc2FVLWNjIiwgInBheW1lbnRfcHJvdmlkZXJfdG9rZW4iOiAiMjg0Njg1MDYzOlRFU1Q6TVRBek1UWTNNR1V5TmpnMCIsICJib3RfbmFtZSI6ICJQVEIgdGVzdHMgb24gVHJhdmlzIHVzaW5nIFB5UHkgMy41IiwgInN1cGVyX2dyb3VwX2lkIjogIi0xMDAxNDA3ODM2NjA1IiwgImJvdF91c2VybmFtZSI6ICJAcHRiX3RyYXZpc19weXB5XzM1X2JvdCJ9LCB7InRva2VuIjogIjY5MDA5MTM0NzpBQUZMbVI1cEFCNVljcGVfbU9oN3pNNEpGQk9oMHozVDBUbyIsICJwYXltZW50X3Byb3ZpZGVyX3Rva2VuIjogIjI4NDY4NTA2MzpURVNUOlpEaGxOekU1TURrd1lXSmkiLCAiYm90X25hbWUiOiAiUFRCIHRlc3RzIG9uIEFwcFZleW9yIHVzaW5nIENQeXRob24gMy40IiwgInN1cGVyX2dyb3VwX2lkIjogIi0xMDAxMjc5NjAwMDI2IiwgImJvdF91c2VybmFtZSI6ICJAcHRiX2FwcHZleW9yX2NweXRob25fMzRfYm90In0sIHsidG9rZW4iOiAiNjk0MzA4MDUyOkFBRUIyX3NvbkNrNTVMWTlCRzlBTy1IOGp4aVBTNTVvb0JBIiwgInBheW1lbnRfcHJvdmlkZXJfdG9rZW4iOiAiMjg0Njg1MDYzOlRFU1Q6WW1aaVlXWm1NakpoWkdNeSIsICJib3RfbmFtZSI6ICJQVEIgdGVzdHMgb24gQXBwVmV5b3IgdXNpbmcgQ1B5dGhvbiAyLjciLCAic3VwZXJfZ3JvdXBfaWQiOiAiLTEwMDEyOTMwNzkxNjUiLCAiYm90X3VzZXJuYW1lIjogIkBwdGJfYXBwdmV5b3JfY3B5dGhvbl8yN19ib3QifSwgeyJ0b2tlbiI6ICIxMDU1Mzk3NDcxOkFBRzE4bkJfUzJXQXd1SjNnN29oS0JWZ1hYY2VNbklPeVNjIiwgInBheW1lbnRfcHJvdmlkZXJfdG9rZW4iOiAiMjg0Njg1MDYzOlRFU1Q6TmpBd056QXpZalZpTkdOayIsICJuYW1lIjogIlBUQiB0ZXN0cyBbMF0iLCAic3VwZXJfZ3JvdXBfaWQiOiAiLTEwMDExODU1MDk2MzYiLCAidXNlcm5hbWUiOiAicHRiXzBfYm90In0sIHsidG9rZW4iOiAiMTA0NzMyNjc3MTpBQUY4bk90ODFGcFg4bGJidno4VWV3UVF2UmZUYkZmQnZ1SSIsICJwYXltZW50X3Byb3ZpZGVyX3Rva2VuIjogIjI4NDY4NTA2MzpURVNUOllUVTFOVEk0WkdSallqbGkiLCAibmFtZSI6ICJQVEIgdGVzdHMgWzFdIiwgInN1cGVyX2dyb3VwX2lkIjogIi0xMDAxNDg0Nzk3NjEyIiwgInVzZXJuYW1lIjogInB0Yl8xX2JvdCJ9LCB7InRva2VuIjogIjk3MTk5Mjc0NTpBQUdPa09hVzBOSGpnSXY1LTlqUWJPajR2R3FkaFNGLVV1cyIsICJwYXltZW50X3Byb3ZpZGVyX3Rva2VuIjogIjI4NDY4NTA2MzpURVNUOk5XWmtNV1ZoWWpsallqVTUiLCAibmFtZSI6ICJQVEIgdGVzdHMgWzJdIiwgInN1cGVyX2dyb3VwX2lkIjogIi0xMDAxNDAyMjU1MDcwIiwgInVzZXJuYW1lIjogInB0Yl8yX2JvdCJ9XQ==
TEST_NO_PYTZ : "true"
TEST_NO_PASSPORT: "true"
TEST_PYTZ : "false"
TEST_PASSPORT: "false"
TEST_RATE_LIMITER: "false"
TEST_BUILD: "true"
shell: bash --noprofile --norc {0}
@ -95,6 +104,7 @@ jobs:
run: |
python -W ignore -m pip install --upgrade pip
python -W ignore -m pip install -r requirements.txt
python -W ignore -m pip install -r requirements-opts.txt
python -W ignore -m pip install -r requirements-dev.txt
- name: Compare to official api
run: |

View file

@ -35,6 +35,7 @@ repos:
- tornado~=6.2
- APScheduler~=3.9.1
- cachetools~=5.2.0
- aiolimiter~=1.0.0
- . # this basically does `pip install -e .`
- repo: https://github.com/pre-commit/mirrors-mypy
rev: v0.971
@ -43,7 +44,6 @@ repos:
name: mypy-ptb
files: ^telegram/.*\.py$
additional_dependencies:
- types-ujson
- types-pytz
- types-cryptography
- types-cachetools
@ -51,6 +51,7 @@ repos:
- tornado~=6.2
- APScheduler~=3.9.1
- cachetools~=5.2.0
- aiolimiter~=1.0.0
- . # this basically does `pip install -e .`
- id: mypy
name: mypy-examples

View file

@ -1 +1 @@
include LICENSE LICENSE.lesser Makefile requirements.txt README_RAW.rst telegram/py.typed
include LICENSE LICENSE.lesser Makefile requirements.txt requirements-opts.txt README_RAW.rst telegram/py.typed

View file

@ -138,6 +138,7 @@ PTB can be installed with optional dependencies:
* ``pip install python-telegram-bot[passport]`` installs the `cryptography>=3.0 <https://cryptography.io/en/stable>`_ library. Use this, if you want to use Telegram Passport related functionality.
* ``pip install python-telegram-bot[socks]`` installs ``httpx[socks]``. Use this, if you want to work behind a Socks5 server.
* ``pip install python-telegram-bot[rate-limiter]`` installs ``aiolimiter~=1.0.0``. Use this, if you want to use ``telegram.ext.AIORateLimiter``.
Quick Start
===========

View file

@ -490,7 +490,7 @@ def autodoc_process_bases(app, name, obj, option, bases: list):
# Now convert `telegram._message.Message` to `telegram.Message` etc
match = re.search(pattern=r"(telegram(\.ext|))\.[_\w\.]+", string=base)
if not match or "_utils" in base:
return
continue
parts = match.group(0).split(".")

View file

@ -0,0 +1,6 @@
telegram.ext.AIORateLimiter
============================
.. autoclass:: telegram.ext.AIORateLimiter
:members:
:show-inheritance:

View file

@ -0,0 +1,6 @@
telegram.ext.BaseRateLimiter
============================
.. autoclass:: telegram.ext.BaseRateLimiter
:members:
:show-inheritance:

View file

@ -3,5 +3,4 @@ telegram.ext.ExtBot
.. autoclass:: telegram.ext.ExtBot
:show-inheritance:
.. autofunction:: telegram.ext.ExtBot.insert_callback_data
:members: insert_callback_data, defaults, rate_limiter, initialize, shutdown

View file

@ -55,3 +55,11 @@ Arbitrary Callback Data
telegram.ext.callbackdatacache
telegram.ext.invalidcallbackdata
Rate Limiting
-------------
.. toctree::
telegram.ext.baseratelimiter
telegram.ext.aioratelimiter

4
requirements-all.txt Normal file
View file

@ -0,0 +1,4 @@
-r requirements.txt
-r requirements-dev.txt
-r requirements-opts.txt
-r docs/requirements-docs.txt

View file

@ -1,6 +1,3 @@
# cryptography is an optional dependency, but running the tests properly requires it
cryptography!=3.4,!=3.4.1,!=3.4.2,!=3.4.3
pre-commit
pytest==7.1.2

7
requirements-opts.txt Normal file
View file

@ -0,0 +1,7 @@
# Format:
# package_name==version # req-1, req-2, req-3!ext
# `pip install ptb-raw[req-1/2]` will install `package_name`
# `pip install ptb[req-1/2/3]` will also install `package_name`
httpx[socks] # socks
cryptography!=3.4,!=3.4.1,!=3.4.2,!=3.4.3,>=3.0 # passport
aiolimiter~=1.0.0 # rate-limiter!ext

View file

@ -2,6 +2,7 @@
"""The setup and build script for the python-telegram-bot library."""
import subprocess
import sys
from collections import defaultdict
from pathlib import Path
from setuptools import find_packages, setup
@ -35,6 +36,25 @@ def get_packages_requirements(raw=False):
return packs, reqs
def get_optional_requirements(raw=False):
"""Build the optional dependencies"""
requirements = defaultdict(list)
with Path("requirements-opts.txt").open() as reqs:
for line in reqs:
if line.startswith("#"):
continue
dependency, names = line.split("#")
dependency = dependency.strip()
for name in names.split(","):
name = name.strip()
if name.endswith("!ext") and raw:
continue
requirements[name].append(dependency)
return requirements
def get_setup_kwargs(raw=False):
"""Builds a dictionary of kwargs for the setup function"""
packages, requirements = get_packages_requirements(raw=raw)
@ -69,11 +89,7 @@ def get_setup_kwargs(raw=False):
long_description_content_type="text/x-rst",
packages=packages,
install_requires=requirements,
extras_require={
"socks": "httpx[socks]",
# 3.4-3.4.3 contained some cyclical import bugs
"passport": "cryptography!=3.4,!=3.4.1,!=3.4.2,!=3.4.3,>=3.0",
},
extras_require=get_optional_requirements(raw=raw),
include_package_data=True,
classifiers=[
"Development Status :: 5 - Production/Stable",

View file

@ -296,6 +296,25 @@ class Bot(TelegramObject, AbstractAsyncContextManager):
# Drop any None values because Telegram doesn't handle them well
data = {key: value for key, value in data.items() if value is not None}
return await self._do_post(
endpoint=endpoint,
data=data,
read_timeout=read_timeout,
write_timeout=write_timeout,
connect_timeout=connect_timeout,
pool_timeout=pool_timeout,
)
async def _do_post(
self,
endpoint: str,
data: JSONDict,
*,
read_timeout: ODVInput[float] = DEFAULT_NONE,
write_timeout: ODVInput[float] = DEFAULT_NONE,
connect_timeout: ODVInput[float] = DEFAULT_NONE,
pool_timeout: ODVInput[float] = DEFAULT_NONE,
) -> Union[bool, JSONDict, None]:
# This also converts datetimes into timestamps.
# We don't do this earlier so that _insert_defaults (see above) has a chance to convert
# to the default timezone in case this is called by ExtBot
@ -2902,7 +2921,7 @@ class Bot(TelegramObject, AbstractAsyncContextManager):
api_kwargs=api_kwargs,
)
return UserProfilePhotos.de_json(result, self) # type: ignore[return-value, arg-type]
return UserProfilePhotos.de_json(result, self) # type: ignore[arg-type,return-value]
@_log
async def get_file(

View file

@ -19,10 +19,13 @@
"""Extensions over the Telegram Bot API to facilitate bot making"""
__all__ = (
"AIORateLimiter",
"Application",
"ApplicationBuilder",
"ApplicationHandlerStop",
"BaseHandler",
"BasePersistence",
"BaseRateLimiter",
"CallbackContext",
"CallbackDataCache",
"CallbackQueryHandler",
@ -36,7 +39,6 @@ __all__ = (
"DictPersistence",
"ExtBot",
"filters",
"BaseHandler",
"InlineQueryHandler",
"InvalidCallbackData",
"Job",
@ -56,9 +58,11 @@ __all__ = (
)
from . import filters
from ._aioratelimiter import AIORateLimiter
from ._application import Application, ApplicationHandlerStop
from ._applicationbuilder import ApplicationBuilder
from ._basepersistence import BasePersistence, PersistenceInput
from ._baseratelimiter import BaseRateLimiter
from ._callbackcontext import CallbackContext
from ._callbackdatacache import CallbackDataCache, InvalidCallbackData
from ._callbackqueryhandler import CallbackQueryHandler

View file

@ -0,0 +1,262 @@
#!/usr/bin/env python
#
# A library that provides a Python interface to the Telegram Bot API
# Copyright (C) 2015-2022
# Leandro Toledo de Souza <devs@python-telegram-bot.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser Public License for more details.
#
# You should have received a copy of the GNU Lesser Public License
# along with this program. If not, see [http://www.gnu.org/licenses/].
"""This module contains an implementation of the BaseRateLimiter class based on the aiolimiter
library.
"""
import asyncio
import contextlib
import logging
import sys
from typing import Any, AsyncIterator, Callable, Coroutine, Dict, Optional, Union
try:
from aiolimiter import AsyncLimiter
AIO_LIMITER_AVAILABLE = True
except ImportError:
AIO_LIMITER_AVAILABLE = False
from telegram._utils.types import JSONDict
from telegram.error import RetryAfter
from telegram.ext._baseratelimiter import BaseRateLimiter
# Useful for something like:
# async with group_limiter if group else null_context():
# so we don't have to differentiate between "I'm using a context manager" and "I'm not"
if sys.version_info >= (3, 10):
null_context = contextlib.nullcontext # pylint: disable=invalid-name
else:
@contextlib.asynccontextmanager
async def null_context() -> AsyncIterator[None]:
yield None
class AIORateLimiter(BaseRateLimiter[int]):
"""
Implementation of :class:`~telegram.ext.BaseRateLimiter` using the library
`aiolimiter <https://aiolimiter.readthedocs.io/>`_.
Important:
If you want to use this class, you must install PTB with the optional requirement
``rate-limiter``, i.e.
.. code-block:: bash
pip install python-telegram-bot[rate-limiter]
The rate limiting is applied by combining two levels of throttling and :meth:`process_request`
roughly boils down to::
async with group_limiter(group_id):
async with overall_limiter:
await callback(*args, **kwargs)
Here, ``group_id`` is determined by checking if there is a ``chat_id`` parameter in the
:paramref:`~telegram.ext.BaseRateLimiter.process_request.data`.
The ``overall_limiter`` is applied only if a ``chat_id`` argument is present at all.
Attention:
* Some bot methods accept a ``chat_id`` parameter in form of a ``@username`` for
supergroups and channels. As we can't know which ``@username`` corresponds to which
integer ``chat_id``, these will be treated as different groups, which may lead to
exceeding the rate limit.
* As channels can't be differentiated from supergroups by the ``@username`` or integer
``chat_id``, this also applies the group related rate limits to channels.
* A :exc:`~telegram.error.RetryAfter` exception will halt *all* requests for
:attr:`~telegram.error.RetryAfter.retry_after` + 0.1 seconds. This may be stricter than
necessary in some cases, e.g. the bot may hit a rate limit in one group but might still
be allowed to send messages in another group.
Note:
This class is to be understood as minimal effort reference implementation.
If you would like to handle rate limiting in a more sophisticated, fine-tuned way, we
welcome you to implement your own subclass of :class:`~telegram.ext.BaseRateLimiter`.
Feel free to check out the source code of this class for inspiration.
.. versionadded:: 20.0
Args:
overall_max_rate (:obj:`float`): The maximum number of requests allowed for the entire bot
per :paramref:`overall_time_period`. When set to 0, no rate limiting will be applied.
Defaults to ``30``.
overall_time_period (:obj:`float`): The time period (in seconds) during which the
:paramref:`overall_max_rate` is enforced. When set to 0, no rate limiting will be
applied. Defaults to 1.
group_max_rate (:obj:`float`): The maximum number of requests allowed for requests related
to groups and channels per :paramref:`group_time_period`. When set to 0, no rate
limiting will be applied. Defaults to 20.
group_time_period (:obj:`float`): The time period (in seconds) during which the
:paramref:`group_time_period` is enforced. When set to 0, no rate limiting will be
applied. Defaults to 60.
max_retries (:obj:`int`): The maximum number of retries to be made in case of a
:exc:`~telegram.error.RetryAfter` exception.
If set to 0, no retries will be made. Defaults to ``0``.
"""
__slots__ = (
"_base_limiter",
"_group_limiters",
"_group_max_rate",
"_group_time_period",
"_logger",
"_max_retries",
"_retry_after_event",
)
def __init__(
self,
overall_max_rate: float = 30,
overall_time_period: float = 1,
group_max_rate: float = 20,
group_time_period: float = 60,
max_retries: int = 0,
) -> None:
if not AIO_LIMITER_AVAILABLE:
raise RuntimeError(
"To use `AIORateLimiter`, PTB must be installed via `pip install "
"python-telegram-bot[rate-limiter]`."
)
if overall_max_rate and overall_time_period:
self._base_limiter: Optional[AsyncLimiter] = AsyncLimiter(
max_rate=overall_max_rate, time_period=overall_time_period
)
else:
self._base_limiter = None
if group_max_rate and group_time_period:
self._group_max_rate = group_max_rate
self._group_time_period = group_time_period
else:
self._group_max_rate = 0
self._group_time_period = 0
self._group_limiters: Dict[Union[str, int], AsyncLimiter] = {}
self._max_retries = max_retries
self._logger = logging.getLogger(__name__)
self._retry_after_event = asyncio.Event()
self._retry_after_event.set()
async def initialize(self) -> None:
"""Does nothing."""
async def shutdown(self) -> None:
"""Does nothing."""
def _get_group_limiter(self, group_id: Union[str, int, bool]) -> "AsyncLimiter":
# Remove limiters that haven't been used for so long that all their capacity is unused
# We only do that if we have a lot of limiters lying around to avoid looping on every call
# This is a minimal effort approach - a full-fledged cache could use a TTL approach
# or at least adapt the threshold dynamically depending on the number of active limiters
if len(self._group_limiters) > 512:
# We copy to avoid modifying the dict while we iterate over it
for key, limiter in self._group_limiters.copy().items():
if key == group_id:
continue
if limiter.has_capacity(limiter.max_rate):
del self._group_limiters[key]
if group_id not in self._group_limiters:
self._group_limiters[group_id] = AsyncLimiter(
max_rate=self._group_max_rate,
time_period=self._group_time_period,
)
return self._group_limiters[group_id]
async def _run_request(
self,
chat: bool,
group: Union[str, int, bool],
callback: Callable[..., Coroutine[Any, Any, Union[bool, JSONDict, None]]],
args: Any,
kwargs: Dict[str, Any],
) -> Union[bool, JSONDict, None]:
base_context = self._base_limiter if (chat and self._base_limiter) else null_context()
group_context = (
self._get_group_limiter(group) if group and self._group_max_rate else null_context()
)
async with group_context: # skipcq: PTC-W0062
async with base_context:
# In case a retry_after was hit, we wait with processing the request
await self._retry_after_event.wait()
return await callback(*args, **kwargs)
# mypy doesn't understand that the last run of the for loop raises an exception
async def process_request( # type: ignore[return]
self,
callback: Callable[..., Coroutine[Any, Any, Union[bool, JSONDict, None]]],
args: Any,
kwargs: Dict[str, Any],
endpoint: str, # skipcq: PYL-W0613
data: Dict[str, Any],
rate_limit_args: Optional[int],
) -> Union[bool, JSONDict, None]:
"""
Processes a request by applying rate limiting.
See :meth:`telegram.ext.BaseRateLimiter.process_request` for detailed information on the
arguments.
Args:
rate_limit_args (:obj:`None` | :obj:`int`): If set, specifies the maximum number of
retries to be made in case of a :exc:`~telegram.error.RetryAfter` exception.
Defaults to :paramref:`AIORateLimiter.max_retries`.
"""
max_retries = rate_limit_args or self._max_retries
group: Union[int, str, bool] = False
chat: bool = False
chat_id = data.get("chat_id")
if chat_id is not None:
chat = True
# In case user passes integer chat id as string
try:
chat_id = int(chat_id)
except (ValueError, TypeError):
pass
if (isinstance(chat_id, int) and chat_id < 0) or isinstance(chat_id, str):
# string chat_id only works for channels and supergroups
# We can't really tell channels from groups though ...
group = chat_id
for i in range(max_retries + 1):
try:
return await self._run_request(
chat=chat, group=group, callback=callback, args=args, kwargs=kwargs
)
except RetryAfter as exc:
if i == max_retries:
self._logger.exception(
"Rate limit hit after maximum of %d retries", max_retries, exc_info=exc
)
raise exc
sleep = exc.retry_after + 0.1
self._logger.info("Rate limit hit. Retrying after %f seconds", sleep)
# Make sure we don't allow other requests to be processed
self._retry_after_event.clear()
await asyncio.sleep(sleep)
finally:
# Allow other requests to be processed
self._retry_after_event.set()

View file

@ -45,7 +45,8 @@ from telegram.request import BaseRequest
from telegram.request._httpxrequest import HTTPXRequest
if TYPE_CHECKING:
from telegram.ext import BasePersistence, CallbackContext, Defaults
from telegram.ext import BasePersistence, BaseRateLimiter, CallbackContext, Defaults
from telegram.ext._utils.types import RLARGS
# Type hinting is a bit complicated here because we try to get to a sane level of
# leveraging generics and therefore need a number of type variables.
@ -81,6 +82,7 @@ _BOT_CHECKS = [
("defaults", "defaults"),
("arbitrary_callback_data", "arbitrary_callback_data"),
("private_key", "private_key"),
("rate_limiter", "rate_limiter instance"),
]
_TWO_ARGS_REQ = "The parameter `{}` may only be set, if no {} was set."
@ -139,6 +141,7 @@ class ApplicationBuilder(Generic[BT, CCT, UD, CD, BD, JQ]):
"_private_key",
"_private_key_password",
"_proxy_url",
"_rate_limiter",
"_read_timeout",
"_request",
"_token",
@ -180,6 +183,7 @@ class ApplicationBuilder(Generic[BT, CCT, UD, CD, BD, JQ]):
self._updater: ODVInput[Updater] = DEFAULT_NONE
self._post_init: Optional[Callable[[Application], Coroutine[Any, Any, None]]] = None
self._post_shutdown: Optional[Callable[[Application], Coroutine[Any, Any, None]]] = None
self._rate_limiter: ODVInput["BaseRateLimiter"] = DEFAULT_NONE
def _build_request(self, get_updates: bool) -> BaseRequest:
prefix = "_get_updates_" if get_updates else "_"
@ -227,6 +231,7 @@ class ApplicationBuilder(Generic[BT, CCT, UD, CD, BD, JQ]):
arbitrary_callback_data=DefaultValue.get_value(self._arbitrary_callback_data),
request=self._build_request(get_updates=False),
get_updates_request=self._build_request(get_updates=True),
rate_limiter=DefaultValue.get_value(self._rate_limiter),
)
def build(
@ -973,10 +978,31 @@ class ApplicationBuilder(Generic[BT, CCT, UD, CD, BD, JQ]):
self._post_shutdown = post_shutdown
return self
def rate_limiter(
self: "ApplicationBuilder[BT, CCT, UD, CD, BD, JQ]",
rate_limiter: "BaseRateLimiter[RLARGS]",
) -> "ApplicationBuilder[ExtBot[RLARGS], CCT, UD, CD, BD, JQ]":
"""Sets a :class:`telegram.ext.BaseRateLimiter` instance for the
:paramref:`telegram.ext.ExtBot.rate_limiter` parameter of
:attr:`telegram.ext.Application.bot`.
Args:
rate_limiter (:class:`telegram.ext.BaseRateLimiter`): The rate limiter.
Returns:
:class:`ApplicationBuilder`: The same builder with the updated argument.
"""
if self._bot is not DEFAULT_NONE:
raise RuntimeError(_TWO_ARGS_REQ.format("rate_limiter", "bot instance"))
if self._updater not in (DEFAULT_NONE, None):
raise RuntimeError(_TWO_ARGS_REQ.format("rate_limiter", "updater"))
self._rate_limiter = rate_limiter
return self # type: ignore[return-value]
InitApplicationBuilder = ( # This is defined all the way down here so that its type is inferred
ApplicationBuilder[ # by Pylance correctly.
ExtBot,
ExtBot[None],
ContextTypes.DEFAULT_TYPE,
Dict,
Dict,

View file

@ -0,0 +1,138 @@
#!/usr/bin/env python
#
# A library that provides a Python interface to the Telegram Bot API
# Copyright (C) 2015-2022
# Leandro Toledo de Souza <devs@python-telegram-bot.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser Public License for more details.
#
# You should have received a copy of the GNU Lesser Public License
# along with this program. If not, see [http://www.gnu.org/licenses/].
"""This module contains a class that allows to rate limit requests to the Bot API."""
from abc import ABC, abstractmethod
from typing import Any, Callable, Coroutine, Dict, Generic, Optional, Union
from telegram._utils.types import JSONDict
from telegram.ext._utils.types import RLARGS
class BaseRateLimiter(ABC, Generic[RLARGS]):
"""
Abstract interface class that allows to rate limit the requests that python-telegram-bot
sends to the Telegram Bot API. An implementation of this class
must implement all abstract methods and properties.
This class is a :class:`~typing.Generic` class and accepts one type variable that specifies
the type of the argument :paramref:`~process_request.rate_limit_args` of
:meth:`process_request` and the methods of :class:`~telegram.ext.ExtBot`.
Hint:
Requests to :meth:`~telegram.Bot.get_updates` are never rate limited.
.. versionadded:: 20.0
"""
__slots__ = ()
@abstractmethod
async def initialize(self) -> None:
"""Initialize resources used by this class. Must be implemented by a subclass."""
@abstractmethod
async def shutdown(self) -> None:
"""Stop & clear resources used by this class. Must be implemented by a subclass."""
@abstractmethod
async def process_request(
self,
callback: Callable[..., Coroutine[Any, Any, Union[bool, JSONDict, None]]],
args: Any,
kwargs: Dict[str, Any],
endpoint: str,
data: Dict[str, Any],
rate_limit_args: Optional[RLARGS],
) -> Union[bool, JSONDict, None]:
"""
Process a request. Must be implemented by a subclass.
This method must call :paramref:`callback` and return the result of the call.
`When` the callback is called is up to the implementation.
Important:
This method must only return once the result of :paramref:`callback` is known!
If a :exc:`~telegram.error.RetryAfter` error is raised, this method may try to make
a new request by calling the callback again.
Warning:
This method *should not* handle any other exception raised by :paramref:`callback`!
There are basically two different approaches how a rate limiter can be implemented:
1. React only if necessary. In this case, the :paramref:`callback` is called without any
precautions. If a :exc:`~telegram.error.RetryAfter` error is raised, processing requests
is halted for the :attr:`~telegram.error.RetryAfter.retry_after` and finally the
:paramref:`callback` is called again. This approach is often amendable for bots that
don't have a large user base and/or don't send more messages than they get updates.
2. Throttle all outgoing requests. In this case the implementation makes sure that the
requests are spread out over a longer time interval in order to stay below the rate
limits. This approach is often amendable for bots that have a large user base and/or
send more messages than they get updates.
An implementation can use the information provided by :paramref:`data`,
:paramref:`endpoint` and :paramref:`rate_limit_args` to handle each request differently.
Examples:
* It is usually desirable to call :meth:`telegram.Bot.answer_inline_query`
as quickly as possible, while delaying :meth:`telegram.Bot.send_message`
is acceptable.
* There are `different <https://core.telegram.org/bots/faq\
#my-bot-is-hitting-limits-how-do-i-avoid-this>`_ rate limits for group chats and
private chats.
* When sending broadcast messages to a large number of users, these requests can
typically be delayed for a longer time than messages that are direct replies to a
user input.
Args:
callback (Callable[..., :term:`coroutine`]): The coroutine function that must be called
to make the request.
args (Tuple[:obj:`object`]): The positional arguments for the :paramref:`callback`
function.
kwargs (Dict[:obj:`str`, :obj:`object`]): The keyword arguments for the
:paramref:`callback` function.
endpoint (:obj:`str`): The endpoint that the request is made for, e.g.
``"sendMessage"``.
data (Dict[:obj:`str`, :obj:`object`]): The parameters that were passed to the method
of :class:`~telegram.ext.ExtBot`. Any ``api_kwargs`` are included in this and
any :paramref:`~telegram.ext.ExtBot.defaults` are already applied.
Example:
When calling::
await ext_bot.send_message(
chat_id=1,
text="Hello world!",
api_kwargs={"custom": "arg"}
)
then :paramref:`data` will be::
{"chat_id": 1, "text": "Hello world!", "custom": "arg"}
rate_limit_args (:obj:`None` | :class:`object`): Custom arguments passed to the methods
of :class:`~telegram.ext.ExtBot`. Can e.g. be used to specify the priority of
the request.
Returns:
:obj:`bool` | Dict[:obj:`str`, :obj:`object`] | :obj:`None`: The result of the
callback function.
"""

File diff suppressed because it is too large Load diff

View file

@ -39,8 +39,10 @@ from typing import (
)
if TYPE_CHECKING:
from typing import Optional
from telegram import Bot
from telegram.ext import CallbackContext, JobQueue
from telegram.ext import BaseRateLimiter, CallbackContext, JobQueue
CCT = TypeVar("CCT", bound="CallbackContext")
"""An instance of :class:`telegram.ext.CallbackContext` or a custom subclass.
@ -101,3 +103,13 @@ JQ = TypeVar("JQ", bound=Union[None, "JobQueue"])
"""Type of the job queue.
.. versionadded:: 20.0"""
RL = TypeVar("RL", bound="Optional[BaseRateLimiter]")
"""Type of the rate limiter.
.. versionadded:: 20.0"""
RLARGS = TypeVar("RLARGS")
"""Type of the rate limiter arguments.
.. versionadded:: 20.0"""

View file

@ -558,6 +558,7 @@ async def check_shortcut_call(
bot: The bot
bot_method_name: The bot methods name, e.g. `'send_message'`
skip_params: Parameters that are allowed to be missing, e.g. `['inline_message_id']`
`rate_limit_args` will be skipped by default
shortcut_kwargs: The kwargs passed by the shortcut directly, e.g. ``chat_id``
Returns:
@ -565,8 +566,13 @@ async def check_shortcut_call(
"""
if not skip_params:
skip_params = set()
else:
skip_params = set(skip_params)
skip_params.add("rate_limit_args")
if not shortcut_kwargs:
shortcut_kwargs = set()
else:
shortcut_kwargs = set(shortcut_kwargs)
orig_bot_method = getattr(bot, bot_method_name)
bot_signature = inspect.signature(orig_bot_method)

View file

@ -23,6 +23,7 @@ import httpx
import pytest
from telegram.ext import (
AIORateLimiter,
Application,
ApplicationBuilder,
ContextTypes,
@ -82,6 +83,7 @@ class TestApplicationBuilder:
assert app.bot.private_key is None
assert app.bot.arbitrary_callback_data is False
assert app.bot.defaults is None
assert app.bot.rate_limiter is None
get_updates_client = app.bot._request[0]._client
assert get_updates_client.limits == httpx.Limits(
@ -196,6 +198,7 @@ class TestApplicationBuilder:
"proxy_url",
"bot",
"update_queue",
"rate_limiter",
]
+ [entry[0] for entry in _BOT_CHECKS],
)
@ -247,10 +250,13 @@ class TestApplicationBuilder:
defaults = Defaults()
request = HTTPXRequest()
get_updates_request = HTTPXRequest()
rate_limiter = AIORateLimiter()
builder.token(bot.token).base_url("base_url").base_file_url("base_file_url").private_key(
PRIVATE_KEY
).defaults(defaults).arbitrary_callback_data(42).request(request).get_updates_request(
get_updates_request
).rate_limiter(
rate_limiter
)
built_bot = builder.build().bot
@ -266,6 +272,7 @@ class TestApplicationBuilder:
assert built_bot._request[0] is get_updates_request
assert built_bot.callback_data_cache.maxsize == 42
assert built_bot.private_key
assert built_bot.rate_limiter is rate_limiter
@dataclass
class Client:

View file

@ -481,9 +481,9 @@ class TestBot:
corresponding methods of tg.Bot.
"""
# Some methods of ext.ExtBot
global_extra_args = set()
global_extra_args = {"rate_limit_args"}
extra_args_per_method = defaultdict(
set, {"__init__": {"arbitrary_callback_data", "defaults"}}
set, {"__init__": {"arbitrary_callback_data", "defaults", "rate_limiter"}}
)
different_hints_per_method = defaultdict(set, {"__setattr__": {"ext_bot"}})
@ -2948,3 +2948,8 @@ class TestBot:
assert (
"api_kwargs" in param_names
), f"{bot_method_name} is missing the parameter `api_kwargs`"
if bot_class is ExtBot and bot_method_name.replace("_", "").lower() != "getupdates":
assert (
"rate_limit_args" in param_names
), f"{bot_method_name} of ExtBot is missing the parameter `rate_limit_args`"

View file

@ -48,16 +48,16 @@ Because imports in pytest are intricate, we just run
pytest -k test_helpers.py
with the TEST_NO_PYTZ environment variable set in addition to the regular test suite.
with the TEST_PYTZ environment variable set to False in addition to the regular test suite.
Because actually uninstalling pytz would lead to errors in the test suite we just mock the
import to raise the expected exception.
Note that a fixture that just does this for every test that needs it is a nice idea, but for some
reason makes test_updater.py hang indefinitely on GitHub Actions (at least when Hinrich tried that)
"""
TEST_NO_PYTZ = env_var_2_bool(os.getenv("TEST_NO_PYTZ", False))
TEST_PYTZ = env_var_2_bool(os.getenv("TEST_PYTZ", True))
if TEST_NO_PYTZ:
if not TEST_PYTZ:
orig_import = __import__
def import_mock(module_name, *args, **kwargs):
@ -72,7 +72,7 @@ if TEST_NO_PYTZ:
class TestDatetime:
def test_helpers_utc(self):
# Here we just test, that we got the correct UTC variant
if TEST_NO_PYTZ:
if not TEST_PYTZ:
assert tg_dtm.UTC is tg_dtm.DTM_UTC
else:
assert tg_dtm.UTC is not tg_dtm.DTM_UTC

View file

@ -24,7 +24,7 @@ Currently this only means that cryptography is not installed.
Because imports in pytest are intricate, we just run
pytest -k test_no_passport.py
with the TEST_NO_PASSPORT environment variable set in addition to the regular test suite.
with the TEST_PASSPORT environment variable set to False in addition to the regular test suite.
Because actually uninstalling the optional dependencies would lead to errors in the test suite we
just mock the import to raise the expected exception.
@ -41,9 +41,9 @@ from telegram import _bot as bot
from telegram._passport import credentials as credentials
from tests.conftest import env_var_2_bool
TEST_NO_PASSPORT = env_var_2_bool(os.getenv("TEST_NO_PASSPORT", False))
TEST_PASSPORT = env_var_2_bool(os.getenv("TEST_PASSPORT", True))
if TEST_NO_PASSPORT:
if not TEST_PASSPORT:
orig_import = __import__
def import_mock(module_name, *args, **kwargs):
@ -58,24 +58,24 @@ if TEST_NO_PASSPORT:
class TestNoPassport:
"""
The monkeypatches simulate cryptography not being installed even when TEST_NO_PASSPORT is
False, though that doesn't test the actual imports
The monkeypatches simulate cryptography not being installed even when TEST_PASSPORT is
True, though that doesn't test the actual imports
"""
def test_bot_init(self, bot_info, monkeypatch):
if not TEST_NO_PASSPORT:
if TEST_PASSPORT:
monkeypatch.setattr(bot, "CRYPTO_INSTALLED", False)
with pytest.raises(RuntimeError, match="passport"):
bot.Bot(bot_info["token"], private_key=1, private_key_password=2)
def test_credentials_decrypt(self, monkeypatch):
if not TEST_NO_PASSPORT:
if TEST_PASSPORT:
monkeypatch.setattr(credentials, "CRYPTO_INSTALLED", False)
with pytest.raises(RuntimeError, match="passport"):
credentials.decrypt(1, 1, 1)
def test_encrypted_credentials_decrypted_secret(self, monkeypatch):
if not TEST_NO_PASSPORT:
if TEST_PASSPORT:
monkeypatch.setattr(credentials, "CRYPTO_INSTALLED", False)
ec = credentials.EncryptedCredentials("data", "hash", "secret")
with pytest.raises(RuntimeError, match="passport"):

364
tests/test_ratelimiter.py Normal file
View file

@ -0,0 +1,364 @@
#!/usr/bin/env python
#
# A library that provides a Python interface to the Telegram Bot API
# Copyright (C) 2015-2022
# Leandro Toledo de Souza <devs@python-telegram-bot.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser Public License for more details.
#
# You should have received a copy of the GNU Lesser Public License
# along with this program. If not, see [http://www.gnu.org/licenses/].
"""
We mostly test on directly on AIORateLimiter here, b/c BaseRateLimiter doesn't contain anything
notable
"""
import asyncio
import json
import os
import platform
import time
from datetime import datetime
from http import HTTPStatus
import pytest
from flaky import flaky
from telegram import BotCommand, Chat, Message, User
from telegram.constants import ParseMode
from telegram.error import RetryAfter
from telegram.ext import AIORateLimiter, BaseRateLimiter, Defaults, ExtBot
from telegram.request import BaseRequest, RequestData
from tests.conftest import env_var_2_bool
TEST_RATE_LIMITER = env_var_2_bool(os.getenv("TEST_RATE_LIMITER", True))
@pytest.mark.skipif(
TEST_RATE_LIMITER, reason="Only relevant if the optional dependency is not installed"
)
class TestNoRateLimiter:
def test_init(self):
with pytest.raises(RuntimeError, match=r"python-telegram-bot\[rate-limiter\]"):
AIORateLimiter()
class TestBaseRateLimiter:
rl_received = None
request_received = None
async def test_no_rate_limiter(self, bot):
with pytest.raises(ValueError, match="if a `ExtBot.rate_limiter` is set"):
await bot.send_message(chat_id=42, text="test", rate_limit_args="something")
async def test_argument_passing(self, bot_info, monkeypatch, bot):
class TestRateLimiter(BaseRateLimiter):
async def initialize(self) -> None:
pass
async def shutdown(self) -> None:
pass
async def process_request(
self,
callback,
args,
kwargs,
endpoint,
data,
rate_limit_args,
):
if TestBaseRateLimiter.rl_received is None:
TestBaseRateLimiter.rl_received = []
TestBaseRateLimiter.rl_received.append((endpoint, data, rate_limit_args))
return await callback(*args, **kwargs)
class TestRequest(BaseRequest):
async def initialize(self) -> None:
pass
async def shutdown(self) -> None:
pass
async def do_request(self, *args, **kwargs):
if TestBaseRateLimiter.request_received is None:
TestBaseRateLimiter.request_received = []
TestBaseRateLimiter.request_received.append((args, kwargs))
# return bot.bot.to_dict() for the `get_me` call in `Bot.initialize`
return 200, json.dumps({"ok": True, "result": bot.bot.to_dict()}).encode()
defaults = Defaults(parse_mode=ParseMode.HTML)
test_request = TestRequest()
standard_bot = ExtBot(token=bot.token, defaults=defaults, request=test_request)
rl_bot = ExtBot(
token=bot.token,
defaults=defaults,
request=test_request,
rate_limiter=TestRateLimiter(),
)
async with standard_bot:
await standard_bot.set_my_commands(
commands=[BotCommand("test", "test")],
language_code="en",
api_kwargs={"api": "kwargs"},
)
async with rl_bot:
await rl_bot.set_my_commands(
commands=[BotCommand("test", "test")],
language_code="en",
rate_limit_args=(43, "test-1"),
api_kwargs={"api": "kwargs"},
)
assert len(self.rl_received) == 2
assert self.rl_received[0] == ("getMe", {}, None)
assert self.rl_received[1] == (
"setMyCommands",
dict(commands=[BotCommand("test", "test")], language_code="en", api="kwargs"),
(43, "test-1"),
)
assert len(self.request_received) == 4
# self.request_received[i] = i-th received request
# self.request_received[i][0] = i-th received request's args
# self.request_received[i][1] = i-th received request's kwargs
assert self.request_received[0][1]["url"].endswith("getMe")
assert self.request_received[2][1]["url"].endswith("getMe")
assert self.request_received[1][0] == self.request_received[3][0]
assert self.request_received[1][1].keys() == self.request_received[3][1].keys()
for key, value in self.request_received[1][1].items():
if isinstance(value, RequestData):
assert value.parameters == self.request_received[3][1][key].parameters
assert value.parameters["api"] == "kwargs"
else:
assert value == self.request_received[3][1][key]
@pytest.mark.skipif(
not TEST_RATE_LIMITER, reason="Only relevant if the optional dependency is installed"
)
@pytest.mark.skipif(
os.getenv("GITHUB_ACTIONS", False) and platform.system() == "Darwin",
reason="The timings are apparently rather inaccurate on MacOS.",
)
@flaky(10, 1) # Timings aren't quite perfect
class TestAIORateLimiter:
count = 0
call_times = []
class CountRequest(BaseRequest):
def __init__(self, retry_after=None):
self.retry_after = retry_after
async def initialize(self) -> None:
pass
async def shutdown(self) -> None:
pass
async def do_request(self, *args, **kwargs):
TestAIORateLimiter.count += 1
TestAIORateLimiter.call_times.append(time.time())
if self.retry_after:
raise RetryAfter(retry_after=1)
url = kwargs.get("url").lower()
if url.endswith("getme"):
return (
HTTPStatus.OK,
json.dumps(
{"ok": True, "result": User(id=1, first_name="bot", is_bot=True).to_dict()}
).encode(),
)
if url.endswith("sendmessage"):
return (
HTTPStatus.OK,
json.dumps(
{
"ok": True,
"result": Message(
message_id=1, date=datetime.now(), chat=Chat(1, "chat")
).to_dict(),
}
).encode(),
)
@pytest.fixture(autouse=True)
def reset(self):
self.count = 0
TestAIORateLimiter.count = 0
self.call_times = []
TestAIORateLimiter.call_times = []
@pytest.mark.parametrize("max_retries", [0, 1, 4])
async def test_max_retries(self, bot, max_retries):
bot = ExtBot(
token=bot.token,
request=self.CountRequest(retry_after=1),
rate_limiter=AIORateLimiter(
max_retries=max_retries, overall_max_rate=0, group_max_rate=0
),
)
with pytest.raises(RetryAfter):
await bot.get_me()
# Check that we retried the request the correct number of times
assert TestAIORateLimiter.count == max_retries + 1
# Check that the retries were delayed correctly
times = TestAIORateLimiter.call_times
if len(times) <= 1:
return
delays = [j - i for i, j in zip(times[:-1], times[1:])]
assert delays == pytest.approx([1.1 for _ in range(max_retries)], rel=0.05)
async def test_delay_all_pending_on_retry(self, bot):
# Makes sure that a RetryAfter blocks *all* pending requests
bot = ExtBot(
token=bot.token,
request=self.CountRequest(retry_after=1),
rate_limiter=AIORateLimiter(max_retries=1, overall_max_rate=0, group_max_rate=0),
)
task_1 = asyncio.create_task(bot.get_me())
await asyncio.sleep(0.1)
task_2 = asyncio.create_task(bot.get_me())
assert not task_1.done()
assert not task_2.done()
await asyncio.sleep(1.1)
assert isinstance(task_1.exception(), RetryAfter)
assert not task_2.done()
await asyncio.sleep(1.1)
assert isinstance(task_2.exception(), RetryAfter)
@pytest.mark.parametrize("group_id", [-1, "-1", "@username"])
@pytest.mark.parametrize("chat_id", [1, "1"])
async def test_basic_rate_limiting(self, bot, group_id, chat_id):
try:
rl_bot = ExtBot(
token=bot.token,
request=self.CountRequest(retry_after=None),
rate_limiter=AIORateLimiter(
overall_max_rate=1,
overall_time_period=1 / 4,
group_max_rate=1,
group_time_period=1 / 2,
),
)
async with rl_bot:
non_group_tasks = {}
group_tasks = {}
for i in range(4):
group_tasks[i] = asyncio.create_task(
rl_bot.send_message(chat_id=group_id, text="test")
)
for i in range(8):
non_group_tasks[i] = asyncio.create_task(
rl_bot.send_message(chat_id=chat_id, text="test")
)
await asyncio.sleep(0.85)
# We expect 5 requests:
# 1: `get_me` from `async with rl_bot`
# 2: `send_message` at time 0.00
# 3: `send_message` at time 0.25
# 4: `send_message` at time 0.50
# 5: `send_message` at time 0.75
assert TestAIORateLimiter.count == 5
assert sum(1 for task in non_group_tasks.values() if task.done()) < 8
assert sum(1 for task in group_tasks.values() if task.done()) < 4
# 3 seconds after start
await asyncio.sleep(3.1 - 0.85)
assert all(task.done() for task in non_group_tasks.values())
assert all(task.done() for task in group_tasks.values())
finally:
# cleanup
await asyncio.gather(*non_group_tasks.values(), *group_tasks.values())
TestAIORateLimiter.count = 0
TestAIORateLimiter.call_times = []
async def test_rate_limiting_no_chat_id(self, bot):
try:
rl_bot = ExtBot(
token=bot.token,
request=self.CountRequest(retry_after=None),
rate_limiter=AIORateLimiter(
overall_max_rate=1,
overall_time_period=1 / 2,
),
)
async with rl_bot:
non_chat_tasks = {}
chat_tasks = {}
for i in range(4):
chat_tasks[i] = asyncio.create_task(
rl_bot.send_message(chat_id=-1, text="test")
)
for i in range(8):
non_chat_tasks[i] = asyncio.create_task(rl_bot.get_me())
await asyncio.sleep(0.6)
# We expect 11 requests:
# 1: `get_me` from `async with rl_bot`
# 2: `send_message` at time 0.00
# 3: `send_message` at time 0.05
# 4: 8 times `get_me`
assert TestAIORateLimiter.count == 11
assert sum(1 for task in non_chat_tasks.values() if task.done()) == 8
assert sum(1 for task in chat_tasks.values() if task.done()) == 2
# 1.6 seconds after start
await asyncio.sleep(1.6 - 0.6)
assert all(task.done() for task in non_chat_tasks.values())
assert all(task.done() for task in chat_tasks.values())
finally:
# cleanup
await asyncio.gather(*non_chat_tasks.values(), *chat_tasks.values())
TestAIORateLimiter.count = 0
TestAIORateLimiter.call_times = []
@pytest.mark.parametrize("intermediate", [True, False])
async def test_group_caching(self, bot, intermediate):
try:
max_rate = 1000
rl_bot = ExtBot(
token=bot.token,
request=self.CountRequest(retry_after=None),
rate_limiter=AIORateLimiter(
overall_max_rate=max_rate,
overall_time_period=1,
group_max_rate=max_rate,
group_time_period=1,
),
)
# Unfortunately, there is no reliable way to test this without checking the internals
assert len(rl_bot.rate_limiter._group_limiters) == 0
await asyncio.gather(
*(rl_bot.send_message(chat_id=-(i + 1), text=f"{i}") for i in range(513))
)
if intermediate:
await rl_bot.send_message(chat_id=-1, text="999")
assert 1 <= len(rl_bot.rate_limiter._group_limiters) <= 513
else:
await asyncio.sleep(1)
await rl_bot.send_message(chat_id=-1, text="999")
assert len(rl_bot.rate_limiter._group_limiters) == 1
finally:
TestAIORateLimiter.count = 0
TestAIORateLimiter.call_times = []