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), )