Версия 0.0.2

SRV/api/v1/websocket.py

@@ -0,0 +1,90 @@
+import asyncio
+import logging
+from typing import Set
+
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+from starlette.websockets import WebSocketState
+
+from .network import update_container_position, get_network_data
+
+router = APIRouter()
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class WebSocketManager:
+    """
+    Класс для управления подключениями к WebSocket-серверу
+    """
+    def __init__(self):
+        """ Инициализация менеджера подключений к WebSocket-серверу """
+        self.active_connections: Set[WebSocket] = set()
+
+    async def connect(self, websocket: WebSocket):
+        """ Подключение клиента к WebSocket-серверу """
+        await websocket.accept()
+        self.active_connections.add(websocket)
+
+    def disconnect(self, websocket: WebSocket):
+        """ Отключение клиента от WebSocket-серверу """
+        self.active_connections.remove(websocket)
+
+    async def broadcast(self, message: dict):
+        """ Трансляция сообщений всем подключенным клиентам """
+        for connection in list(self.active_connections):
+            try:
+                if connection.client_state == WebSocketState.CONNECTED:
+                    await connection.send_json(message)
+            except WebSocketDisconnect:
+                self.disconnect(connection)
+            except Exception as e:
+                logger.error(f"Сообщение об ошибке трансляции: {e}")
+                self.disconnect(connection)
+
+manager = WebSocketManager()
+
+@router.websocket("/api/v1")
+async def websocket_endpoint(websocket: WebSocket):
+    """ Обработчик WebSocket-соединения """
+    await manager.connect(websocket)
+    try:
+        while True:
+            data = await websocket.receive_json()
+            logger.info(f"Полученные данные: {data}")
+
+            if all(key in data for key in ["container_id", "container_x", "container_y"]):
+                await update_container_position(data)
+
+            network_data = await get_network_data()
+            logger.info(f"Текущие сетевые данные: {network_data}")
+
+            await manager.broadcast(network_data)
+            await asyncio.sleep(2)
+    except WebSocketDisconnect:
+        logger.info("Клиент отключен")
+    except Exception as e:
+        logger.error(f"Ошибка в websocket_endpoint:{e}")
+    finally:
+        manager.disconnect(websocket)
+        if websocket.client_state != WebSocketState.DISCONNECTED:
+            await websocket.close()
+
+@router.on_event("startup")
+async def startup_event():
+    """ Запуск таска для трансляции сетевых данных """
+    asyncio.create_task(broadcast_data())
+
+async def broadcast_data():
+    """ Таск для трансляции сетевых данных """
+    while True:
+        try:
+            network_data = await get_network_data()
+            if "activation" not in network_data:
+                logger.error("Ошибка: Ключ 'activation' отсутствует в полученных данных сети")
+                await asyncio.sleep(2)
+                continue
+
+            await manager.broadcast(network_data)
+        except Exception as e:
+            logger.error(f"Ошибка в broadcast_data:{e}")
+        await asyncio.sleep(2)