如何用pytest测试滑动窗口限流算法的回调与循环逻辑？

这是滑动窗口限流算法的代码实现：

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

针对这个时间依赖的算法，测试核心是模拟时间流逝和验证回调触发逻辑，具体方案如下：

1. 用工具模拟时间，避免依赖真实时间 #

直接用真实时间会导致测试不稳定、耗时久，推荐两种方案：

• 安装freezegun：pip install freezegun，可以精准冻结/快进时间；
• 用pytest-mock替换time.time()函数，手动返回指定时间戳。

2. 用Mock对象验证回调 #

用unittest.mock.Mock替代真实的forward和drop函数，直接断言回调的调用次数、传入参数，无需依赖打印输出判断结果。

3. 模拟循环场景 #

在测试函数中用for循环批量调用handle()，配合时间模拟，复现真实请求流的循环场景。

具体测试用例示例 #

from unittest.mock import Mock
from freezegun import freeze_time
import pytest
from your_module import SlidingWindow  # 替换为你的代码模块名

def test_normal_flow_all_forward():
    # 初始化Mock回调
    forward_mock = Mock()
    drop_mock = Mock()
    # 配置：1秒内允许5个请求
    sw = SlidingWindow(capacity=5, time_unit=1, forward_callback=forward_mock, drop_callback=drop_mock)

    # 连续发送5个请求，全部应该通过
    for packet_id in range(5):
        sw.handle(packet_id)
    
    # 断言回调调用情况
    assert forward_mock.call_count == 5
    assert drop_mock.call_count == 0
    # 验证每个请求都正确传入回调
    for packet_id in range(5):
        forward_mock.assert_any_call(packet_id)

def test_exceed_capacity_trigger_drop():
    forward_mock = Mock()
    drop_mock = Mock()
    sw = SlidingWindow(capacity=5, time_unit=1, forward_callback=forward_mock, drop_callback=drop_mock)

    # 前5个请求通过
    for packet_id in range(5):
        sw.handle(packet_id)
    
    # 第6个请求触发限流
    sw.handle(5)
    assert forward_mock.call_count == 5
    assert drop_mock.call_count == 1
    drop_mock.assert_called_once_with(5)

def test_time_reset_restore_capacity():
    forward_mock = Mock()
    drop_mock = Mock()
    sw = SlidingWindow(capacity=5, time_unit=1, forward_callback=forward_mock, drop_callback=drop_mock)

    # 初始时间发送5个请求
    with freeze_time("2024-01-01 00:00:00"):
        for packet_id in range(5):
            sw.handle(packet_id)
        assert forward_mock.call_count == 5
    
    # 时间流逝1秒，窗口重置，可再接收5个请求
    with freeze_time("2024-01-01 00:00:01"):
        for packet_id in range(5, 10):
            sw.handle(packet_id)
        assert forward_mock.call_count == 10
        assert drop_mock.call_count == 0

def test_partial_time_pass_limit():
    forward_mock = Mock()
    drop_mock = Mock()
    sw = SlidingWindow(capacity=5, time_unit=1, forward_callback=forward_mock, drop_callback=drop_mock)

    with freeze_time("2024-01-01 00:00:00") as frozen_time:
        # 初始发送5个请求
        for packet_id in range(5):
            sw.handle(packet_id)
        assert forward_mock.call_count == 5

        # 时间过了0.5秒，此时窗口剩余容量为2.5，允许2个新请求
        frozen_time.tick(delta=0.5)
        sw.handle(5)
        sw.handle(6)
        # 第7个请求被限流
        sw.handle(7)

        assert forward_mock.call_count == 7
        assert drop_mock.call_count == 1
        drop_mock.assert_called_once_with(7)

模拟大量循环请求的技巧 #

如果需要测试几百次循环的流量场景，直接用嵌套循环配合时间模拟即可：

def test_large_loop_traffic():
    forward_mock = Mock()
    drop_mock = Mock()
    sw = SlidingWindow(capacity=10, time_unit=2, forward_callback=forward_mock, drop_callback=drop_mock)

    with freeze_time("2024-01-01 00:00:00") as frozen_time:
        # 模拟每秒发送8个请求，持续3秒
        for second in range(3):
            for _ in range(8):
                sw.handle(f"{second}_packet")
            frozen_time.tick(delta=1)
        
        # 计算总通过数：初始窗口10，第1秒后剩余2，第2秒新增8个刚好到上限，第3秒仅能通过2个
        assert forward_mock.call_count == 10 + 8 + 2
        assert drop_mock.call_count == 8 - 2

内容的提问来源于stack exchange，提问作者Joshua Leung