mirror of
https://github.com/andwati/telecom-platform.git
synced 2026-10-18 19:12:54 +00:00
nutcas3
758bd5f075
- Add pyjwt>=2.10.0 dependency to requirements.txt - Add AuthProvider class with API key and JWT token support (create_token, validate_token, refresh_token, get_headers methods) - Add APIClient base class with retry logic, error handling, and async session management - Add SubscriberAPI module (get, list, create, update, suspend, activate, terminate methods) - Add UsageAPI module (get_stats, list_events, get_realtime methods
101 lines
3.4 KiB
Python
101 lines
3.4 KiB
Python
import json
|
|
import logging
|
|
from typing import Optional, Callable
|
|
import websockets
|
|
from websockets.exceptions import ConnectionClosed
|
|
|
|
from .types import WebSocketMessage
|
|
from .exceptions import NetworkError
|
|
|
|
|
|
class WebSocketClient:
|
|
"""WebSocket client for real-time updates from the Telecom Platform."""
|
|
|
|
def __init__(
|
|
self,
|
|
api_url: str,
|
|
auth_provider,
|
|
enable_logging: bool = False,
|
|
):
|
|
"""
|
|
Initialize the WebSocket client.
|
|
|
|
Args:
|
|
api_url: Base URL for the API server
|
|
auth_provider: Authentication provider instance
|
|
enable_logging: Enable debug logging
|
|
"""
|
|
self.api_url = api_url.rstrip('/')
|
|
self.auth_provider = auth_provider
|
|
self.enable_logging = enable_logging
|
|
|
|
if enable_logging:
|
|
self.logger = logging.getLogger(__name__)
|
|
else:
|
|
self.logger = logging.getLogger(__name__)
|
|
self.logger.disabled = True
|
|
|
|
self._websocket: Optional[websockets.WebSocketServerProtocol] = None
|
|
self._connected = False
|
|
|
|
async def connect(self, message_handler: Callable[[WebSocketMessage], None]) -> None:
|
|
"""
|
|
Connect to WebSocket for real-time updates.
|
|
|
|
Args:
|
|
message_handler: Async function to handle WebSocket messages
|
|
"""
|
|
ws_url = self.api_url.replace("http://", "ws://") + "/ws"
|
|
headers = self.auth_provider.get_headers()
|
|
|
|
try:
|
|
self._websocket = await websockets.connect(ws_url, extra_headers=headers)
|
|
self._connected = True
|
|
self.logger.info("WebSocket connected")
|
|
|
|
async for message in self._websocket:
|
|
try:
|
|
data = json.loads(message)
|
|
ws_message = WebSocketMessage(**data)
|
|
await message_handler(ws_message)
|
|
except json.JSONDecodeError as e:
|
|
self.logger.error(f"Failed to decode WebSocket message: {e}")
|
|
except Exception as e:
|
|
self.logger.error(f"Error handling WebSocket message: {e}")
|
|
|
|
except ConnectionClosed:
|
|
self.logger.warning("WebSocket connection closed")
|
|
self._connected = False
|
|
except Exception as e:
|
|
self.logger.error(f"WebSocket error: {e}")
|
|
self._connected = False
|
|
raise NetworkError(f"WebSocket connection failed: {e}")
|
|
|
|
async def disconnect(self) -> None:
|
|
"""Disconnect from WebSocket."""
|
|
if self._websocket:
|
|
await self._websocket.close()
|
|
self._websocket = None
|
|
self._connected = False
|
|
self.logger.info("WebSocket disconnected")
|
|
|
|
def is_connected(self) -> bool:
|
|
"""Check if WebSocket is connected."""
|
|
return self._connected
|
|
|
|
async def send(self, data: dict) -> None:
|
|
"""
|
|
Send data through the WebSocket connection.
|
|
|
|
Args:
|
|
data: Dictionary data to send
|
|
"""
|
|
if not self._connected or not self._websocket:
|
|
raise NetworkError("WebSocket is not connected")
|
|
|
|
try:
|
|
await self._websocket.send(json.dumps(data))
|
|
except Exception as e:
|
|
self.logger.error(f"Failed to send WebSocket message: {e}")
|
|
raise NetworkError(f"Failed to send message: {e}")
|