sanguo_vnpy_v2/tests/verify_websocket.py

#!/usr/bin/env python3
"""
WebSocket 验证脚本
验证 WebSocket 模块的基本结构和功能
"""
import sys
import os

# 添加项目路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../vnpy_v4.4.0'))


def test_imports():
    """测试模块导入"""
    print("Testing imports...")

    try:
        from sanguo_web.websocket.manager import ConnectionManager, manager
        print("  ✓ ConnectionManager imported")

        from sanguo_web.websocket.routes import router
        print("  ✓ WebSocket router imported")

        from sanguo_web.websocket.events import (
            EventMonitorManager,
            serialize_tick_data,
            serialize_order_data,
            serialize_trade_data
        )
        print("  ✓ Event monitors imported")

        from sanguo_web.websocket import (
            ConnectionManager as CM,
            manager as mgr,
            router as ws_router,
            EventMonitorManager as EMM
        )
        print("  ✓ WebSocket module __init__ exports")

        return True
    except Exception as e:
        print(f"  ✗ Import failed: {e}")
        return False


def test_manager():
    """测试连接管理器"""
    print("\nTesting ConnectionManager...")

    try:
        from sanguo_web.websocket.manager import ConnectionManager

        manager = ConnectionManager()

        # 测试连接计数
        count = manager.get_connection_count()
        assert count == 0, f"Expected 0 connections, got {count}"
        print("  ✓ Connection count initialized correctly")

        # 测试用户连接计数
        user_count = manager.get_user_connection_count("test_user")
        assert user_count == 0, f"Expected 0 user connections, got {user_count}"
        print("  ✓ User connection count initialized correctly")

        return True
    except Exception as e:
        print(f"  ✗ Manager test failed: {e}")
        return False


def test_serialization():
    """测试数据序列化函数"""
    print("\nTesting serialization functions...")

    try:
        from sanguo_web.websocket.events import serialize_datetime
        from datetime import datetime

        # 测试 datetime 序列化
        dt = datetime(2025, 1, 1, 12, 30, 45)
        serialized = serialize_datetime(dt)
        assert serialized == "2025-01-01T12:30:45", f"Expected ISO format, got {serialized}"
        print("  ✓ datetime serialization works")

        # 测试 None 处理
        none_result = serialize_datetime(None)
        assert none_result is None, f"Expected None for None input, got {none_result}"
        print("  ✓ None datetime handled correctly")

        return True
    except Exception as e:
        print(f"  ✗ Serialization test failed: {e}")
        return False


def test_router():
    """测试路由"""
    print("\nTesting WebSocket router...")

    try:
        from sanguo_web.websocket.routes import router

        # 检查路由
        routes = [route.path for route in router.routes]
        assert "/ws" in routes, "WebSocket route not found"
        print("  ✓ WebSocket /ws route exists")

        assert "/ws/status" in routes, "WebSocket status route not found"
        print("  ✓ WebSocket /ws/status route exists")

        return True
    except Exception as e:
        print(f"  ✗ Router test failed: {e}")
        return False


def test_event_monitor_classes():
    """测试事件监听器类"""
    print("\nTesting event monitor classes...")

    try:
        from sanguo_web.websocket.events import (
            TickEventMonitor,
            OrderEventMonitor,
            TradeEventMonitor,
            PositionEventMonitor,
            AccountEventMonitor,
            LogEventMonitor,
            ContractEventMonitor,
            EventMonitorManager
        )

        # 检查类是否具有必要的方法
        monitors = [
            TickEventMonitor,
            OrderEventMonitor,
            TradeEventMonitor,
            PositionEventMonitor,
            AccountEventMonitor,
            LogEventMonitor,
            ContractEventMonitor
        ]

        for monitor_class in monitors:
            # 检查是否有 stop 方法
            assert hasattr(monitor_class, 'stop'), f"{monitor_class.__name__} missing stop method"
            print(f"  ✓ {monitor_class.__name__} has stop method")

        # 检查 EventMonitorManager
        assert hasattr(EventMonitorManager, 'start_all'), "EventMonitorManager missing start_all method"
        assert hasattr(EventMonitorManager, 'stop_all'), "EventMonitorManager missing stop_all method"
        print("  ✓ EventMonitorManager has start_all and stop_all methods")

        return True
    except Exception as e:
        print(f"  ✗ Event monitor test failed: {e}")
        return False


def main():
    """主测试函数"""
    print("="*50)
    print("WebSocket Phase 3 Verification")
    print("="*50)

    tests = [
        ("Import Test", test_imports),
        ("Manager Test", test_manager),
        ("Serialization Test", test_serialization),
        ("Router Test", test_router),
        ("Event Monitor Test", test_event_monitor_classes),
    ]

    results = []

    for name, test_func in tests:
        print(f"\n{name}:")
        try:
            result = test_func()
            results.append((name, result))
        except Exception as e:
            print(f"  ✗ Test crashed: {e}")
            results.append((name, False))

    # 汇总结果
    print("\n" + "="*50)
    print("Test Results")
    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")

    if passed == total:
        print("\n✓ All tests passed! Phase 3 implementation verified.")
        return 0
    else:
        print(f"\n✗ {total - passed} test(s) failed")
        return 1


if __name__ == "__main__":
    sys.exit(main())