How to test a rate limiting algorithm?

Question

This is the sliding window rate limiting algorithm I got from this page :

from time import time, sleep

class SlidingWindow:

    def __init__(self, capacity, time_unit, forward_callback, drop_callback):
        self.capacity = capacity
        self.time_unit = time_unit
        self.forward_callback = forward_callback
        self.drop_callback = drop_callback

        self.cur_time = time()
        self.pre_count = capacity
        self.cur_count = 0

    def handle(self, packet):

        if (time() - self.cur_time) > self.time_unit:
            self.cur_time = time()
            self.pre_count = self.cur_count
            self.cur_count = 0

        ec = (self.pre_count * (self.time_unit - (time() - self.cur_time)) / self.time_unit) + self.cur_count

        if (ec > self.capacity):
            return self.drop_callback(packet)

        self.cur_count += 1
        return self.forward_callback(packet)


def forward(packet):
    print("Packet Forwarded: " + str(packet))


def drop(packet):
    print("Packet Dropped: " + str(packet))


throttle = SlidingWindow(5, 1, forward, drop)

packet = 0

while True:
    sleep(0.1)
    throttle.handle(packet)
    packet += 1

How can I test this time sensitive algorithm in pytest? I want to test forward_callback is getting called when the request (by calling handle() ) passes, and drop_callback is getting called when the request fails. How could I simulate a loop in pytest?

Answer 1

You should test the SlidingWindow class directly, without testing the loop in the end. The reason is that the SlidingWindow is independent and it should work regardless of how it's used.

Since the class is heavily dependent on the current time, you can use pytest-freezegun to simulate the time of the requests.

A test example would be something similar to:

from datetime import datetime, timedelta
from mock import MagicMock, call

from sliding_window import SlidingWindow


def test_moving_date(freezer, mg):
    start = datetime(year=2012, month=1, day=14, hour=12, minute=0, second=1)
    freezer.move_to(start)
    forward, drop = MagicMock(), MagicMock()
    throttle = SlidingWindow(5, 1, forward, drop)  # allow rate ~5 per second

    # first two should be accepted, while the 3rd should be dropped
    throttle.handle(0)
    freezer.move_to(start + timedelta(milliseconds=200))
    throttle.handle(1)
    freezer.move_to(start + timedelta(milliseconds=250))
    throttle.handle(2)

    forward.assert_has_calls(calls=[call(0), call(1)])
    drop.assert_called_once_with(2)