novalon-manage-system/api_integration_tests/tests/test_websocket.py

"""
WebSocket实时通信测试用例
"""

import pytest
import asyncio
import json
import os
from websockets.client import connect
from websockets.exceptions import ConnectionClosed


@pytest.mark.websocket
@pytest.mark.regression
class TestWebSocket:
    """WebSocket实时通信测试类"""
    
    @pytest.fixture
    def websocket_url(self):
        """WebSocket连接URL"""
        api_base_url = os.getenv("API_BASE_URL", "http://localhost:8084")
        ws_url = api_base_url.replace("http://", "ws://")
        return f"{ws_url}/ws"
    
    @pytest.fixture
    async def websocket_connection(self, websocket_url):
        """WebSocket连接fixture"""
        async with connect(websocket_url) as websocket:
            yield websocket
    
    @pytest.fixture
    async def authenticated_websocket(self, websocket_url, authenticated_client):
        """带认证的WebSocket连接"""
        token = authenticated_client.headers.get("Authorization")
        url_with_token = f"{websocket_url}?token={token.replace('Bearer ', '')}"
        
        async with connect(url_with_token) as websocket:
            yield websocket
    
    @pytest.mark.asyncio
    async def test_websocket_connection(self, websocket_url):
        """测试WebSocket连接建立"""
        try:
            async with connect(websocket_url) as websocket:
                assert websocket.open
        except ConnectionRefusedError:
            pytest.skip("WebSocket服务未启动")
    
    @pytest.mark.asyncio
    async def test_websocket_ping_pong(self, websocket_connection):
        """测试WebSocket心跳机制"""
        ping_message = {
            "type": "ping",
            "timestamp": 1234567890
        }
        
        await websocket_connection.send(json.dumps(ping_message))
        
        response = await asyncio.wait_for(websocket_connection.recv(), timeout=5.0)
        pong_message = json.loads(response)
        
        assert pong_message["type"] == "pong"
        assert "timestamp" in pong_message
    
    @pytest.mark.asyncio
    async def test_websocket_subscribe(self, websocket_connection):
        """测试WebSocket订阅"""
        subscribe_message = {
            "type": "subscribe",
            "channel": "notifications"
        }
        
        await websocket_connection.send(json.dumps(subscribe_message))
        
        response = await asyncio.wait_for(websocket_connection.recv(), timeout=5.0)
        message = json.loads(response)
        
        assert message["type"] in ["subscribe_ack", "pong"]
    
    @pytest.mark.asyncio
    async def test_websocket_multiple_messages(self, websocket_connection):
        """测试WebSocket多消息处理"""
        messages = [
            {"type": "ping"},
            {"type": "subscribe", "channel": "test"},
            {"type": "ping"}
        ]
        
        responses = []
        for msg in messages:
            await websocket_connection.send(json.dumps(msg))
            try:
                response = await asyncio.wait_for(websocket_connection.recv(), timeout=2.0)
                responses.append(json.loads(response))
            except asyncio.TimeoutError:
                pass
        
        assert len(responses) >= 2
    
    @pytest.mark.asyncio
    async def test_websocket_invalid_message(self, websocket_connection):
        """测试WebSocket无效消息处理"""
        invalid_messages = [
            "invalid json",
            "",
            json.dumps({"type": "unknown_type"}),
            json.dumps({})
        ]
        
        for msg in invalid_messages:
            try:
                await websocket_connection.send(msg)
                await asyncio.sleep(0.5)
            except Exception:
                pass
    
    @pytest.mark.asyncio
    async def test_websocket_connection_close(self, websocket_url):
        """测试WebSocket连接关闭"""
        async with connect(websocket_url) as websocket:
            assert websocket.open
            await websocket.close()
            assert not websocket.open
    
    @pytest.mark.asyncio
    async def test_websocket_timeout(self, websocket_url):
        """测试WebSocket超时"""
        try:
            async with connect(websocket_url, ping_timeout=2.0) as websocket:
                await asyncio.sleep(3.0)
        except (ConnectionClosed, asyncio.TimeoutError):
            pass
    
    @pytest.mark.asyncio
    async def test_websocket_concurrent_connections(self, websocket_url):
        """测试WebSocket并发连接"""
        async def create_connection():
            try:
                async with connect(websocket_url) as websocket:
                    await websocket.send(json.dumps({"type": "ping"}))
                    await asyncio.wait_for(websocket_connection.recv(), timeout=2.0)
            except Exception:
                pass
        
        connections = [create_connection() for _ in range(5)]
        await asyncio.gather(*connections, return_exceptions=True)
    
    @pytest.mark.asyncio
    async def test_websocket_large_message(self, websocket_connection):
        """测试WebSocket大消息处理"""
        large_data = "x" * 10000
        message = {
            "type": "test",
            "data": large_data
        }
        
        await websocket_connection.send(json.dumps(message))
        
        try:
            response = await asyncio.wait_for(websocket_connection.recv(), timeout=5.0)
            assert response
        except asyncio.TimeoutError:
            pass
    
    @pytest.mark.asyncio
    async def test_websocket_reconnect(self, websocket_url):
        """测试WebSocket重连"""
        for i in range(3):
            try:
                async with connect(websocket_url) as websocket:
                    await websocket.send(json.dumps({"type": "ping"}))
                    response = await asyncio.wait_for(websocket.recv(), timeout=2.0)
                    assert response
            except Exception:
                pass
    
    @pytest.mark.asyncio
    async def test_websocket_unicode_message(self, websocket_connection):
        """测试WebSocket Unicode消息"""
        unicode_message = {
            "type": "test",
            "content": "测试中文🎉🚀"
        }
        
        await websocket_connection.send(json.dumps(unicode_message))
        
        try:
            response = await asyncio.wait_for(websocket_connection.recv(), timeout=2.0)
            assert response
        except asyncio.TimeoutError:
            pass