hair/tests/test_gateway_pool.py

"""gateway/pool.py 隔离期逻辑单测（不依赖真实网络）。

覆盖《docs/网关摘除隔离期实施方案.md》§3.3 的三个用例：
- 被动摘除设置隔离期（offline_until > now）
- 隔离期内探测成功不复活
- 隔离期满后探测成功恢复上线
外加：_try_revive 的阈值/在线短路、mark_worker_unhealthy 对已下线 worker 的幂等、
_health_check_loop 整链路验证、_cooldown_seconds 配置读取。

与 test_gateway_serialization.py 同款：asyncio.run 套同步测试，不依赖 pytest-asyncio。
"""

import asyncio
import time
from unittest.mock import patch

from gateway import pool
from gateway.pool import WorkerState

_HC_CFG = {
    "health_check": {
        "path": "/health",
        "interval_seconds": 0.05,  # 测试用：快速轮询
        "timeout_seconds": 3,
        "unhealthy_threshold": 2,
        "healthy_threshold": 1,
    },
    "shared_password": "test-token",
}


# ---------------------------------------------------------------------------
# 被动摘除 → 隔离期
# ---------------------------------------------------------------------------

def test_passive_ejection_sets_cooldown():
    w = WorkerState(url="http://127.0.0.1:9000", online=True)

    async def run():
        pool._pool_condition = asyncio.Condition()
        with patch.object(pool, "_cooldown_seconds", return_value=60.0):
            await pool.mark_worker_unhealthy(w)

    asyncio.run(run())
    assert w.online is False
    assert w.offline_until > time.monotonic(), "被动摘除必须设置未来的隔离截止时刻"
    assert w.consecutive_failures == 1


def test_passive_ejection_idempotent_when_offline():
    """已下线的 worker 再次被标记：不重复计失败、不续期隔离。"""
    w = WorkerState(url="http://127.0.0.1:9000", online=False)
    w.offline_until = time.monotonic() + 30
    before = w.offline_until

    async def run():
        pool._pool_condition = asyncio.Condition()
        with patch.object(pool, "_cooldown_seconds", return_value=60.0):
            await pool.mark_worker_unhealthy(w)

    asyncio.run(run())
    assert w.online is False
    assert w.consecutive_failures == 0
    assert w.offline_until == before


# ---------------------------------------------------------------------------
# _try_revive：探测成功后是否允许复活
# ---------------------------------------------------------------------------

def test_probe_success_blocked_during_cooldown():
    w = WorkerState(url="http://127.0.0.1:9000")  # offline
    w.consecutive_successes = 1
    w.offline_until = time.monotonic() + 60
    assert pool._try_revive(w, healthy_threshold=1) is False
    assert w.online is False


def test_probe_success_revives_after_cooldown():
    w = WorkerState(url="http://127.0.0.1:9000")
    w.consecutive_successes = 1
    w.offline_until = time.monotonic() - 1  # 隔离期已过
    assert pool._try_revive(w, healthy_threshold=1) is True


def test_try_revive_requires_success_threshold():
    w = WorkerState(url="http://127.0.0.1:9000")
    w.consecutive_successes = 1
    w.offline_until = 0.0
    assert pool._try_revive(w, healthy_threshold=2) is False


def test_try_revive_noop_when_already_online():
    w = WorkerState(url="http://127.0.0.1:9000", online=True)
    w.consecutive_successes = 1
    assert pool._try_revive(w, healthy_threshold=1) is False


# ---------------------------------------------------------------------------
# _health_check_loop 整链路：隔离期内探测 200 不上线，期满后上线
# ---------------------------------------------------------------------------

