claude_dev
918bbed0fc
- 修复 deps.py 中 get_vn_service 的引用错误 (vn_service.vn_service -> vn_service) - 移除登录页面上的明文密码提示 - 改进前端错误处理，避免数据加载失败导致登录显示错误
417 lines
11 KiB
Python
417 lines
11 KiB
Python
"""
|
|
WebSocket 客户端测试脚本
|
|
用于测试 WebSocket 实时数据推送功能
|
|
"""
|
|
import asyncio
|
|
import websockets
|
|
import json
|
|
from typing import Optional
|
|
|
|
|
|
class WebSocketTestClient:
|
|
"""WebSocket 测试客户端"""
|
|
|
|
def __init__(self, url: str = "ws://localhost:8000/ws"):
|
|
"""
|
|
初始化测试客户端
|
|
|
|
- **url**: WebSocket 服务器 URL
|
|
"""
|
|
self.url = url
|
|
self.websocket: Optional[websockets.WebSocketClientProtocol] = None
|
|
self.connected = False
|
|
|
|
async def connect(self, token: Optional[str] = None):
|
|
"""
|
|
连接到 WebSocket 服务器
|
|
|
|
- **token**: 可选的 JWT Token
|
|
"""
|
|
uri = f"{self.url}"
|
|
if token:
|
|
uri += f"?token={token}"
|
|
|
|
try:
|
|
self.websocket = await websockets.connect(uri)
|
|
self.connected = True
|
|
print(f"[+] Connected to {self.url}")
|
|
|
|
# 接收欢迎消息
|
|
welcome_msg = await self.websocket.recv()
|
|
print(f"[+] Welcome message: {welcome_msg}")
|
|
return json.loads(welcome_msg)
|
|
except Exception as e:
|
|
print(f"[-] Failed to connect: {e}")
|
|
raise
|
|
|
|
async def disconnect(self):
|
|
"""断开连接"""
|
|
if self.websocket:
|
|
await self.websocket.close()
|
|
self.connected = False
|
|
print("[+] Disconnected from server")
|
|
|
|
async def subscribe(self, subscription_types: list):
|
|
"""
|
|
订阅数据类型
|
|
|
|
- **subscription_types**: 订阅类型列表，如 ["tick", "order", "trade"]
|
|
"""
|
|
message = {
|
|
"type": "subscribe",
|
|
"data": {
|
|
"subscription": subscription_types
|
|
}
|
|
}
|
|
await self.send(message)
|
|
print(f"[+] Subscribed to: {subscription_types}")
|
|
|
|
async def unsubscribe(self, subscription_types: list):
|
|
"""
|
|
取消订阅
|
|
|
|
- **subscription_types**: 订阅类型列表
|
|
"""
|
|
message = {
|
|
"type": "unsubscribe",
|
|
"data": {
|
|
"subscription": subscription_types
|
|
}
|
|
}
|
|
await self.send(message)
|
|
print(f"[+] Unsubscribed from: {subscription_types}")
|
|
|
|
async def subscribe_symbol(self, symbols: list):
|
|
"""
|
|
订阅品种行情
|
|
|
|
- **symbols**: 品种代码列表，如 ["IF2501.CFFEX", "IH2501.CFFEX"]
|
|
"""
|
|
message = {
|
|
"type": "subscribe_symbol",
|
|
"data": {
|
|
"symbol": symbols
|
|
}
|
|
}
|
|
await self.send(message)
|
|
print(f"[+] Subscribed to symbols: {symbols}")
|
|
|
|
async def unsubscribe_symbol(self, symbols: list):
|
|
"""
|
|
取消品种订阅
|
|
|
|
- **symbols**: 品种代码列表
|
|
"""
|
|
message = {
|
|
"type": "unsubscribe_symbol",
|
|
"data": {
|
|
"symbol": symbols
|
|
}
|
|
}
|
|
await self.send(message)
|
|
print(f"[+] Unsubscribed from symbols: {symbols}")
|
|
|
|
async def send_ping(self):
|
|
"""发送心跳"""
|
|
message = {
|
|
"type": "ping",
|
|
"data": {}
|
|
}
|
|
await self.send(message)
|
|
print("[+] Ping sent")
|
|
|
|
async def send(self, message: dict):
|
|
"""
|
|
发送消息
|
|
|
|
- **message**: 消息字典
|
|
"""
|
|
if not self.websocket or not self.connected:
|
|
raise Exception("Not connected to WebSocket server")
|
|
|
|
await self.websocket.send(json.dumps(message))
|
|
|
|
async def receive(self, timeout: Optional[float] = None):
|
|
"""
|
|
接收消息
|
|
|
|
- **timeout**: 超时时间（秒）
|
|
"""
|
|
if not self.websocket or not self.connected:
|
|
raise Exception("Not connected to WebSocket server")
|
|
|
|
try:
|
|
message = await asyncio.wait_for(self.websocket.recv(), timeout=timeout)
|
|
return json.loads(message)
|
|
except asyncio.TimeoutError:
|
|
return None
|
|
|
|
async def listen(self, duration: int = 10, print_messages: bool = True):
|
|
"""
|
|
监听消息
|
|
|
|
- **duration**: 监听时长（秒）
|
|
- **print_messages**: 是否打印消息
|
|
"""
|
|
print(f"\n[*] Listening for messages ({duration}s)...")
|
|
messages = []
|
|
|
|
try:
|
|
while True:
|
|
message = await asyncio.wait_for(self.websocket.recv(), timeout=duration)
|
|
data = json.loads(message)
|
|
messages.append(data)
|
|
|
|
if print_messages:
|
|
msg_type = data.get("type", "unknown")
|
|
print(f"[*] Received {msg_type}: {json.dumps(data, ensure_ascii=False)[:200]}...")
|
|
|
|
except asyncio.TimeoutError:
|
|
print(f"\n[+] Listening finished. Received {len(messages)} messages")
|
|
return messages
|
|
|
|
|
|
async def test_basic_connection():
|
|
"""测试基本连接功能"""
|
|
print("\n" + "="*50)
|
|
print("TEST 1: Basic Connection")
|
|
print("="*50)
|
|
|
|
client = WebSocketTestClient()
|
|
|
|
try:
|
|
# 连接
|
|
await client.connect()
|
|
await asyncio.sleep(1)
|
|
|
|
# 测试心跳
|
|
await client.send_ping()
|
|
response = await client.receive(timeout=2)
|
|
if response and response.get("type") == "pong":
|
|
print("[+] Ping/Pong test passed")
|
|
|
|
# 断开
|
|
await client.disconnect()
|
|
print("[+] Test passed")
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"[-] Test failed: {e}")
|
|
return False
|
|
|
|
|
|
async def test_subscription():
|
|
"""测试订阅功能"""
|
|
print("\n" + "="*50)
|
|
print("TEST 2: Subscription")
|
|
print("="*50)
|
|
|
|
client = WebSocketTestClient()
|
|
|
|
try:
|
|
# 连接
|
|
await client.connect()
|
|
|
|
# 订阅多个数据类型
|
|
await client.subscribe(["tick", "order", "trade", "position", "account", "log"])
|
|
await asyncio.sleep(0.5)
|
|
|
|
# 订阅品种
|
|
await client.subscribe_symbol(["IF2501.CFFEX", "IH2501.CFFEX"])
|
|
await asyncio.sleep(0.5)
|
|
|
|
# 取消订阅
|
|
await client.unsubscribe(["log"])
|
|
await asyncio.sleep(0.5)
|
|
|
|
# 断开
|
|
await client.disconnect()
|
|
print("[+] Test passed")
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"[-] Test failed: {e}")
|
|
return False
|
|
|
|
|
|
async def test_message_reception():
|
|
"""测试消息接收功能"""
|
|
print("\n" + "="*50)
|
|
print("TEST 3: Message Reception")
|
|
print("="*50)
|
|
|
|
client = WebSocketTestClient()
|
|
|
|
try:
|
|
# 连接
|
|
await client.connect()
|
|
|
|
# 订阅所有类型
|
|
await client.subscribe(["tick", "order", "trade", "position", "account", "log", "contract"])
|
|
|
|
# 监听消息
|
|
messages = await client.listen(duration=5)
|
|
|
|
print(f"\n[+] Received messages by type:")
|
|
msg_types = {}
|
|
for msg in messages:
|
|
msg_type = msg.get("type", "unknown")
|
|
msg_types[msg_type] = msg_types.get(msg_type, 0) + 1
|
|
|
|
for msg_type, count in msg_types.items():
|
|
print(f" - {msg_type}: {count}")
|
|
|
|
# 断开
|
|
await client.disconnect()
|
|
print("[+] Test passed")
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"[-] Test failed: {e}")
|
|
return False
|
|
|
|
|
|
async def test_heartbeat():
|
|
"""测试心跳机制"""
|
|
print("\n" + "="*50)
|
|
print("TEST 4: Heartbeat")
|
|
print("="*50)
|
|
|
|
client = WebSocketTestClient()
|
|
|
|
try:
|
|
# 连接
|
|
await client.connect()
|
|
|
|
# 发送多次心跳
|
|
for i in range(5):
|
|
await client.send_ping()
|
|
response = await client.receive(timeout=2)
|
|
if response and response.get("type") == "pong":
|
|
print(f"[+] Heartbeat {i+1}/5 successful")
|
|
await asyncio.sleep(1)
|
|
|
|
# 断开
|
|
await client.disconnect()
|
|
print("[+] Test passed")
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"[-] Test failed: {e}")
|
|
return False
|
|
|
|
|
|
async def test_error_handling():
|
|
"""测试错误处理"""
|
|
print("\n" + "="*50)
|
|
print("TEST 5: Error Handling")
|
|
print("="*50)
|
|
|
|
client = WebSocketTestClient()
|
|
|
|
try:
|
|
# 连接
|
|
await client.connect()
|
|
|
|
# 发送未知消息类型
|
|
await client.send({
|
|
"type": "unknown_type",
|
|
"data": {}
|
|
})
|
|
|
|
response = await client.receive(timeout=2)
|
|
if response and response.get("type") == "error":
|
|
print("[+] Error response received correctly")
|
|
|
|
# 断开
|
|
await client.disconnect()
|
|
print("[+] Test passed")
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"[-] Test failed: {e}")
|
|
return False
|
|
|
|
|
|
async def test_authenticated_connection():
|
|
"""测试认证连接"""
|
|
print("\n" + "="*50)
|
|
print("TEST 6: Authenticated Connection")
|
|
print("="*50)
|
|
|
|
# 注意：需要有效的 JWT Token
|
|
# 这里测试无效 Token 的情况
|
|
client = WebSocketTestClient()
|
|
|
|
try:
|
|
# 使用无效 Token 连接
|
|
await client.connect(token="invalid_token")
|
|
print("[-] Should have failed with invalid token")
|
|
await client.disconnect()
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"[+] Correctly rejected invalid token: {e}")
|
|
return True
|
|
|
|
|
|
async def run_all_tests():
|
|
"""运行所有测试"""
|
|
print("\n" + "="*50)
|
|
print("WebSocket Test Suite")
|
|
print("="*50)
|
|
|
|
tests = [
|
|
("Basic Connection", test_basic_connection),
|
|
("Subscription", test_subscription),
|
|
("Message Reception", test_message_reception),
|
|
("Heartbeat", test_heartbeat),
|
|
("Error Handling", test_error_handling),
|
|
("Authenticated Connection", test_authenticated_connection),
|
|
]
|
|
|
|
results = []
|
|
|
|
for name, test_func in tests:
|
|
try:
|
|
result = await test_func()
|
|
results.append((name, result))
|
|
except Exception as e:
|
|
print(f"[-] Test '{name}' crashed: {e}")
|
|
results.append((name, False))
|
|
|
|
await asyncio.sleep(1)
|
|
|
|
# 汇总结果
|
|
print("\n" + "="*50)
|
|
print("Test Results Summary")
|
|
print("="*50)
|
|
|
|
passed = sum(1 for _, result in results if result)
|
|
total = len(results)
|
|
|
|
for name, result in results:
|
|
status = "✓ PASS" if result else "✗ FAIL"
|
|
print(f"{status}: {name}")
|
|
|
|
print(f"\nTotal: {passed}/{total} tests passed")
|
|
|
|
return passed == total
|
|
|
|
|
|
def main():
|
|
"""主函数"""
|
|
# 运行测试
|
|
success = asyncio.run(run_all_tests())
|
|
|
|
if success:
|
|
print("\n[+] All tests passed!")
|
|
exit(0)
|
|
else:
|
|
print("\n[-] Some tests failed")
|
|
exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|