"""限流中间件测试 — 覆盖内存后端和 Redis 后端降级场景。""" import asyncio import time import pytest from app.middleware.rate_limit import ( MemoryRateLimitBackend, RateLimitBackend, RedisRateLimitBackend, _create_backend, ) # --------------------------------------------------------------------------- # MemoryRateLimitBackend 测试 # --------------------------------------------------------------------------- class TestMemoryRateLimitBackend: """内存后端核心逻辑测试。""" def setup_method(self): self.backend = MemoryRateLimitBackend() @pytest.mark.asyncio async def test_under_limit_not_blocked(self): """未达限流阈值时不被限流。""" now = time.time() for i in range(5): result = await self.backend.is_rate_limited("test:key", now + i * 0.001, 5, 60) assert result is False @pytest.mark.asyncio async def test_at_limit_blocked(self): """达到限流阈值后被限流。""" now = time.time() # 先发送 max_requests 个请求 for i in range(5): await self.backend.is_rate_limited("test:key", now + i * 0.001, 5, 60) # 第 6 个请求应被限流 result = await self.backend.is_rate_limited("test:key", now + 0.006, 5, 60) assert result is True @pytest.mark.asyncio async def test_window_expiry_allows_new_requests(self): """窗口过期后允许新请求。""" now = time.time() # 在窗口开始时发送 5 个请求 for i in range(5): await self.backend.is_rate_limited("test:key", now + i * 0.001, 5, 60) # 窗口过期后发送新请求 new_now = now + 61 # 超过 60 秒窗口 result = await self.backend.is_rate_limited("test:key", new_now, 5, 60) assert result is False @pytest.mark.asyncio async def test_different_keys_independent(self): """不同 key 的限流状态独立。""" now = time.time() # key1 达到限流 for i in range(5): await self.backend.is_rate_limited("key1", now + i * 0.001, 5, 60) result1 = await self.backend.is_rate_limited("key1", now + 0.006, 5, 60) assert result1 is True # key2 不受限流影响 result2 = await self.backend.is_rate_limited("key2", now, 5, 60) assert result2 is False @pytest.mark.asyncio async def test_reset_clears_key(self): """reset 清除指定 key 的限流状态。""" now = time.time() for i in range(5): await self.backend.is_rate_limited("test:key", now + i * 0.001, 5, 60) # 限流中 result = await self.backend.is_rate_limited("test:key", now + 0.006, 5, 60) assert result is True # 重置后不再限流 await self.backend.reset("test:key") result = await self.backend.is_rate_limited("test:key", now + 0.007, 5, 60) assert result is False @pytest.mark.asyncio async def test_cleanup_removes_expired_entries(self): """后台清理任务移除过期记录。""" now = time.time() # 添加一些记录 await self.backend.is_rate_limited("old_key", now - 3700, 5, 60) await self.backend.is_rate_limited("recent_key", now, 5, 60) # 模拟清理 self.backend._requests["old_key"] = [now - 3700] self.backend._requests["recent_key"] = [now] # 手动触发清理逻辑 expired_keys = [] for key in list(self.backend._requests.keys()): self.backend._requests[key] = [t for t in self.backend._requests[key] if now - t < 3600] if not self.backend._requests[key]: expired_keys.append(key) for key in expired_keys: del self.backend._requests[key] assert "old_key" not in self.backend._requests assert "recent_key" in self.backend._requests # --------------------------------------------------------------------------- # RedisRateLimitBackend 测试（无 Redis 连接时的降级行为） # --------------------------------------------------------------------------- class TestRedisRateLimitBackendFallback: """Redis 后端在连接失败时的降级行为。""" @pytest.mark.asyncio async def test_redis_unavailable_allows_request(self): """Redis 不可用时放行请求（不阻塞服务）。""" backend = RedisRateLimitBackend() now = time.time() # Redis 连接失败时应返回 False（不被限流） result = await backend.is_rate_limited("test:key", now, 5, 60) assert result is False @pytest.mark.asyncio async def test_redis_reset_silently_fails(self): """Redis reset 失败时静默处理。""" backend = RedisRateLimitBackend() # 不应抛出异常 await backend.reset("test:key") # --------------------------------------------------------------------------- # Backend 工厂函数测试 # --------------------------------------------------------------------------- class TestCreateBackend: """_create_backend 工厂函数测试。""" def test_default_creates_memory_backend(self, monkeypatch): """默认配置创建内存后端。""" monkeypatch.setattr("app.middleware.rate_limit.settings.RATE_LIMIT_BACKEND", "memory") backend = _create_backend() assert isinstance(backend, MemoryRateLimitBackend) def test_redis_with_empty_url_falls_back_to_memory(self, monkeypatch): """RATE_LIMIT_BACKEND=redis 但 REDIS_URL 为空时降级到内存。""" monkeypatch.setattr("app.middleware.rate_limit.settings.RATE_LIMIT_BACKEND", "redis") monkeypatch.setattr("app.middleware.rate_limit.settings.REDIS_URL", "") backend = _create_backend() assert isinstance(backend, MemoryRateLimitBackend) def test_redis_with_url_creates_redis_backend(self, monkeypatch): """RATE_LIMIT_BACKEND=redis 且 REDIS_URL 非空时创建 Redis 后端。""" monkeypatch.setattr("app.middleware.rate_limit.settings.RATE_LIMIT_BACKEND", "redis") monkeypatch.setattr("app.middleware.rate_limit.settings.REDIS_URL", "redis://localhost:6379/0") backend = _create_backend() assert isinstance(backend, RedisRateLimitBackend) def test_unknown_backend_defaults_to_memory(self, monkeypatch): """未知后端类型默认使用内存。""" monkeypatch.setattr("app.middleware.rate_limit.settings.RATE_LIMIT_BACKEND", "unknown") backend = _create_backend() assert isinstance(backend, MemoryRateLimitBackend) # --------------------------------------------------------------------------- # RateLimitBackend 接口测试 # --------------------------------------------------------------------------- class TestRateLimitBackendInterface: """验证 RateLimitBackend 抽象接口。""" def test_cannot_instantiate_abstract_class(self): """不能直接实例化抽象类。""" with pytest.raises(TypeError): RateLimitBackend() def test_memory_backend_is_subclass(self): """MemoryRateLimitBackend 是 RateLimitBackend 的子类。""" assert issubclass(MemoryRateLimitBackend, RateLimitBackend) def test_redis_backend_is_subclass(self): """RedisRateLimitBackend 是 RateLimitBackend 的子类。""" assert issubclass(RedisRateLimitBackend, RateLimitBackend)