"""KeyRing 多 key 轮换单元测试。""" from __future__ import annotations import asyncio import time import pytest from src.data.key_ring import KeyRing, parse_keys, get_key_ring class TestParseKeys: def test_single_key(self): assert parse_keys("abc123") == ["abc123"] def test_comma_separated(self): assert parse_keys("k1, k2,k3") == ["k1", "k2", "k3"] def test_semicolon_separated(self): assert parse_keys("k1;k2;k3") == ["k1", "k2", "k3"] def test_newline_separated(self): assert parse_keys("k1\nk2\nk3") == ["k1", "k2", "k3"] def test_mixed_separators(self): assert parse_keys("k1, k2; k3\nk4") == ["k1", "k2", "k3", "k4"] def test_empty(self): assert parse_keys("") == [] assert parse_keys(None) == [] assert parse_keys(" , ; ") == [] def test_strips_whitespace(self): assert parse_keys(" a , b ") == ["a", "b"] class TestKeyRingSingleKey: """单 key 场景:行为与之前一致。""" def test_get_returns_key(self): ring = KeyRing(["only-key"]) assert ring.get() == "only-key" assert ring.active_key == "only-key" def test_no_rotation(self): ring = KeyRing(["key"]) ring.report_rate_limited() # 单 key 切换后仍是自己 assert ring.get() == "key" def test_empty_keys(self): ring = KeyRing([]) assert ring.get() is None assert ring.active_key is None def test_has_multiple_false(self): ring = KeyRing(["key"]) assert ring.has_multiple is False class TestKeyRingMultiKey: """多 key 场景:429 自动轮换。""" def test_get_rounds_robin(self): ring = KeyRing(["a", "b", "c"]) # 前三次 get 依次返回 a, b, c assert ring.get() == "a" assert ring.get() == "a" # 不报告限流时保持当前 key # 手动推进:通过 report 后 get ring.report_rate_limited("a") # a 被冷却,下一个可用的是 b assert ring.get() == "b" def test_rate_limit_skips_key(self): ring = KeyRing(["a", "b", "c"], cooldown_seconds=60.0) key = ring.get() assert key == "a" new_key = ring.report_rate_limited("a") assert new_key == "b" # 再次 get 应继续是 b(可用) assert ring.get() == "b" def test_cycle_back_to_first(self): ring = KeyRing(["a", "b"], cooldown_seconds=0.1) ring.report_rate_limited("a") # b 可用 assert ring.get() == "b" ring.report_rate_limited("b") # a 仍在冷却,b 也在冷却 → 选最早过期的(可能是 a) key = ring.get() assert key in ("a", "b") def test_cooldown_expires(self): ring = KeyRing(["a", "b"], cooldown_seconds=0.05) ring.report_rate_limited("a") assert ring.get() == "b" # 等 a 的冷却过期 time.sleep(0.08) # 现在 get 应该能找到可用的 key(b 或 a 都行,取决于指针) key = ring.get() assert key in ("a", "b") def test_wait_if_all_blocked(self): ring = KeyRing(["a", "b"], cooldown_seconds=1.0) ring.report_rate_limited("a") ring.report_rate_limited("b") wait = ring.wait_if_all_blocked() assert wait > 0 # 应返回正数等待时间 def test_wait_if_not_all_blocked(self): ring = KeyRing(["a", "b"], cooldown_seconds=1.0) ring.report_rate_limited("a") # b 仍可用 assert ring.wait_if_all_blocked() == 0.0 def test_stats(self): ring = KeyRing(["a" * 12, "b" * 12], cooldown_seconds=1.0) ring.report_rate_limited("a" * 12) st = ring.stats() assert st["total"] == 2 assert st["keys"][0]["blocked_remaining"] > 0 assert st["keys"][1]["blocked_remaining"] == 0 # 脱敏 assert "***" in st["keys"][0]["masked"] or "..." in st["keys"][0]["masked"] def test_all_keys_property(self): ring = KeyRing(["x", "y"]) assert ring.all_keys == ["x", "y"] class TestKeyRingHotUpdate: """热更新 key 列表。""" def test_setter_clears_state(self): ring = KeyRing(["a", "b"]) ring.report_rate_limited("a") assert ring.get() == "b" # 更新 key 列表 ring._keys = ["c", "d"] ring._blocked_until.clear() ring._index = 0 assert ring.get() == "c" class TestGetKeyRing: def test_singleton_per_base(self): r1 = get_key_ring("https://api.test.com", "k1, k2") r2 = get_key_ring("https://api.test.com", "k1, k2") assert r1 is r2 def test_different_base_isolated(self): r1 = get_key_ring("https://a.com", "k1") r2 = get_key_ring("https://b.com", "k2") assert r1 is not r2 assert r1.get() == "k1" assert r2.get() == "k2" def test_hot_update_keys(self): ring = get_key_ring("https://hot.com", "k1, k2") assert set(ring.all_keys) == {"k1", "k2"} # 更新(同 base 会命中缓存,触发热更新) ring2 = get_key_ring("https://hot.com", "k3, k4") assert ring2 is ring assert set(ring.all_keys) == {"k3", "k4"} class TestKeyRingAsyncSafety: """async 并发场景下单 event loop 不需要锁,但验证交替 429 不会死锁。""" @pytest.mark.asyncio async def test_concurrent_get(self): ring = KeyRing(["a", "b", "c"]) async def worker(): for _ in range(20): key = ring.get() assert key in ("a", "b", "c") # 模拟偶发 429 if hash(key) % 3 == 0: ring.report_rate_limited(key) await asyncio.sleep(0) await asyncio.gather(*(worker() for _ in range(5)))