fix(test): Implement robust WebSocket echo server with proper error handling and logging

Signed-off-by: Xin Liu <[email protected]>

Author: apepkuss

.github/workflows/test.yml

@@ -87,17 +87,40 @@ jobs:
 
       - name: Start WebSocket echo server
         run: |
-          # Use simple Python WebSocket echo server
+          # Use robust Python WebSocket echo server
           cat > ws-echo.py << 'EOF'
           import asyncio
           import websockets
+          import logging
+
+          logging.basicConfig(level=logging.INFO)
+          logger = logging.getLogger(__name__)
 
           async def echo(websocket, path):
-              async for message in websocket:
-                  await websocket.send(message)
+              logger.info(f"New WebSocket connection from {websocket.remote_address}")
+              try:
+                  async for message in websocket:
+                      # Echo back the message
+                      await websocket.send(message)
+                      logger.info(f"Echoed message: {len(message)} bytes")
+              except websockets.exceptions.ConnectionClosed:
+                  logger.info("Connection closed normally")
+              except Exception as e:
+                  logger.error(f"Error in echo handler: {e}")
+              finally:
+                  logger.info("Connection ended")
 
           async def main():
-              async with websockets.serve(echo, "0.0.0.0", 8890):
+              # Configure server with proper settings
+              async with websockets.serve(
+                  echo,
+                  "0.0.0.0",
+                  8890,
+                  ping_interval=20,
+                  ping_timeout=20,
+                  close_timeout=10
+              ):
+                  logger.info("WebSocket echo server started on port 8890")
                   await asyncio.Future()
 
           if __name__ == "__main__":

docker-compose.test.yml

@@ -36,20 +36,22 @@ services:
       - test-network
 
   # WebSocket Echo 服务（替代 echo.websocket.org）
-  # 使用简单的 WebSocket echo 服务器
+  # 使用 Python websockets 库实现的健壮 echo 服务器
   ws-echo:
-    image: jmalloc/echo-server:latest
+    image: python:3.11-alpine
     container_name: ss-proxy-test-ws
     ports:
-      - "8890:8080"
-    environment:
-      - PORT=8080
+      - "8890:8890"
+    volumes:
+      - ./tests/mock-data/ws-echo.py:/app/ws-echo.py:ro
+    working_dir: /app
+    command: sh -c "pip install --no-cache-dir websockets && python ws-echo.py"
     healthcheck:
-      test: ["CMD", "wget", "--spider", "-q", "http://localhost:8080"]
+      test: ["CMD", "sh", "-c", "timeout 5 sh -c 'cat < /dev/null > /dev/tcp/localhost/8890' || exit 1"]
       interval: 10s
       timeout: 5s
       retries: 3
-      start_period: 5s
+      start_period: 10s
     networks:
       - test-network
 

tests/mock-data/ws-echo.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+"""
+Robust WebSocket Echo Server for Testing
+
+This server echoes back any messages it receives.
+It properly handles WebSocket control frames and connection lifecycle.
+"""
+
+import asyncio
+import logging
+import sys
+
+import websockets
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+async def echo(websocket, path):
+    """
+    Echo handler that receives messages and sends them back.
+
+    Args:
+        websocket: The WebSocket connection
+        path: The request path
+    """
+    client_addr = websocket.remote_address
+    logger.info(f"New WebSocket connection from {client_addr} (path: {path})")
+
+    try:
+        async for message in websocket:
+            # Determine message type
+            if isinstance(message, bytes):
+                msg_type = "binary"
+                msg_size = len(message)
+            else:
+                msg_type = "text"
+                msg_size = len(message)
+
+            # Echo back the message
+            await websocket.send(message)
+            logger.info(f"Echoed {msg_type} message: {msg_size} bytes from {client_addr}")
+
+    except websockets.exceptions.ConnectionClosedOK:
+        logger.info(f"Connection closed normally by {client_addr}")
+    except websockets.exceptions.ConnectionClosedError as e:
+        logger.warning(f"Connection closed with error from {client_addr}: {e}")
+    except Exception as e:
+        logger.error(f"Unexpected error in echo handler from {client_addr}: {e}", exc_info=True)
+    finally:
+        logger.info(f"Connection ended with {client_addr}")
+
+
+async def main():
+    """
+    Start the WebSocket echo server.
+    """
+    host = "0.0.0.0"
+    port = 8890
+
+    # Configure server with proper timeouts and settings
+    server = await websockets.serve(
+        echo,
+        host,
+        port,
+        # Ping settings to keep connection alive
+        ping_interval=20,  # Send ping every 20 seconds
+        ping_timeout=20,   # Wait 20 seconds for pong response
+        close_timeout=10,  # Wait 10 seconds for close handshake
+        # Max message size (10MB)
+        max_size=10 * 1024 * 1024,
+        # Max queue size for incoming messages
+        max_queue=32,
+    )
+
+    logger.info(f"WebSocket echo server started on {host}:{port}")
+    logger.info("Press Ctrl+C to stop")
+
+    # Keep the server running
+    try:
+        await asyncio.Future()
+    except KeyboardInterrupt:
+        logger.info("Shutting down server...")
+        server.close()
+        await server.wait_closed()
+        logger.info("Server stopped")
+
+
+if __name__ == "__main__":
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        logger.info("Server interrupted by user")
+        sys.exit(0)
+        sys.exit(0)