def test_health_check_loop_respects_cooldown(monkeypatch):
    w = WorkerState(url="http://127.0.0.1:9000", online=False)
    w.offline_until = time.monotonic() + 60  # 处于隔离期

    class _FakeResp:
        status_code = 200

    class _FakeClient:
        async def __aenter__(self):
            return self

        async def __aexit__(self, *args):
            return False

        async def get(self, url, headers=None, timeout=None):
            return _FakeResp()

    async def run():
        monkeypatch.setattr(pool, "_workers", {"w1": w})
        monkeypatch.setattr(pool, "_shutdown_event", asyncio.Event())
        monkeypatch.setattr(pool.httpx, "AsyncClient", _FakeClient)
        task = asyncio.create_task(pool._health_check_loop(_HC_CFG))
        try:
            await asyncio.sleep(0.2)  # 跑若干轮探测
            assert w.consecutive_successes >= 1, "探测应已成功过"
            assert w.online is False, "隔离期内即使探测成功也不得复活"

            w.offline_until = time.monotonic() - 1  # 隔离期结束
            await asyncio.sleep(0.2)
            assert w.online is True, "隔离期满后下一次探测成功应上线"
        finally:
            pool._shutdown_event.set()
            await asyncio.wait_for(task, timeout=2)

    asyncio.run(run())


# ---------------------------------------------------------------------------
# _cooldown_seconds 配置读取
# ---------------------------------------------------------------------------

def test_cooldown_seconds_reads_config(monkeypatch):
    from gateway import config as gw_config
    monkeypatch.setattr(gw_config, "_config_cache",
                        {"health_check": {"cooldown_seconds": 120}})
    assert pool._cooldown_seconds() == 120.0


def test_cooldown_seconds_falls_back_to_60(monkeypatch):
    from gateway import config as gw_config
    monkeypatch.setattr(gw_config, "_config_cache", None)  # get_config 会抛 RuntimeError
    assert pool._cooldown_seconds() == 60.0


# ---------------------------------------------------------------------------
# 主备分组派发（primary_workers）
# ---------------------------------------------------------------------------

_DISPATCH_CFG = {
    "dispatch": {
        "queue_wait_seconds": 0.05,  # 测试用：快速超时
    }
}


def _make_pool(specs):
    """specs: [(url, online, busy, primary), ...] → 构造 pool._workers。"""
    return {
        url: WorkerState(url=url, online=online, busy=busy, primary=primary)
        for url, online, busy, primary in specs
    }


def _run_acquire(workers):
    async def run():
        pool._pool_condition = asyncio.Condition()
        orig = pool._workers
        pool._workers = workers
        try:
            return await pool.acquire_worker(_DISPATCH_CFG)
        finally:
            pool._workers = orig
    return asyncio.run(run())


def test_primary_preferred_over_backup():
    """主 worker 在线时空闲主优先于备 worker。"""
    workers = _make_pool([
        ("http://backup:1", True, False, False),
        ("http://primary:1", True, False, True),
    ])
    w = _run_acquire(workers)
    assert w.url == "http://primary:1"
    assert w.busy is True


def test_backup_used_only_when_all_primary_down():
    """主 worker 全部下线时才用备 worker。"""
    workers = _make_pool([
        ("http://backup:1", True, False, False),
        ("http://primary:1", False, False, True),
        ("http://primary:2", False, False, True),
    ])
    w = _run_acquire(workers)
    assert w.url == "http://backup:1"


def test_backup_not_used_when_primary_busy():
    """主 worker 在线但全忙时不退到备 worker（等待后超时拒绝）。"""
    workers = _make_pool([
        ("http://backup:1", True, False, False),
        ("http://primary:1", True, True, True),
    ])

    async def run():
        pool._pool_condition = asyncio.Condition()
        try:
            await pool.acquire_worker(_DISPATCH_CFG)
        except pool.NoWorkerAvailable:
            return True
        return False

    assert asyncio.run(run()) is True, "主组全忙时不得使用备 worker"
    assert workers["http://backup:1"].busy is False


def test_no_primary_config_all_are_primary():
    """未配置 primary_workers 时全部视为主 worker（兼容旧行为）。"""
    workers = _make_pool([
        ("http://w1:1", True, False, True),
        ("http://w2:1", True, False, True),
    ])
    w = _run_acquire(workers)
    assert w.url == "http://w1:1"


# ---------------------------------------------------------------------------
# 优先级插队（接口5）：全局槽 + worker 池
# ---------------------------------------------------------------------------

