def test_using_redis_time():
redis.flushall()
rl = RateLimiter(tiers=[simple_daily_tier], use_redis_time=True)
# don't loop infinitely if test is failing
count = 0
while count < 20:
try:
rl.check_limit("test-key", simple_daily_tier.name)
count += 1
except RateLimitExceeded:
break
assert count == 10
def test_check_limit_per_minute(tier, reset_time):
redis.flushall()
rl = RateLimiter(tiers=[tier], use_redis_time=False)
count = 0
with freeze_time() as frozen:
# don't loop infinitely if test is failing
while count < 20:
try:
rl.check_limit("test-key", tier.name)
count += 1
except RateLimitExceeded as e:
print(e)
break
# assert that we broke after 10
assert count == 10
# resets after a given time
frozen.tick(reset_time)
assert rl.check_limit("test-key", tier.name)
def test_get_daily_usage():
redis.flushall()
rl = RateLimiter(tiers=[unlimited_tier],
use_redis_time=False,
track_daily_usage=True)
# make Nth day have N calls
for n in range(1, 10):
with freeze_time(f"2020-01-0{n}"):
for _ in range(n):
rl.check_limit("test-key", unlimited_tier.name)
with freeze_time("2020-01-15"):
usage = rl.get_usage_since("test-key", datetime.date(2020, 1, 1))
assert usage[0] == DailyUsage(datetime.date(2020, 1, 1), 1)
assert usage[3] == DailyUsage(datetime.date(2020, 1, 4), 4)
assert usage[8] == DailyUsage(datetime.date(2020, 1, 9), 9)
assert usage[9] == DailyUsage(datetime.date(2020, 1, 10), 0)
assert usage[14] == DailyUsage(datetime.date(2020, 1, 15), 0)
assert len(usage) == 15
def test_multiple_prefix():
redis.flushall()
rl1 = RateLimiter(tiers=[simple_daily_tier],
use_redis_time=True,
prefix="zone1")
rl2 = RateLimiter(tiers=[simple_daily_tier],
use_redis_time=True,
prefix="zone2")
# don't loop infinitely if test is failing
count = 0
while count < 20:
try:
rl1.check_limit("test-key", simple_daily_tier.name)
rl2.check_limit("test-key", simple_daily_tier.name)
count += 1
except RateLimitExceeded:
break
assert count == 10
def test_get_daily_usage_untracked():
redis.flushall()
rl = RateLimiter(tiers=[unlimited_tier],
use_redis_time=False,
track_daily_usage=False)
# make Nth day have N calls
for n in range(1, 10):
with freeze_time(f"2020-01-0{n}"):
for _ in range(n):
rl.check_limit("test-key", unlimited_tier.name)
# values would be incorrect (likely zero), warn the caller
with pytest.raises(RuntimeError):
rl.get_usage_since("test-key", datetime.date(2020, 1, 1))
def test_invalid_tier():
redis.flushall()
rl = RateLimiter(tiers=[simple_daily_tier], use_redis_time=True)
with pytest.raises(ValueError):
rl.check_limit("test-key", "non-existent-tier")
import json
from django.conf import settings
from django.contrib.auth.decorators import login_required
from django.http import HttpResponse, JsonResponse
from django.shortcuts import render, get_object_or_404
from django.views.decorators.csrf import ensure_csrf_cookie
from .models import WidgetConfig, WidgetType
from rrl import RateLimiter, Tier, RateLimitExceeded
limiter = RateLimiter(
prefix="widgets",
tiers=[Tier("default", 120, 0, 2000)],
use_redis_time=False,
track_daily_usage=True,
)
def index(request):
if request.user.is_authenticated:
your_widgets = list(WidgetConfig.objects.filter(owner=request.user))
else:
your_widgets = []
return render(request, "index.html", {"your_widgets": your_widgets})
@ensure_csrf_cookie
@login_required
def configure(request):
if request.method == "GET":
def handle(self, *args, **options):
# day -> (key, endpoint) -> #
self.count_by_day = defaultdict(Counter)
self.duration_by_day = defaultdict(
lambda: defaultdict(lambda: defaultdict(float))
)
self.lines = 0
for filename in options["filenames"]:
self.process_file(filename)
print(f"processed {self.lines} lines")
# don't use the oldest day, it is probably a partial and will overwrite good data
newest_day = "2000-01-01"
oldest_day = "2100-01-01"
for day in self.count_by_day.keys():
if day > newest_day:
newest_day = day
if day < oldest_day:
oldest_day = day
print(f"found logs from {oldest_day} to {newest_day}, dropping {oldest_day}")
keys = {key.api_key: key for key in Profile.objects.all()}
limiter = RateLimiter(
prefix="v3", tiers=[], use_redis_time=False, track_daily_usage=True
)
for key in keys:
usage = limiter.get_usage_since(
key, datetime.date.today() - datetime.timedelta(days=7)
)
for daily_usage in usage:
UsageReport.objects.update_or_create(
profile=keys[key],
date=daily_usage.date,
endpoint="v3",
defaults=dict(calls=daily_usage.calls, total_duration_seconds=0),
)
for day, counter in self.count_by_day.items():
# skip oldest day
if day == oldest_day:
continue
# build log-based usage reports
for (key, endpoint), calls in counter.items():
duration = self.duration_by_day[day][key][endpoint]
# convert key
try:
profile = keys[key]
except KeyError:
print(f"unknown key {key} with {calls} calls")
continue
UsageReport.objects.update_or_create(
profile=profile,
date=day,
endpoint=endpoint,
defaults=dict(calls=calls, total_duration_seconds=duration),
)