zTools2/tests/manual_whiteboard_hub.py

"""白板 hub 内存与心跳验证脚本（手动烟测，非 pytest）。

验证点：
1. 连接接入 -> hub._boards 出现该 board 的 set
2. disconnect -> 连接从 set 移除，空 set 从 dict 删除（无内存泄漏）
3. 心跳失活：reaper 扫描超过 timeout 的连接并移除
4. 删除白板 -> close_board 踢出该 board 所有连接并清理 dict
5. broadcast 对失败连接自动 disconnect

用伪造的 WebSocket 对象（不依赖真实网络）跑，直接断言 hub 内部状态。
运行：.venv/bin/python tests/manual_whiteboard_hub.py
"""

from __future__ import annotations

import asyncio
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from app.services.whiteboard_hub import Connection, WhiteboardHub  # noqa: E402


class FakeWS:
    """最小化的 WebSocket 替身：记录发送/关闭，receive 永不返回。"""

    def __init__(self) -> None:
        self.sent: list[dict] = []
        self.closed = False
        self._fail_send = False

    async def accept(self) -> None:
        pass

    async def send_json(self, msg: dict) -> None:
        if self._fail_send:
            raise RuntimeError("send failed (simulated)")
        self.sent.append(msg)

    async def close(self) -> None:
        self.closed = True


async def main() -> None:
    # 用一个独立 hub 实例，避免污染全局单例；构造时读 config 默认 3s*5=15s
    hub = WhiteboardHub()
    print(f"timeout_seconds = {hub.timeout_seconds} (期望 15)")
    assert hub.timeout_seconds == 15, "心跳超时应为 3*5=15s"

    # ---- 1. 接入 ----
    ws1 = FakeWS()
    c1 = Connection(websocket=ws1, board_id="b1", client_id="c1")
    await hub.register(c1)
    assert "b1" in hub._boards, "接入后 b1 应存在"
    assert hub.connection_count("b1") == 1
    print("1. 接入 OK")

    # ---- 2. disconnect 清理空 set ----
    ws2 = FakeWS()
    c2 = Connection(websocket=ws2, board_id="b1", client_id="c2")
    await hub.register(c2)
    assert hub.connection_count("b1") == 2
    await hub.disconnect(c2)
    assert hub.connection_count("b1") == 1, "c2 移除后应剩 1"
    assert ws2.closed, "c2 的 ws 应被关闭"
    await hub.disconnect(c1)
    assert "b1" not in hub._boards, "空 set 应从 dict 删除（防泄漏）"
    print("2. disconnect 清理空 set OK")

    # ---- 3. 心跳失活 reaper ----
    # 手动把 last_heartbeat 调到很久以前，触发 reaper 移除
    ws3 = FakeWS()
    c3 = Connection(websocket=ws3, board_id="b2", client_id="c3")
    # 模拟 30s 前心跳
    c3.last_heartbeat = asyncio.get_event_loop().time() - 30
    await hub.register(c3)
    assert hub.connection_count("b2") == 1
    await hub._reap_once()
    assert hub.connection_count("b2") == 0, "失活连接应被 reaper 移除"
    assert ws3.closed, "失活连接的 ws 应被关闭"
    assert "b2" not in hub._boards, "移除后空 set 应清理"
    print("3. 心跳失活 reaper OK")

    # ---- 4. 心跳未超时的连接不被移除 ----
    ws4 = FakeWS()
    c4 = Connection(websocket=ws4, board_id="b3", client_id="c4")
    await hub.register(c4)
    await hub._reap_once()
    assert hub.connection_count("b3") == 1, "正常心跳连接不应被移除"
    print("4. 正常连接保留 OK")

    # ---- 5. close_board 踢出所有连接 ----
    ws5 = FakeWS()
    c5 = Connection(websocket=ws5, board_id="b3", client_id="c5")
    await hub.register(c5)
    assert hub.connection_count("b3") == 2
    await hub.close_board("b3")
    assert hub.connection_count("b3") == 0, "close_board 后连接应清空"
    assert ws4.closed and ws5.closed, "该 board 所有 ws 应被关闭"
    assert "b3" not in hub._boards, "close_board 后 dict 应清理"
    assert any(m.get("type") == "error" for m in ws4.sent), "应先发 error 通知"
    print("5. close_board 踢出连接 OK")

    # ---- 6. broadcast 排除发送者 + 失败连接自动清理 ----
    wsA = FakeWS()
    wsB = FakeWS()
    wsB._fail_send = True  # 模拟 B 发送失败
    cA = Connection(websocket=wsA, board_id="b4", client_id="A")
    cB = Connection(websocket=wsB, board_id="b4", client_id="B")
    await hub.register(cA)
    await hub.register(cB)
    await hub.broadcast("b4", {"type": "stroke"}, exclude=cA)
    assert not wsA.sent, "排除发送者：A 不应收到"
    assert wsB.closed, "发送失败的 B 应被自动 disconnect"
    assert hub.connection_count("b4") == 1, "B 移除后应剩 A"
    print("6. broadcast 排除 + 失败自动清理 OK")

    # ---- 7. 幂等 disconnect ----
    await hub.disconnect(cA)
    await hub.disconnect(cA)  # 重复调用不应报错
    assert "b4" not in hub._boards
    print("7. 幂等 disconnect OK")

    print("\n全部通过 ✅")


if __name__ == "__main__":
    asyncio.run(main())