def test_peek(self):
PyRateLimit.init(redis_host="localhost", redis_port=6379)
limit = PyRateLimit(period=10, limit=10)
for x in range(0, 10):
self.assertTrue(limit.attempt('test_namespace2'))
self.assertTrue(limit.is_rate_limited('test_namespace2'))
time.sleep(10)
self.assertFalse(limit.is_rate_limited('test_namespace2'))
def test_peek(self):
PyRateLimit.init(redis_host="localhost", redis_port=6379)
limit = PyRateLimit()
limit.create(
10, # rate limit period in seconds
10) # no of attempts in the time period
for x in range(0, 10):
self.assertTrue(limit.attempt('test_namespace2'))
self.assertTrue(limit.is_rate_limited('test_namespace2'))
time.sleep(10)
self.assertFalse(limit.is_rate_limited('test_namespace2'))
def test_throttle(self):
PyRateLimit.init(redis_host="localhost", redis_port=6379)
limit = PyRateLimit(period=10, limit=10)
for x in range(0, 20):
time.sleep(.5)
if x < 10:
self.assertTrue(limit.attempt('test_namespace'))
else:
self.assertFalse(limit.attempt('test_namespace'))
time.sleep(6)
self.assertTrue(limit.attempt('test_namespace'))
def wait_for_token(self, key):
pylimit = PyRateLimit(int(self.throttle_period.total_seconds()),
self.rate)
throttle_times = 0
current_time = int(round(time.time() * 1000000))
while not pylimit.attempt(key, timestamp=current_time):
throttle_times += 1
_logger.debug("Function {} being throttle".format(key))
time.sleep(5)
return throttle_times
def test_throttle(self):
PyRateLimit.init(redis_host="localhost", redis_port=6379)
limit = PyRateLimit()
limit.create(
10, # rate limit period in seconds
10) # no of attempts in the time period
for x in range(0, 20):
time.sleep(.5)
if x < 10:
self.assertTrue(limit.attempt('test_namespace'))
else:
self.assertFalse(limit.attempt('test_namespace'))
time.sleep(6)
self.assertTrue(limit.attempt('test_namespace'))
def __init__(self, *args):
self.logger = logging.getLogger(__name__)
self.base_url = 'https://www.virustotal.com/vtapi/v2/'
self.headers = {
"Accept-Encoding": "gzip, deflate",
"User-Agent": "Anteater"
}
self.HTTP_OK = 200
self.HTTP_BAD = 400
self.HTTP_FORBIDDEN = 403
self.HTTP_RATE_EXCEEDED = 204
self.public_api_sleep_time = 20
self.logger = logging.getLogger(__name__)
self.uuid = uuid.uuid4()
self.config = six.moves.configparser.SafeConfigParser()
self.config.read('anteater.conf')
try:
conn = redis.StrictRedis(host='localhost', port=6379, password='')
conn.ping()
PyRateLimit.init(redis_host="localhost", redis_port=6379)
except Exception as ex:
self.logger.error('Error: %s', ex)
exit('Failed to connect, terminating.')
self.limit = PyRateLimit()
try:
vt_rate_type = self.config.get('config', 'vt_rate_type')
except six.moves.configparser.NoSectionError:
self.logger.error(
"A config section is required for vt_rate_type with a public | private option "
)
sys.exit(1)
patten = re.compile(r'\bpublic\b|\bprivate\b')
if not patten.match(vt_rate_type):
self.logger.error("Unrecognized %s option for vt_rate_type",
vt_rate_type)
sys.exit(1)
if vt_rate_type == 'public':
self.limit.create(21, 1)
elif vt_rate_type == 'private':
self.limit.create(1, 1)
def test_exception(self):
limit = PyRateLimit()
self.assertRaises(PyRateLimitException, limit.attempt,
'test_namespace')