test: conn_limiter 差异覆盖与导入清理
Co-Authored-By: Claude Code <noreply@anthropic.com>
This commit is contained in:
@@ -6,7 +6,7 @@ from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|||||||
from app.engine.models import ConversationSource, ConversationStatus
|
from app.engine.models import ConversationSource, ConversationStatus
|
||||||
from app.engine.routes.runtime import get_runtime
|
from app.engine.routes.runtime import get_runtime
|
||||||
from app.engine.services.conn_limiter import DEFAULT_WS_CONN_LIMIT
|
from app.engine.services.conn_limiter import DEFAULT_WS_CONN_LIMIT
|
||||||
from app.engine.services.conn_limiter import limiter as conn_limiter
|
from app.engine.services.conn_limiter import limiter
|
||||||
from app.engine.services.conversation import ConversationManager
|
from app.engine.services.conversation import ConversationManager
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -27,7 +27,7 @@ async def ws_chat(websocket: WebSocket):
|
|||||||
|
|
||||||
limit = runtime.config.data["engine"].get(
|
limit = runtime.config.data["engine"].get(
|
||||||
"ws_conn_limit", DEFAULT_WS_CONN_LIMIT)
|
"ws_conn_limit", DEFAULT_WS_CONN_LIMIT)
|
||||||
if not conn_limiter.try_acquire(runtime.config.slug, client_ip, limit=limit):
|
if not limiter.try_acquire(runtime.config.slug, client_ip, limit=limit):
|
||||||
await websocket.close(code=1013, reason="too many connections")
|
await websocket.close(code=1013, reason="too many connections")
|
||||||
return
|
return
|
||||||
|
|
||||||
@@ -123,4 +123,4 @@ async def ws_chat(websocket: WebSocket):
|
|||||||
finally:
|
finally:
|
||||||
# Task 8: 正常关闭与异常断开都归还名额——acquire 成功后所有路径
|
# Task 8: 正常关闭与异常断开都归还名额——acquire 成功后所有路径
|
||||||
# 必经此处,不泄漏(accept 失败/循环中抛错同样覆盖)
|
# 必经此处,不泄漏(accept 失败/循环中抛错同样覆盖)
|
||||||
conn_limiter.release(runtime.config.slug, client_ip)
|
limiter.release(runtime.config.slug, client_ip)
|
||||||
|
|||||||
@@ -17,6 +17,7 @@ DEFAULT_CONFIG: dict = {
|
|||||||
"max_attempts": 3, "retry_backoff_base": 1.0},
|
"max_attempts": 3, "retry_backoff_base": 1.0},
|
||||||
"ai": {"provider": "glm", "api_key": None, "base_url": None, "model": "glm-4-flash"},
|
"ai": {"provider": "glm", "api_key": None, "base_url": None, "model": "glm-4-flash"},
|
||||||
"engine": {"context_seed_keys": {}, "rate_limit_per_hour": 30,
|
"engine": {"context_seed_keys": {}, "rate_limit_per_hour": 30,
|
||||||
|
# 双源防漂移: 默认值须与 conn_limiter.DEFAULT_WS_CONN_LIMIT 一致
|
||||||
"ws_conn_limit": 10},
|
"ws_conn_limit": 10},
|
||||||
"misc": {"name_to_region_extras": {}},
|
"misc": {"name_to_region_extras": {}},
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -46,6 +46,18 @@ def test_release_unheld_key_is_safe_noop():
|
|||||||
assert lim.try_acquire("t", "1.1.1.1") # 状态仍可用
|
assert lim.try_acquire("t", "1.1.1.1") # 状态仍可用
|
||||||
|
|
||||||
|
|
||||||
|
def test_per_call_limit_overrides_constructor():
|
||||||
|
# 按次 cap 覆盖构造值(chat.py 传租户配置 limit 必须真生效,封住
|
||||||
|
# "try_acquire 静默忽略 limit 参数"的盲区); 不传 limit 的键仍受构造值管束
|
||||||
|
lim = ConnectionLimiter(limit=2)
|
||||||
|
for _ in range(5):
|
||||||
|
assert lim.try_acquire("t", "1.1.1.1", limit=5) # 第 3-5 个也能进
|
||||||
|
assert not lim.try_acquire("t", "1.1.1.1", limit=5) # 第 6 个才拒绝
|
||||||
|
assert lim.try_acquire("t", "2.2.2.2") # 不传 limit 的键: 第 1、2 个进
|
||||||
|
assert lim.try_acquire("t", "2.2.2.2")
|
||||||
|
assert not lim.try_acquire("t", "2.2.2.2") # 仍按构造值 2 拒第 3 个
|
||||||
|
|
||||||
|
|
||||||
# ---- 接入: /ws/chat 握手期限流 ----------------------------------------------
|
# ---- 接入: /ws/chat 握手期限流 ----------------------------------------------
|
||||||
# TestClient 的 WS scope client 恒为 ("testclient", 50000) 且测试不带 XFF,
|
# TestClient 的 WS scope client 恒为 ("testclient", 50000) 且测试不带 XFF,
|
||||||
# 故所有连接同落 ip="testclient";app_client 默认 jiancai 域——同租户同 IP,
|
# 故所有连接同落 ip="testclient";app_client 默认 jiancai 域——同租户同 IP,
|
||||||
|
|||||||
Reference in New Issue
Block a user