e407b78506
* Remove six for PY2 compatability since our dependencies are not, PY2 is legacy. * Removed usages of utils from future/past libraries since they are deprecated. This includes 'from __future__ ...' and 'from past...' statements. * Removed compatibility imports for PY2, switched from six imports to standard library imports. * Removed utils code blocks that handle operations depending on PY2/3 versions. * Removed 'from __future__ ...' lines from templates/code generators * Used PY3 syntaxes in place of PY2 compatible blocks. eg: metaclass
116 lines
3.5 KiB
Python
116 lines
3.5 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
# Copyright (c) 2020, Frappe Technologies Pvt. Ltd. and Contributors
|
|
# MIT License. See license.txt
|
|
|
|
import unittest
|
|
from werkzeug.wrappers import Response
|
|
import time
|
|
|
|
import frappe
|
|
import frappe.rate_limiter
|
|
from frappe.rate_limiter import RateLimiter
|
|
from frappe.utils import cint
|
|
|
|
|
|
class TestRateLimiter(unittest.TestCase):
|
|
def setUp(self):
|
|
pass
|
|
|
|
def test_apply_with_limit(self):
|
|
frappe.conf.rate_limit = {"window": 86400, "limit": 1}
|
|
frappe.rate_limiter.apply()
|
|
|
|
self.assertTrue(hasattr(frappe.local, "rate_limiter"))
|
|
self.assertIsInstance(frappe.local.rate_limiter, RateLimiter)
|
|
|
|
frappe.cache().delete(frappe.local.rate_limiter.key)
|
|
delattr(frappe.local, "rate_limiter")
|
|
|
|
def test_apply_without_limit(self):
|
|
frappe.conf.rate_limit = None
|
|
frappe.rate_limiter.apply()
|
|
|
|
self.assertFalse(hasattr(frappe.local, "rate_limiter"))
|
|
|
|
def test_respond_over_limit(self):
|
|
limiter = RateLimiter(0.01, 86400)
|
|
time.sleep(0.01)
|
|
limiter.update()
|
|
|
|
frappe.conf.rate_limit = {"window": 86400, "limit": 0.01}
|
|
self.assertRaises(frappe.TooManyRequestsError, frappe.rate_limiter.apply)
|
|
frappe.rate_limiter.update()
|
|
|
|
response = frappe.rate_limiter.respond()
|
|
|
|
self.assertIsInstance(response, Response)
|
|
self.assertEqual(response.status_code, 429)
|
|
|
|
headers = frappe.local.rate_limiter.headers()
|
|
self.assertIn("Retry-After", headers)
|
|
self.assertNotIn("X-RateLimit-Used", headers)
|
|
self.assertIn("X-RateLimit-Reset", headers)
|
|
self.assertIn("X-RateLimit-Limit", headers)
|
|
self.assertIn("X-RateLimit-Remaining", headers)
|
|
self.assertTrue(int(headers["X-RateLimit-Reset"]) <= 86400)
|
|
self.assertEqual(int(headers["X-RateLimit-Limit"]), 10000)
|
|
self.assertEqual(int(headers["X-RateLimit-Remaining"]), 0)
|
|
|
|
frappe.cache().delete(limiter.key)
|
|
frappe.cache().delete(frappe.local.rate_limiter.key)
|
|
delattr(frappe.local, "rate_limiter")
|
|
|
|
def test_respond_under_limit(self):
|
|
frappe.conf.rate_limit = {"window": 86400, "limit": 0.01}
|
|
frappe.rate_limiter.apply()
|
|
frappe.rate_limiter.update()
|
|
response = frappe.rate_limiter.respond()
|
|
self.assertEqual(response, None)
|
|
|
|
frappe.cache().delete(frappe.local.rate_limiter.key)
|
|
delattr(frappe.local, "rate_limiter")
|
|
|
|
def test_headers_under_limit(self):
|
|
frappe.conf.rate_limit = {"window": 86400, "limit": 0.01}
|
|
frappe.rate_limiter.apply()
|
|
frappe.rate_limiter.update()
|
|
headers = frappe.local.rate_limiter.headers()
|
|
self.assertNotIn("Retry-After", headers)
|
|
self.assertIn("X-RateLimit-Reset", headers)
|
|
self.assertTrue(int(headers["X-RateLimit-Reset"] < 86400))
|
|
self.assertEqual(int(headers["X-RateLimit-Used"]), frappe.local.rate_limiter.duration)
|
|
self.assertEqual(int(headers["X-RateLimit-Limit"]), 10000)
|
|
self.assertEqual(int(headers["X-RateLimit-Remaining"]), 10000)
|
|
|
|
frappe.cache().delete(frappe.local.rate_limiter.key)
|
|
delattr(frappe.local, "rate_limiter")
|
|
|
|
def test_reject_over_limit(self):
|
|
limiter = RateLimiter(0.01, 86400)
|
|
time.sleep(0.01)
|
|
limiter.update()
|
|
|
|
limiter = RateLimiter(0.01, 86400)
|
|
self.assertRaises(frappe.TooManyRequestsError, limiter.apply)
|
|
|
|
frappe.cache().delete(limiter.key)
|
|
|
|
def test_do_not_reject_under_limit(self):
|
|
limiter = RateLimiter(0.01, 86400)
|
|
time.sleep(0.01)
|
|
limiter.update()
|
|
|
|
limiter = RateLimiter(0.02, 86400)
|
|
self.assertEqual(limiter.apply(), None)
|
|
|
|
frappe.cache().delete(limiter.key)
|
|
|
|
def test_update_method(self):
|
|
limiter = RateLimiter(0.01, 86400)
|
|
time.sleep(0.01)
|
|
limiter.update()
|
|
|
|
self.assertEqual(limiter.duration, cint(frappe.cache().get(limiter.key)))
|
|
|
|
frappe.cache().delete(limiter.key)
|