75 lines
2.5 KiB
Python
75 lines
2.5 KiB
Python
"""HA 阶段 2 单测：DB 共享限流。
|
||
|
||
DbRateLimiter 与内存版同接口：窗口内达到阈值即 False、窗口
|
||
滑过去后恢复、不同 key 互不干扰；关键一条是多实例共享计数
|
||
（两个 server 进程各建一个实例，计数仍合流——多 server 限流
|
||
的前提）。SQLite 跑同款 SQL，Postgres 方言差异由 db.py 适配层
|
||
覆盖（本文件不连真实 Postgres）。
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import tempfile
|
||
import unittest
|
||
from pathlib import Path
|
||
|
||
try:
|
||
from cloud.server import DbRateLimiter, RateLimiter
|
||
except ImportError:
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||
from cloud.server import DbRateLimiter, RateLimiter
|
||
|
||
|
||
class DbRateLimiterTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
self.now = [1000.0]
|
||
|
||
def tearDown(self):
|
||
self.tmp.cleanup()
|
||
|
||
def _limiter(self, per_minute=3, window=60):
|
||
return DbRateLimiter(self.db_path, per_minute=per_minute,
|
||
window_seconds=window,
|
||
time_fn=lambda: self.now[0])
|
||
|
||
def test_allow_up_to_limit_then_false(self):
|
||
lim = self._limiter()
|
||
self.assertTrue(lim.allow("k1"))
|
||
self.assertTrue(lim.allow("k1"))
|
||
self.assertTrue(lim.allow("k1"))
|
||
self.assertFalse(lim.allow("k1"))
|
||
|
||
def test_window_slides_and_recovers(self):
|
||
lim = self._limiter()
|
||
for _ in range(3):
|
||
self.assertTrue(lim.allow("k1"))
|
||
self.assertFalse(lim.allow("k1"))
|
||
self.now[0] += 61
|
||
self.assertTrue(lim.allow("k1"))
|
||
|
||
def test_keys_are_independent(self):
|
||
lim = self._limiter()
|
||
for _ in range(3):
|
||
self.assertTrue(lim.allow("a"))
|
||
self.assertFalse(lim.allow("a"))
|
||
self.assertTrue(lim.allow("b"))
|
||
|
||
def test_two_instances_share_counts(self):
|
||
# 多 server 的核心性质：两个实例（两个进程）计数合流。
|
||
a = self._limiter()
|
||
b = self._limiter()
|
||
self.assertTrue(a.allow("shared"))
|
||
self.assertTrue(a.allow("shared"))
|
||
self.assertTrue(b.allow("shared"))
|
||
self.assertFalse(b.allow("shared"))
|
||
self.assertFalse(a.allow("shared"))
|
||
|
||
def test_memory_limiter_unchanged(self):
|
||
lim = RateLimiter(per_minute=2, window_seconds=60,
|
||
time_fn=lambda: self.now[0])
|
||
self.assertTrue(lim.allow("k"))
|
||
self.assertTrue(lim.allow("k"))
|
||
self.assertFalse(lim.allow("k"))
|