def should_rate_limit(key_id: str) -> bool: """Check if a particular key_id should be rate limited Args: key_id: The key id to check if it needs to be rate limited Returns: boolean true if key should be rate limited, false if not """ # Don't rate limit for 0 (not enabled) if not RATE_LIMIT: return False redis_key = f"{REQUEST_PREFIX_KEY}{key_id}" current_time = time.time() # Get the oldest relevant call oldest = redis.lindex_sync(redis_key, RATE_LIMIT - 1, decode=False) # If this oldest request has happened less than 1 minute ago, then rate limit if oldest and float(oldest) > current_time - 60: return True # Add this as a request for this key redis.lpush_sync(redis_key, str(current_time)) # Trim the list down to the last <rate_limit> calls (no other calls are relevant) # This frees memory from redis, but slows down authorization, so we only do it occasionally if random.randint( 0, 9) == 0: # nosec (this isn't needed for cryptographic purposes) redis.ltrim_sync(redis_key, 0, RATE_LIMIT - 1) return False
def test_lindex(self): redis.lindex_sync("banana", 2) redis.redis_client.lindex.assert_called_once_with("banana", 2)