def setUp(self):
cache.init(
host=self.host,
port=self.port,
namespace=self.namespace,
)
# Making sure there are no items in cache before we run each test
cache.flush_all()
def setUp(self):
cache.init(
host=self.host,
port=self.port,
namespace=self.namespace,
)
# Making sure there are no items in cache before we run each test
cache.flush_all()
def test_ratelimit(self):
""" Tests that the ratelimit decorator works
"""
# Set the limits as per defines in this class
set_rate_limits(self.max_token_requests, self.max_ip_requests,
self.ratelimit_window)
# create an app
app = flask.CustomFlask(__name__)
self.assertIsNotNone(app)
app.debug = True
app.config['SECRET_KEY'] = 'this is a totally secret key btw'
app.init_debug_toolbar()
@app.after_request
def after_request_callbacks(response):
return inject_x_rate_headers(response)
# add a dummy route
@app.route('/')
@ratelimit()
def index():
return '<html><body>test</body></html>'
def print_headers(response):
print("X-RateLimit-Remaining",
response.headers['X-RateLimit-Remaining'])
print("X-RateLimit-Limit", response.headers['X-RateLimit-Limit'])
print("X-RateLimit-Reset", response.headers['X-RateLimit-Reset'])
print("X-RateLimit-Reset-In",
response.headers['X-RateLimit-Reset-In'])
print()
def make_requests(client, nominal_num_requests, token=None):
print("===== make %d requests" % nominal_num_requests)
# make one more than the allowed number of requests to catch the 429
num_requests = nominal_num_requests + 1
# make a specified number of requests
while True:
reset_time = 0
restart = False
for i in range(num_requests):
if token:
response = client.get('/',
headers={'Authorization': token})
else:
response = client.get('/')
if reset_time == 0:
reset_time = response.headers['X-RateLimit-Reset']
if reset_time != response.headers['X-RateLimit-Reset']:
# Whoops, we didn't get our tests done before the window expired. start over.
restart = True
# when restarting we need to do one request less, since the current requests counts to the new window
num_requests = nominal_num_requests
break
if i == num_requests - 1:
self.assertEqual(response.status_code, 429)
else:
self.assertEqual(response.status_code, 200)
self.assertEqual(
int(response.headers['X-RateLimit-Remaining']),
num_requests - i - 2)
print_headers(response)
sleep(1.1)
if not restart:
break
client = app.test_client()
# Make a pile of requests based on IP address
make_requests(client, self.max_ip_requests)
# Set a user token and make requests based on the token
cache.flush_all()
set_user_validation_function(validate_user)
set_rate_limits(self.max_token_requests, self.max_ip_requests,
self.ratelimit_window)
make_requests(client,
self.max_token_requests,
token="Token %s" % valid_user)
def clear_cache():
"""Clear the cache"""
# inits the db engine and cache
create_app()
cache.flush_all()
def clear_cache():
"""Clear the cache."""
cache.flush_all()
def clear_memcached():
with frontend.create_app().app_context():
cache.flush_all()
click.echo("Flushed everything from memcached.")
def clear_memcached():
with frontend.create_app().app_context():
cache.flush_all()
click.echo("Flushed everything from memcached.")