rate-limiter.py

from abc import ABC, abstractmethod
import time
import unittest
import heapq

class RateLimiter(ABC):
    @abstractmethod
    def apply(self) -> bool:
        pass

class TokenBucket(RateLimiter):
    def __init__(self):
        self.bucket = 3 * 1000
        self.last = []

    def apply(self, now: int = 0) -> bool:
        if not now:
            now = int(round(time.time() * 1000))
        elapsed = 0
        while self.last and self.last[0] <= now - 1000:
            elapsed += (now - heapq.heappop(self.last))
        heapq.heappush(self.last, now)
        self.bucket += elapsed
        self.bucket = min(self.bucket, 3 * 1000)

        self.bucket -= 1000
        allow = False
        if self.bucket >= 0:
            allow = True
        else:
            self.bucket = 0
        
        return allow

class TokenBucketTests(unittest.TestCase):
    def testLimit(self):
        limiter = TokenBucket()
        self.assertTrue(limiter.apply(100))
        self.assertTrue(limiter.apply(300))
        self.assertTrue(limiter.apply(500))
        self.assertFalse(limiter.apply(700))
        self.assertFalse(limiter.apply(800))
        self.assertTrue(limiter.apply(1100))
        self.assertFalse(limiter.apply(1101))
        self.assertTrue(limiter.apply(1300))

def main():
    unittest.main()

if __name__ == "__main__":
    main()