_PRIO_CFG = {
    "dispatch": {
        "queue_wait_seconds": 5,   # 测试用：等待足够长，避免排队超时
        "max_global_concurrency": 1,
        "max_queue_wait_seconds": 5,
    }
}


def test_global_slot_priority_jumps_queue():
    """普通请求已排队时，优先级请求插队到最前，先获得槽位。"""
    pool.init_global_slot(_PRIO_CFG)
    order = []

    async def run():
        assert await pool.acquire_global_slot(_PRIO_CFG) is True  # 占住唯一槽

        async def normal():
            if await pool.acquire_global_slot(_PRIO_CFG):
                order.append("normal")
                pool.release_global_slot()

        async def prio():
            if await pool.acquire_global_slot(_PRIO_CFG, priority=1):
                order.append("prio")
                pool.release_global_slot()

        t1 = asyncio.create_task(normal())
        await asyncio.sleep(0.05)  # 确保 normal 先进入等待队列
        t2 = asyncio.create_task(prio())
        await asyncio.sleep(0.05)  # prio 此时应已插队到队首
        pool.release_global_slot()  # 释放槽位 → 应唤醒 prio 而非 normal
        await asyncio.wait_for(asyncio.gather(t1, t2), timeout=3)

    asyncio.run(run())
    assert order == ["prio", "normal"], f"优先级请求应先获得槽位，实际顺序={order}"
    assert pool.get_pool_status()["global_busy"] == 0


def test_global_slot_same_priority_fifo():
    """同优先级仍按 FIFO，不插队。"""
    pool.init_global_slot(_PRIO_CFG)
    order = []

    async def run():
        assert await pool.acquire_global_slot(_PRIO_CFG) is True

        async def waiter(tag):
            if await pool.acquire_global_slot(_PRIO_CFG):
                order.append(tag)
                pool.release_global_slot()

        tasks = []
        for tag in ("a", "b", "c"):
            tasks.append(asyncio.create_task(waiter(tag)))
            await asyncio.sleep(0.05)  # 保证入队顺序
        pool.release_global_slot()
        await asyncio.wait_for(asyncio.gather(*tasks), timeout=3)

    asyncio.run(run())
    assert order == ["a", "b", "c"], f"同优先级应保持 FIFO，实际顺序={order}"


def _run_worker_prio_scenario(acquire_specs):
    """构造「唯一 worker 忙 → 多个请求排队 → 释放」场景，返回获得顺序。

    acquire_specs: [(tag, priority), ...]，按此顺序发起请求。
    """
    workers = _make_pool([
        ("http://w1:1", True, True, True),   # 唯一 worker，初始即忙
    ])
    order = []

    async def run():
        pool._pool_condition = asyncio.Condition()
        pool._worker_waiters.clear()
        orig = pool._workers
        pool._workers = workers
        try:
            async def waiter(tag, prio_):
                w = await pool.acquire_worker(_PRIO_CFG, priority=prio_)
                order.append(tag)
                await pool.release_worker(w)

            tasks = []
            for tag, prio_ in acquire_specs:
                tasks.append(asyncio.create_task(waiter(tag, prio_)))
                await asyncio.sleep(0.05)  # 保证入队顺序
            # 释放 worker → 唤醒优先级最高的等待者
            await pool.release_worker(workers["http://w1:1"])
            await asyncio.wait_for(asyncio.gather(*tasks), timeout=3)
        finally:
            pool._workers = orig
            pool._worker_waiters.clear()

    asyncio.run(run())
    return order


def test_worker_queue_priority_jumps_queue():
    """普通请求已排队时，优先级请求插队到最前，先获得 worker。"""
    order = _run_worker_prio_scenario([("normal", 0), ("prio", 1)])
    assert order == ["prio", "normal"], f"优先级请求应先获得 worker，实际顺序={order}"


def test_worker_queue_same_priority_fifo():
    """同优先级仍按 FIFO，不插队。"""
    order = _run_worker_prio_scenario([("a", 0), ("b", 0), ("c", 0)])
    assert order == ["a", "b", "c"], f"同优先级应保持 FIFO，实际顺序={order}"