├── .github └── workflows │ └── test.yml ├── .gitignore ├── CODE_OF_CONDUCT.rst ├── CONTRIBUTING.rst ├── LICENSE.txt ├── MANIFEST.in ├── README.rst ├── pytest.ini ├── ratelimit ├── __init__.py ├── decorators.py └── exception.py ├── requirements.txt ├── setup.py └── tests ├── __init__.py └── unit ├── __init__.py ├── decorator_test.py └── threading_test.py /.github/workflows/test.yml: -------------------------------------------------------------------------------- 1 | name: Test 2 | 3 | on: [push, pull_request] 4 | 5 | jobs: 6 | build: 7 | 8 | runs-on: ubuntu-latest 9 | strategy: 10 | fail-fast: false 11 | matrix: 12 | python-version: ["3.6", "3.10"] 13 | env: 14 | CC_TEST_REPORTER_ID: 607cbaa7bcd1b5511e4bbfe9d85ffcc257c2c458fdcc36534b39cfd1fddf29d6 15 | 16 | steps: 17 | - uses: actions/checkout@v2 18 | - name: Set up Python ${{ matrix.python-version }} 19 | uses: actions/setup-python@v2 20 | with: 21 | python-version: ${{ matrix.python-version }} 22 | - name: Install dependencies 23 | run: | 24 | python -m pip install --upgrade pip 25 | pip install -r requirements.txt 26 | curl -L https://codeclimate.com/downloads/test-reporter/test-reporter-latest-linux-amd64 > ./cc-test-reporter 27 | chmod +x ./cc-test-reporter 28 | - name: Lint with black 29 | run: | 30 | black --check ratelimit tests 31 | - name: Test with pytest 32 | run: | 33 | ./cc-test-reporter before-build 34 | pytest 35 | ./cc-test-reporter after-build -t coverage.py --exit-code $? 36 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | 58 | # Flask stuff: 59 | instance/ 60 | .webassets-cache 61 | 62 | # Scrapy stuff: 63 | .scrapy 64 | 65 | # Sphinx documentation 66 | docs/_build/ 67 | 68 | # PyBuilder 69 | target/ 70 | 71 | # Jupyter Notebook 72 | .ipynb_checkpoints 73 | 74 | # pyenv 75 | .python-version 76 | 77 | # celery beat schedule file 78 | celerybeat-schedule 79 | 80 | # SageMath parsed files 81 | *.sage.py 82 | 83 | # Environments 84 | .env 85 | .venv 86 | env/ 87 | venv/ 88 | ENV/ 89 | env.bak/ 90 | venv.bak/ 91 | 92 | # Spyder project settings 93 | .spyderproject 94 | .spyproject 95 | 96 | # Rope project settings 97 | .ropeproject 98 | 99 | # mkdocs documentation 100 | /site 101 | 102 | # mypy 103 | .mypy_cache/ 104 | -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.rst: -------------------------------------------------------------------------------- 1 | Contributor Covenant Code of Conduct 2 | ==================================== 3 | 4 | Our Pledge 5 | ---------- 6 | 7 | In the interest of fostering an open and welcoming environment, we as 8 | contributors and maintainers pledge to making participation in our project and 9 | our community a harassment-free experience for everyone, regardless of age, 10 | body size, disability, ethnicity, gender identity and expression, level of 11 | experience, nationality, personal appearance, race, religion, or sexual 12 | identity and orientation. 13 | 14 | Our Standards 15 | ------------- 16 | 17 | Examples of behavior that contributes to creating a positive environment 18 | include: 19 | 20 | * Using welcoming and inclusive language 21 | * Being respectful of differing viewpoints and experiences 22 | * Gracefully accepting constructive criticism 23 | * Focusing on what is best for the community 24 | * Showing empathy towards other community members 25 | 26 | Examples of unacceptable behavior by participants include: 27 | 28 | * The use of sexualized language or imagery and unwelcome sexual attention or 29 | advances 30 | * Trolling, insulting/derogatory comments, and personal or political attacks 31 | * Public or private harassment 32 | * Publishing others' private information, such as a physical or electronic 33 | address, without explicit permission 34 | * Other conduct which could reasonably be considered inappropriate in a 35 | professional setting 36 | 37 | Our Responsibilities 38 | -------------------- 39 | 40 | Project maintainers are responsible for clarifying the standards of acceptable 41 | behavior and are expected to take appropriate and fair corrective action in 42 | response to any instances of unacceptable behavior. 43 | 44 | Project maintainers have the right and responsibility to remove, edit, or 45 | reject comments, commits, code, wiki edits, issues, and other contributions 46 | that are not aligned to this Code of Conduct, or to ban temporarily or 47 | permanently any contributor for other behaviors that they deem inappropriate, 48 | threatening, offensive, or harmful. 49 | 50 | Scope 51 | ----- 52 | 53 | This Code of Conduct applies both within project spaces and in public spaces 54 | when an individual is representing the project or its community. Examples of 55 | representing a project or community include using an official project e-mail 56 | address, posting via an official social media account, or acting as an 57 | appointed representative at an online or offline event. Representation of a 58 | project may be further defined and clarified by project maintainers. 59 | 60 | Enforcement 61 | ----------- 62 | 63 | Instances of abusive, harassing, or otherwise unacceptable behavior may be 64 | reported by contacting the project team at jared@shademaps.com. All 65 | complaints will be reviewed and investigated and will result in a response that 66 | is deemed necessary and appropriate to the circumstances. The project team is 67 | obligated to maintain confidentiality with regard to the reporter of an 68 | incident. Further details of specific enforcement policies may be posted 69 | separately. 70 | 71 | Project maintainers who do not follow or enforce the Code of Conduct in good 72 | faith may face temporary or permanent repercussions as determined by other 73 | members of the project's leadership. 74 | 75 | Attribution 76 | ----------- 77 | 78 | This Code of Conduct is adapted from the `Contributor Covenant`_, version 1.4, 79 | available at `http://contributor-covenant.org/version/1/4`_ 80 | 81 | .. _Contributor Covenant: http://contributor-covenant.org 82 | .. _`http://contributor-covenant.org/version/1/4`: http://contributor-covenant.org/version/1/4/ 83 | -------------------------------------------------------------------------------- /CONTRIBUTING.rst: -------------------------------------------------------------------------------- 1 | How To Contribute 2 | ================= 3 | 4 | Installation 5 | ------------ 6 | 7 | * ``git clone `` this repository 8 | * ``cd ratelimit`` 9 | * ``pip install -r requirements.txt`` 10 | 11 | Linting 12 | ------- 13 | 14 | * ``black ratelimit`` 15 | * ``black tests`` 16 | 17 | Running tests 18 | ------------- 19 | 20 | * ``py.test`` – Runs the pytest test suite 21 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2019 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy of 6 | this software and associated documentation files (the "Software"), to deal in 7 | the Software without restriction, including without limitation the rights to 8 | use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies 9 | of the Software, and to permit persons to whom the Software is furnished to do 10 | so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include CODE_OF_CONDUCT.rst 2 | include CONTRIBUTING.rst 3 | include LICENSE.txt 4 | include README.rst 5 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | ratelimit |build| |maintainability| |coverage| 2 | ============================================== 3 | 4 | This project is a fork of `tomasbasham/ratelimit `_ 5 | that implements a `sliding log `_ 6 | for correctness and provides persistance via sqlite. See the usage section on 7 | `Persistence <#persistence>`_ for more details. Turning on persistence is highly 8 | recommended, especially during development, to ensure rate limits are respected 9 | between application restarts. 10 | 11 | APIs are a very common way to interact with web services. As the need to 12 | consume data grows, so does the number of API calls necessary to remain up to 13 | date with data sources. However many API providers constrain developers from 14 | making too many API calls. This is know as rate limiting and in a worst case 15 | scenario your application can be banned from making further API calls if it 16 | abuses these limits. 17 | 18 | This packages introduces a function decorator preventing a function from being 19 | called more often than that allowed by the API provider. This should prevent 20 | API providers from banning your applications by conforming to their rate 21 | limits. 22 | 23 | Installation 24 | ------------ 25 | 26 | PyPi 27 | ~~~~ 28 | 29 | Add this line to your application's requirements.txt: 30 | 31 | .. code:: python 32 | 33 | deckar01-ratelimit 34 | 35 | And then execute: 36 | 37 | .. code:: bash 38 | 39 | $ pip install -r requirements.txt 40 | 41 | Or install it yourself: 42 | 43 | .. code:: bash 44 | 45 | $ pip install deckar01-ratelimit 46 | 47 | GitHub 48 | ~~~~~~ 49 | 50 | Installing the latest version from Github: 51 | 52 | .. code:: bash 53 | 54 | $ git clone https://github.com/deckar01/ratelimit 55 | $ cd ratelimit 56 | $ python setup.py install 57 | 58 | Usage 59 | ----- 60 | 61 | To use this package simply decorate any function that makes an API call: 62 | 63 | .. code:: python 64 | 65 | from ratelimit import limits 66 | 67 | import requests 68 | 69 | FIFTEEN_MINUTES = 900 70 | 71 | @limits(calls=15, period=FIFTEEN_MINUTES) 72 | def call_api(url): 73 | response = requests.get(url) 74 | 75 | if response.status_code != 200: 76 | raise Exception('API response: {}'.format(response.status_code)) 77 | return response 78 | 79 | This function will not be able to make more then 15 API call within a 15 minute 80 | time period. 81 | 82 | The arguments passed into the decorator describe the number of function 83 | invocation allowed over a specified time period (in seconds). If no time period 84 | is specified then it defaults to 15 minutes (the time window imposed by 85 | Twitter). 86 | 87 | If a decorated function is called more times than that allowed within the 88 | specified time period then a ``ratelimit.RateLimitException`` is raised. This 89 | may be used to implement a retry strategy such as an `expoential backoff 90 | `_ 91 | 92 | .. code:: python 93 | 94 | from ratelimit import limits, RateLimitException 95 | from backoff import on_exception, expo 96 | 97 | import requests 98 | 99 | FIFTEEN_MINUTES = 900 100 | 101 | @on_exception(expo, RateLimitException, max_tries=8) 102 | @limits(calls=15, period=FIFTEEN_MINUTES) 103 | def call_api(url): 104 | response = requests.get(url) 105 | 106 | if response.status_code != 200: 107 | raise Exception('API response: {}'.format(response.status_code)) 108 | return response 109 | 110 | Alternatively to cause the current thread to sleep until the specified time 111 | period has ellapsed and then retry the function use the ``sleep_and_retry`` 112 | decorator. This ensures that every function invocation is successful at the 113 | cost of halting the thread. 114 | 115 | .. code:: python 116 | 117 | from ratelimit import limits, sleep_and_retry 118 | 119 | import requests 120 | 121 | FIFTEEN_MINUTES = 900 122 | 123 | @sleep_and_retry 124 | @limits(calls=15, period=FIFTEEN_MINUTES) 125 | def call_api(url): 126 | response = requests.get(url) 127 | 128 | if response.status_code != 200: 129 | raise Exception('API response: {}'.format(response.status_code)) 130 | return response 131 | 132 | Persistence 133 | ~~~~~~~~~~~ 134 | 135 | If a limit needs to be respected between application restarts or shared by 136 | multiple processes, the ``storage`` argument can be used to save the limit 137 | state to disk and load it automatically. 138 | 139 | .. code:: python 140 | 141 | from ratelimit import limits, sleep_and_retry 142 | 143 | import requests 144 | 145 | FIFTEEN_MINUTES = 900 146 | 147 | @sleep_and_retry 148 | @limits(calls=15, period=FIFTEEN_MINUTES, storage='ratelimit.db') 149 | def call_api(url): 150 | response = requests.get(url) 151 | 152 | if response.status_code != 200: 153 | raise Exception('API response: {}'.format(response.status_code)) 154 | return response 155 | 156 | If multiple limits need to be persisted, the ``name`` argument can be used to 157 | store them in the same database using different tables. 158 | 159 | .. code:: python 160 | 161 | from ratelimit import limits, sleep_and_retry 162 | 163 | import requests 164 | 165 | HOUR = 3600 166 | DAY = 24*HOUR 167 | 168 | @sleep_and_retry 169 | @limits(calls=15, period=HOUR, storage='ratelimit.db', name='hourly_limit') 170 | @sleep_and_retry 171 | @limits(calls=100, period=DAY, storage='ratelimit.db', name='daily_limit') 172 | def call_api(url): 173 | response = requests.get(url) 174 | 175 | if response.status_code != 200: 176 | raise Exception('API response: {}'.format(response.status_code)) 177 | return response 178 | 179 | License 180 | ------- 181 | 182 | This project is licensed under the `MIT License `_. 183 | 184 | .. |build| image:: https://github.com/deckar01/ratelimit/actions/workflows/test.yml/badge.svg 185 | :target: https://github.com/deckar01/ratelimit/actions/workflows/test.yml 186 | 187 | .. |maintainability| image:: https://api.codeclimate.com/v1/badges/8bf92a976a1763a93339/maintainability 188 | :target: https://codeclimate.com/github/deckar01/ratelimit/maintainability 189 | :alt: Maintainability 190 | 191 | .. |coverage| image:: https://api.codeclimate.com/v1/badges/8bf92a976a1763a93339/test_coverage 192 | :target: https://codeclimate.com/github/deckar01/ratelimit/test_coverage 193 | :alt: Test Coverage 194 | -------------------------------------------------------------------------------- /pytest.ini: -------------------------------------------------------------------------------- 1 | [pytest] 2 | norecursedirs = venv* .* 3 | addopts = 4 | -r fEsxXw 5 | -vvv 6 | --doctest-modules 7 | --ignore setup.py 8 | --cov-report=term-missing 9 | --cov-report=xml 10 | --cov=ratelimit 11 | -------------------------------------------------------------------------------- /ratelimit/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Function decorator for rate limiting. 3 | 4 | This module provides a function decorator that can be used to wrap a function 5 | such that it will raise an exception if the number of calls to that function 6 | exceeds a maximum within a specified time window. 7 | 8 | For examples and full documentation see the README at 9 | https://github.com/deckar01/ratelimit 10 | """ 11 | from ratelimit.decorators import RateLimitDecorator, sleep_and_retry 12 | from ratelimit.exception import RateLimitException 13 | 14 | limits = RateLimitDecorator 15 | rate_limited = RateLimitDecorator # For backwards compatibility 16 | 17 | __all__ = ["RateLimitException", "limits", "rate_limited", "sleep_and_retry"] 18 | 19 | __version__ = "3.0.2" 20 | -------------------------------------------------------------------------------- /ratelimit/decorators.py: -------------------------------------------------------------------------------- 1 | """ 2 | Rate limit public interface. 3 | 4 | This module includes the decorator used to rate limit function invocations. 5 | Additionally this module includes a naive retry strategy to be used in 6 | conjunction with the rate limit decorator. 7 | """ 8 | from functools import wraps 9 | from math import floor 10 | 11 | import time 12 | import sys 13 | import threading 14 | import sqlite3 15 | 16 | from ratelimit.exception import RateLimitException 17 | 18 | 19 | class RateLimitDecorator(object): 20 | """ 21 | Rate limit decorator class. 22 | """ 23 | 24 | def __init__( 25 | self, 26 | calls=15, 27 | period=900, 28 | raise_on_limit=True, 29 | storage="file:ratelimit?mode=memory&cache=shared", 30 | name="main_limit", 31 | ): 32 | """ 33 | Instantiate a RateLimitDecorator with some sensible defaults. By 34 | default the Twitter rate limiting window is respected (15 calls every 35 | 15 minutes). 36 | 37 | :param int calls: Maximum function invocations allowed within a time period. 38 | :param float period: An upper bound time period (in seconds) before the rate limit resets. 39 | :param bool raise_on_limit: A boolean allowing the caller to avoiding rasing an exception. 40 | :param string storage: An sqlite3 database path for storing the call history. 41 | :param string name: The name of the sqlite3 table. 42 | """ 43 | self.clamped_calls = max(1, min(sys.maxsize, floor(calls))) 44 | self.period = period 45 | self.raise_on_limit = raise_on_limit 46 | 47 | self.database = sqlite3.connect(storage, uri=True, check_same_thread=False) 48 | self.name = name 49 | 50 | # Add thread safety. 51 | self.lock = threading.RLock() 52 | 53 | try: 54 | with self.database: 55 | self.database.execute( 56 | """ 57 | CREATE TABLE {} 58 | (time DATETIME DEFAULT(julianday('now'))) 59 | """.format( 60 | self.name 61 | ) 62 | ) 63 | except sqlite3.OperationalError: 64 | pass 65 | 66 | @property 67 | def _offset(self): 68 | return str(-self.period) + " seconds" 69 | 70 | @property 71 | def _num_calls(self): 72 | query = self.database.execute( 73 | """ 74 | SELECT count(*) FROM {} 75 | WHERE time > julianday('now', '{}') 76 | """.format( 77 | self.name, self._offset 78 | ) 79 | ) 80 | return int(query.fetchone()[0]) 81 | 82 | @property 83 | def _period_remaining(self): 84 | query = self.database.execute( 85 | """ 86 | SELECT julianday('now') - time FROM {} 87 | WHERE time > julianday('now', '{}') 88 | LIMIT 1 89 | """.format( 90 | self.name, self._offset 91 | ) 92 | ) 93 | (age,) = query.fetchone() 94 | oldest_age = 24 * 60 * 60 * float(age) 95 | return max(0, self.period - oldest_age) 96 | 97 | @property 98 | def _too_many_calls(self): 99 | # If the number of attempts to call the function exceeds the 100 | # maximum then raise an exception. 101 | if self._num_calls >= self.clamped_calls: 102 | if self.raise_on_limit: 103 | raise RateLimitException("too many calls", self._period_remaining) 104 | return True 105 | 106 | def _can_call(self): 107 | while True: 108 | try: 109 | with self.database, self.lock: 110 | if self._too_many_calls: 111 | return False 112 | # Clean old calls 113 | query = "DELETE FROM {} WHERE time <= julianday('now', '{}')" 114 | self.database.execute(query.format(self.name, self._offset)) 115 | # Log call 116 | self.database.execute( 117 | "INSERT INTO {} DEFAULT VALUES".format(self.name) 118 | ) 119 | return True 120 | except sqlite3.OperationalError: 121 | pass 122 | 123 | def __call__(self, func): 124 | """ 125 | Return a wrapped function that prevents further function invocations if 126 | previously called within a specified period of time. 127 | 128 | :param function func: The function to decorate. 129 | :return: Decorated function. 130 | :rtype: function 131 | """ 132 | 133 | @wraps(func) 134 | def wrapper(*args, **kargs): 135 | """ 136 | Extend the behaviour of the decorated function, forwarding function 137 | invocations previously called no sooner than a specified period of 138 | time. The decorator will raise an exception if the function cannot 139 | be called so the caller may implement a retry strategy such as an 140 | exponential backoff. 141 | 142 | :param args: non-keyword variable length argument list to the decorated function. 143 | :param kargs: keyworded variable length argument list to the decorated function. 144 | :raises: RateLimitException 145 | """ 146 | if self._can_call(): 147 | return func(*args, **kargs) 148 | 149 | return wrapper 150 | 151 | 152 | def sleep_and_retry(func): 153 | """ 154 | Return a wrapped function that rescues rate limit exceptions, sleeping the 155 | current thread until rate limit resets. 156 | 157 | :param function func: The function to decorate. 158 | :return: Decorated function. 159 | :rtype: function 160 | """ 161 | 162 | @wraps(func) 163 | def wrapper(*args, **kargs): 164 | """ 165 | Call the rate limited function. If the function raises a rate limit 166 | exception sleep for the remaing time period and retry the function. 167 | 168 | :param args: non-keyword variable length argument list to the decorated function. 169 | :param kargs: keyworded variable length argument list to the decorated function. 170 | """ 171 | while True: 172 | try: 173 | return func(*args, **kargs) 174 | except RateLimitException as exception: 175 | time.sleep(exception.period_remaining) 176 | 177 | return wrapper 178 | -------------------------------------------------------------------------------- /ratelimit/exception.py: -------------------------------------------------------------------------------- 1 | """ 2 | Rate limit exceptions. 3 | """ 4 | 5 | 6 | class RateLimitException(Exception): 7 | """ 8 | Rate limit exception class. 9 | """ 10 | 11 | def __init__(self, message, period_remaining): 12 | """ 13 | Custom exception raise when the number of function invocations exceeds 14 | that imposed by a rate limit. Additionally the exception is aware of 15 | the remaining time period after which the rate limit is reset. 16 | 17 | :param string message: Custom exception message. 18 | :param float period_remaining: The time remaining until the rate limit is reset. 19 | """ 20 | super(RateLimitException, self).__init__(message) 21 | self.period_remaining = period_remaining 22 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | black==21.11b1 2 | mock==4.0.3 3 | pytest==6.2.5 4 | pytest-cov==2.10.1 5 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | import ratelimit 4 | 5 | def readme(): 6 | '''Read README file''' 7 | with open('README.rst') as infile: 8 | return infile.read() 9 | 10 | setup( 11 | name='deckar01-ratelimit', 12 | version=ratelimit.__version__, 13 | description='API rate limit decorator', 14 | long_description=readme().strip(), 15 | long_description_content_type='text/x-rst', 16 | author='Jared Deckard', 17 | author_email='jared@shademaps.com', 18 | url='https://github.com/deckar01/ratelimit', 19 | license='MIT', 20 | packages=['ratelimit'], 21 | python_requires='>3.6.0', 22 | install_requires=[], 23 | keywords=[ 24 | 'ratelimit', 25 | 'api', 26 | 'decorator' 27 | ], 28 | classifiers=[ 29 | 'Development Status :: 5 - Production/Stable', 30 | 'Intended Audience :: Developers', 31 | 'License :: OSI Approved :: MIT License', 32 | 'Natural Language :: English', 33 | 'Programming Language :: Python', 34 | 'Topic :: Software Development' 35 | ], 36 | include_package_data=True, 37 | zip_safe=False 38 | ) 39 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/deckar01/ratelimit/d9ee2ba149d4f75e60f2ab4651ccac18d54e7c24/tests/__init__.py -------------------------------------------------------------------------------- /tests/unit/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/deckar01/ratelimit/d9ee2ba149d4f75e60f2ab4651ccac18d54e7c24/tests/unit/__init__.py -------------------------------------------------------------------------------- /tests/unit/decorator_test.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | import sqlite3 3 | from mock import patch 4 | from ratelimit import limits, RateLimitException, sleep_and_retry 5 | 6 | 7 | class TestDecorator(unittest.TestCase): 8 | def sleep(self, name, n): 9 | # Push calls back in time to simulate sleep 10 | self.database.execute( 11 | "UPDATE {} SET time = julianday(time, '{} seconds')".format(name, -n) 12 | ) 13 | self.database.commit() 14 | 15 | @limits(calls=1, period=10) 16 | def increment_a(self): 17 | """ 18 | Increment the counter at most once every 10 seconds. 19 | """ 20 | self.a += 1 21 | 22 | @limits(calls=1, period=10, name="b") 23 | def increment_b(self): 24 | """ 25 | Increment the counter at most once every 10 seconds. 26 | """ 27 | self.b += 1 28 | 29 | @limits(calls=1, period=10, raise_on_limit=False) 30 | def increment_no_exception(self): 31 | """ 32 | Increment the counter at most once every 10 seconds, but w/o rasing an 33 | exception when reaching limit. 34 | """ 35 | self.count += 1 36 | 37 | def setup_method(self, _): 38 | self.count = 0 39 | self.a = 0 40 | self.b = 0 41 | self.database = sqlite3.connect( 42 | "file:ratelimit?mode=memory&cache=shared", uri=True 43 | ) 44 | 45 | def teardown_method(self, _): 46 | self.database.execute("DELETE FROM main_limit") 47 | self.database.commit() 48 | 49 | def test_increment(self): 50 | self.increment_a() 51 | self.assertEqual(self.a, 1) 52 | 53 | def test_exception(self): 54 | self.increment_a() 55 | self.assertRaises(RateLimitException, self.increment_a) 56 | 57 | def test_reset(self): 58 | self.increment_a() 59 | self.sleep("main_limit", 10) 60 | 61 | self.increment_a() 62 | self.assertEqual(self.a, 2) 63 | 64 | def test_no_exception(self): 65 | self.increment_no_exception() 66 | self.increment_no_exception() 67 | 68 | self.assertEqual(self.count, 1) 69 | 70 | def test_separate_decorators(self): 71 | self.increment_a() 72 | self.increment_b() 73 | self.sleep("main_limit", 10) 74 | self.sleep("b", 10) 75 | 76 | self.increment_a() 77 | self.assertEqual(self.a, 2) 78 | self.assertEqual(self.b, 1) 79 | 80 | @patch("ratelimit.decorators.time.sleep") 81 | def test_separate_decorators(self, sleep): 82 | sleep.side_effect = lambda t: self.sleep("main_limit", t) 83 | increment_a = sleep_and_retry(self.increment_a) 84 | 85 | increment_a() 86 | sleep.assert_not_called() 87 | self.assertEqual(self.a, 1) 88 | 89 | increment_a() 90 | sleep.assert_called_once() 91 | assert 9.9 < sleep.call_args[0][0] <= 10 92 | self.assertEqual(self.a, 2) 93 | -------------------------------------------------------------------------------- /tests/unit/threading_test.py: -------------------------------------------------------------------------------- 1 | from threading import Thread 2 | from queue import Queue 3 | import unittest 4 | import sqlite3 5 | from ratelimit import limits, RateLimitException 6 | 7 | 8 | @limits(calls=2, period=10) 9 | def decorate_in_shared(q): 10 | """ 11 | Increment the counter at most once every 10 seconds. 12 | """ 13 | q.put(1) 14 | 15 | 16 | def decorate_in_thread(q): 17 | """ 18 | Increment the counter at most once every 10 seconds. 19 | """ 20 | 21 | @limits(calls=2, period=10) 22 | def run(): 23 | q.put(1) 24 | 25 | run() 26 | 27 | 28 | class TestPreThreading(unittest.TestCase): 29 | def setup_method(self, _): 30 | self.shared = {"count": 0} 31 | self.queue = Queue() 32 | self.threads = [ 33 | Thread(target=decorate_in_shared, args=(self.queue,)) for _ in range(2) 34 | ] 35 | 36 | def teardown_method(self, _): 37 | database = sqlite3.connect("file:ratelimit?mode=memory&cache=shared", uri=True) 38 | database.execute("DELETE FROM main_limit") 39 | database.commit() 40 | database.close() 41 | 42 | def test_increment(self): 43 | self.threads[0].start() 44 | self.threads[1].start() 45 | self.threads[0].join() 46 | self.threads[1].join() 47 | self.assertEqual(self.queue.qsize(), 2) 48 | 49 | def test_exception(self): 50 | self.threads[0].start() 51 | self.threads[1].start() 52 | self.threads[0].join() 53 | self.threads[1].join() 54 | database = sqlite3.connect( 55 | "file:ratelimit?mode=memory&cache=shared", 56 | uri=True, 57 | ) 58 | a = database.execute("SELECT time from main_limit").fetchall() 59 | self.assertRaises(RateLimitException, decorate_in_shared, q=self.queue) 60 | 61 | 62 | class TestPostThreading(unittest.TestCase): 63 | def setup_method(self, _): 64 | self.shared = {"count": 0} 65 | self.queue = Queue() 66 | self.threads = [ 67 | Thread(target=decorate_in_thread, args=(self.queue,)) for _ in range(2) 68 | ] 69 | 70 | def teardown_method(self, _): 71 | database = sqlite3.connect("file:ratelimit?mode=memory&cache=shared", uri=True) 72 | database.execute("DELETE FROM main_limit") 73 | database.commit() 74 | database.close() 75 | 76 | def test_increment(self): 77 | self.threads[0].start() 78 | self.threads[1].start() 79 | self.threads[0].join() 80 | self.threads[1].join() 81 | self.assertEqual(self.queue.qsize(), 2) 82 | 83 | def test_exception(self): 84 | self.threads[0].start() 85 | self.threads[1].start() 86 | self.threads[0].join() 87 | self.threads[1].join() 88 | database = sqlite3.connect( 89 | "file:ratelimit?mode=memory&cache=shared", 90 | uri=True, 91 | ) 92 | a = database.execute("SELECT time from main_limit").fetchall() 93 | self.assertRaises(RateLimitException, decorate_in_thread, q=self.queue) 94 | --------------------------------------------------------------------------------