Profeto/tests/test_key_ring.py

"""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)))