浅念
81 lines
3.5 KiB
Python
81 lines
3.5 KiB
Python
# Task 8: WS 每 IP 并发连接上限(spec §6.1 P2)。
|
|
#
|
|
# 两层测试:
|
|
# - 纯单元: ConnectionLimiter 算法(名额/键隔离/归零清 key/no-op release);
|
|
# - 接入: /ws/chat 握手期按 (租户, IP) 限流——第 limit+1 条并发连接被
|
|
# 1013 拒(close 代 accept),连接关闭后名额必须释放。
|
|
import json
|
|
from contextlib import ExitStack
|
|
|
|
import pytest
|
|
from starlette.websockets import WebSocketDisconnect
|
|
|
|
from app.engine.services.conn_limiter import ConnectionLimiter
|
|
|
|
# ---- 纯单元: 算法 -----------------------------------------------------------
|
|
|
|
|
|
def test_limit_per_ip():
|
|
lim = ConnectionLimiter(limit=2)
|
|
assert lim.try_acquire("t1", "1.1.1.1") and lim.try_acquire("t1", "1.1.1.1")
|
|
assert not lim.try_acquire("t1", "1.1.1.1") # 第 3 个拒绝
|
|
lim.release("t1", "1.1.1.1")
|
|
assert lim.try_acquire("t1", "1.1.1.1") # 释放后可再进
|
|
assert lim.try_acquire("t2", "1.1.1.1") # 其他租户不受影响
|
|
assert lim.try_acquire("t1", "2.2.2.2") # 其他 IP 不受影响
|
|
|
|
|
|
def test_count_zero_cleans_key():
|
|
# 计数归零删 key: 长期运行下内部 dict 不随历史 (租户,IP) 对无限增长
|
|
lim = ConnectionLimiter(limit=3)
|
|
lim.try_acquire("t", "1.1.1.1")
|
|
lim.try_acquire("t", "1.1.1.1")
|
|
lim.release("t", "1.1.1.1")
|
|
assert ("t", "1.1.1.1") in lim._counts # 还剩 1 个,键保留
|
|
lim.release("t", "1.1.1.1")
|
|
assert lim._counts == {} # 归零即清空
|
|
|
|
|
|
def test_release_unheld_key_is_safe_noop():
|
|
lim = ConnectionLimiter(limit=2)
|
|
lim.release("nobody", "9.9.9.9") # 从未 acquire 过
|
|
lim.try_acquire("t", "1.1.1.1")
|
|
lim.release("t", "1.1.1.1")
|
|
lim.release("t", "1.1.1.1") # 双重 release
|
|
assert lim._counts == {} # 不出负计数
|
|
assert lim.try_acquire("t", "1.1.1.1") # 状态仍可用
|
|
|
|
|
|
# ---- 接入: /ws/chat 握手期限流 ----------------------------------------------
|
|
# TestClient 的 WS scope client 恒为 ("testclient", 50000) 且测试不带 XFF,
|
|
# 故所有连接同落 ip="testclient";app_client 默认 jiancai 域——同租户同 IP,
|
|
# 正好是被限的键。默认 ws_conn_limit=10(defaults.py engine 组)。
|
|
|
|
|
|
def test_ws_over_limit_rejected_with_1013(app_client):
|
|
"""填满 10 个名额 → 第 11 条握手即被拒(1013, 未 accept)。"""
|
|
with ExitStack() as stack:
|
|
conns = [
|
|
stack.enter_context(app_client.websocket_connect("/ws/chat"))
|
|
for _ in range(10)
|
|
]
|
|
# 前 10 条真实可用(走完 init),不是恰好没崩
|
|
conns[0].send_text(json.dumps({"source": "standalone"}))
|
|
assert conns[0].receive_json()["type"] == "init"
|
|
|
|
with pytest.raises(WebSocketDisconnect) as ei:
|
|
with app_client.websocket_connect("/ws/chat"):
|
|
pass # pragma: no cover — 未 accept 即 close,进不来
|
|
assert ei.value.code == 1013
|
|
assert ei.value.reason == "too many connections"
|
|
|
|
|
|
def test_ws_conn_release_allows_reentry(app_client):
|
|
with ExitStack() as stack:
|
|
for _ in range(10):
|
|
stack.enter_context(app_client.websocket_connect("/ws/chat"))
|
|
# ExitStack 退出 = 10 条全部关闭,服务端 finally 已 release
|
|
with app_client.websocket_connect("/ws/chat") as ws:
|
|
ws.send_text(json.dumps({"source": "standalone"}))
|
|
assert ws.receive_json()["type"] == "init"
|