143 lines
4.5 KiB
Python
143 lines
4.5 KiB
Python
import json
|
|
import asyncio
|
|
import pytest
|
|
from websockets import ConnectionClosed
|
|
|
|
pytestmark = pytest.mark.asyncio
|
|
|
|
|
|
# ============================================================
|
|
# 认证测试
|
|
# ============================================================
|
|
|
|
async def test_auth_success(raw_ws):
|
|
"""正确凭据应返回 auth_ok"""
|
|
msg = json.dumps({"type": "auth", "username": "admin", "password": "admin123"})
|
|
await raw_ws.send(msg)
|
|
resp = await raw_ws.recv()
|
|
data = json.loads(resp)
|
|
assert data["type"] == "auth_ok"
|
|
|
|
|
|
async def test_auth_wrong_password(raw_ws):
|
|
"""错误密码应返回 auth_err"""
|
|
msg = json.dumps({"type": "auth", "username": "admin", "password": "wrong"})
|
|
await raw_ws.send(msg)
|
|
resp = await raw_ws.recv()
|
|
data = json.loads(resp)
|
|
assert data["type"] == "auth_err"
|
|
|
|
|
|
async def test_auth_nonexistent_user(raw_ws):
|
|
"""不存在的用户名应返回 auth_err"""
|
|
msg = json.dumps({"type": "auth", "username": "no_such_user", "password": "123456"})
|
|
await raw_ws.send(msg)
|
|
resp = await raw_ws.recv()
|
|
data = json.loads(resp)
|
|
assert data["type"] == "auth_err"
|
|
|
|
|
|
async def test_auth_non_json(raw_ws):
|
|
"""首条消息不是 JSON，服务端应关闭连接"""
|
|
await raw_ws.send("not json")
|
|
with pytest.raises(ConnectionClosed):
|
|
await raw_ws.recv()
|
|
|
|
|
|
async def test_auth_missing_type(raw_ws):
|
|
"""缺少 type 字段，服务端应关闭连接"""
|
|
msg = json.dumps({"username": "admin", "password": "admin123"})
|
|
await raw_ws.send(msg)
|
|
with pytest.raises(ConnectionClosed):
|
|
await raw_ws.recv()
|
|
|
|
|
|
# ============================================================
|
|
# 消息回显测试
|
|
# ============================================================
|
|
|
|
async def test_echo_text(authenticated_ws):
|
|
"""发送文本消息，验证回显一致"""
|
|
await authenticated_ws.send("hello world")
|
|
resp = await authenticated_ws.recv()
|
|
assert resp == "hello world"
|
|
|
|
|
|
async def test_echo_binary(authenticated_ws):
|
|
"""发送二进制消息，验证回显一致"""
|
|
data = b"\x00\x01\x02\xff\xfe\xfd"
|
|
await authenticated_ws.send(data)
|
|
resp = await authenticated_ws.recv()
|
|
assert resp == data
|
|
|
|
|
|
async def test_echo_multiple_messages(authenticated_ws):
|
|
"""连续发送多条消息，逐一验证回显"""
|
|
messages = ["msg1", "msg2", "msg3", "hello", "世界"]
|
|
for m in messages:
|
|
await authenticated_ws.send(m)
|
|
resp = await authenticated_ws.recv()
|
|
assert resp == m
|
|
|
|
|
|
async def test_echo_json_message(authenticated_ws):
|
|
"""发送 JSON 消息，验证回显一致"""
|
|
payload = json.dumps({"cmd": "ping", "seq": 1, "data": "test"})
|
|
await authenticated_ws.send(payload)
|
|
resp = await authenticated_ws.recv()
|
|
assert resp == payload
|
|
|
|
|
|
async def test_echo_large_message(authenticated_ws):
|
|
"""发送大消息（64KB），验证回显一致"""
|
|
large_msg = "A" * 65536
|
|
await authenticated_ws.send(large_msg)
|
|
resp = await authenticated_ws.recv()
|
|
assert resp == large_msg
|
|
|
|
|
|
# ============================================================
|
|
# 并发测试
|
|
# ============================================================
|
|
|
|
async def test_concurrent_connections():
|
|
"""同时建立 5 个认证连接，各自发送消息验证回显"""
|
|
from websockets.asyncio.client import connect
|
|
import os
|
|
|
|
ws_url = os.environ.get("LMVPN_WS_URL", "ws://localhost:8080/ws")
|
|
|
|
async def one_session(uid):
|
|
async with connect(ws_url, ping_interval=None) as ws:
|
|
auth = json.dumps({"type": "auth", "username": "admin", "password": "admin123"})
|
|
await ws.send(auth)
|
|
resp = await ws.recv()
|
|
assert json.loads(resp)["type"] == "auth_ok"
|
|
|
|
msg = f"msg_from_{uid}"
|
|
await ws.send(msg)
|
|
resp = await ws.recv()
|
|
assert resp == msg
|
|
|
|
tasks = [one_session(i) for i in range(5)]
|
|
await asyncio.gather(*tasks)
|
|
|
|
|
|
# ============================================================
|
|
# 协议测试
|
|
# ============================================================
|
|
|
|
async def test_connection_close_after_auth_failure(raw_ws):
|
|
"""认证失败后服务端应关闭连接"""
|
|
msg = json.dumps({"type": "auth", "username": "admin", "password": "wrong"})
|
|
await raw_ws.send(msg)
|
|
resp = await raw_ws.recv()
|
|
assert json.loads(resp)["type"] == "auth_err"
|
|
with pytest.raises(ConnectionClosed):
|
|
await raw_ws.recv()
|
|
|
|
|
|
async def test_ping_pong(authenticated_ws):
|
|
"""发送 ping，服务端应在超时前回应 pong"""
|
|
await authenticated_ws.ping()
|