71 lines
2.4 KiB
Python
71 lines
2.4 KiB
Python
"""WebSocket keyword library for Robot Framework tests."""
|
|
|
|
import asyncio
|
|
from typing import Optional
|
|
|
|
import websockets
|
|
from robot.api import logger
|
|
|
|
|
|
class WebSocketKeywords:
|
|
"""Keywords for WebSocket testing."""
|
|
|
|
ROBOT_AUTO_KEYWORDS = True
|
|
|
|
def __init__(self):
|
|
self._connection: Optional[websockets.WebSocketClientProtocol] = None
|
|
|
|
def connect_to_websocket(self, ws_url: str) -> str:
|
|
"""Establish a WebSocket connection and return connection info.
|
|
|
|
Args:
|
|
ws_url: Full WebSocket URL including token query param.
|
|
|
|
Returns:
|
|
String representation of the connected WebSocket.
|
|
"""
|
|
asyncio.run(self._async_connect(ws_url))
|
|
return f"Connected to {ws_url}"
|
|
|
|
connect_to_websocket.robot_name = "Connect To WebSocket"
|
|
|
|
async def _async_connect(self, ws_url: str) -> None:
|
|
self._connection = await websockets.connect(ws_url)
|
|
logger.info(f"WebSocket connected: {self._connection}")
|
|
|
|
def send_websocket_message(self, message: str) -> None:
|
|
"""Send a text message over the active WebSocket connection.
|
|
|
|
Args:
|
|
message: Text message to send.
|
|
"""
|
|
if self._connection is None:
|
|
raise RuntimeError("No active WebSocket connection. Use Connect To WebSocket first.")
|
|
asyncio.run(self._connection.send(message))
|
|
logger.info(f"Sent message: {message}")
|
|
|
|
def receive_websocket_message(self, timeout: float = 5.0) -> str:
|
|
"""Receive a text message from the active WebSocket connection.
|
|
|
|
Args:
|
|
timeout: Maximum seconds to wait for a message.
|
|
|
|
Returns:
|
|
The received text message.
|
|
"""
|
|
if self._connection is None:
|
|
raise RuntimeError("No active WebSocket connection. Use Connect To WebSocket first.")
|
|
try:
|
|
message = asyncio.run(asyncio.wait_for(self._connection.recv(), timeout=timeout))
|
|
logger.info(f"Received message: {message}")
|
|
return message
|
|
except asyncio.TimeoutError:
|
|
raise TimeoutError(f"No message received within {timeout} seconds")
|
|
|
|
def close_websocket(self) -> None:
|
|
"""Close the active WebSocket connection."""
|
|
if self._connection is not None:
|
|
asyncio.run(self._connection.close())
|
|
self._connection = None
|
|
logger.info("WebSocket closed")
|