From 126163cb2f73b56cf05d23dd083d40c9684a0cc0 Mon Sep 17 00:00:00 2001 From: glukhov324 Date: Sat, 4 Oct 2025 17:33:24 +0700 Subject: [PATCH] add websockets --- hw2/hw/client_websockets.py | 68 ++++++++++++++++++ hw2/hw/requirements.txt | 3 +- hw2/hw/shop_api/chat_manager/__init__.py | 1 + hw2/hw/shop_api/chat_manager/chat_manager.py | 73 ++++++++++++++++++++ hw2/hw/shop_api/main.py | 5 +- hw2/hw/shop_api/routers/__init__.py | 3 +- hw2/hw/shop_api/routers/chat.py | 22 ++++++ hw2/hw/test_websockets.py | 36 ++++++++++ 8 files changed, 207 insertions(+), 4 deletions(-) create mode 100644 hw2/hw/client_websockets.py create mode 100644 hw2/hw/shop_api/chat_manager/__init__.py create mode 100644 hw2/hw/shop_api/chat_manager/chat_manager.py create mode 100644 hw2/hw/shop_api/routers/chat.py create mode 100644 hw2/hw/test_websockets.py diff --git a/hw2/hw/client_websockets.py b/hw2/hw/client_websockets.py new file mode 100644 index 00000000..016d972a --- /dev/null +++ b/hw2/hw/client_websockets.py @@ -0,0 +1,68 @@ +import asyncio +import websockets + + +DEFAULT_HOST = "localhost" +DEFAULT_PORT = 8000 + + +async def chat_client(room_name: str, host: str = DEFAULT_HOST, port: int = DEFAULT_PORT): + uri = f"ws://{host}:{port}/chat/{room_name}" + print(f"Подключение к комнате '{room_name}' по адресу {uri}...") + + try: + async with websockets.connect(uri) as websocket: + + welcome = await websocket.recv() + print(f"{welcome}\n") + print("Чат запущен! Введите сообщение и нажмите Enter.") + print("Чтобы выйти — введите 'quit' или нажмите Ctrl+C.\n") + + async def receive_messages(): + try: + while True: + message = await websocket.recv() + print(f"\r{message}") + print("Вы: ", end="", flush=True) + except websockets.exceptions.ConnectionClosed: + print("\nСоединение закрыто сервером.") + return + + async def send_messages(): + loop = asyncio.get_event_loop() + while True: + + msg = await loop.run_in_executor(None, input, "Вы: ") + if msg.strip().lower() == 'quit': + break + if msg.strip(): + await websocket.send(msg.strip()) + + await asyncio.gather(receive_messages(), send_messages()) + + except KeyboardInterrupt: + print("\nВыход по запросу пользователя.") + except Exception as e: + print(f"\nОшибка подключения: {e}") + print("Убедитесь, что сервер запущен: uvicorn main:app") + + +def main(): + import argparse + + parser = argparse.ArgumentParser(description="WebSocket чат-клиент") + parser.add_argument("room", help="Имя комнаты для подключения") + parser.add_argument("--host", default=DEFAULT_HOST, help="Хост сервера (по умолчанию: localhost)") + parser.add_argument("--port", type=int, default=DEFAULT_PORT, help="Порт сервера (по умолчанию: 8000)") + + args = parser.parse_args() + + try: + asyncio.run(chat_client(args.room, args.host, args.port)) + except KeyboardInterrupt: + + print("\nДо встречи!") + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/hw2/hw/requirements.txt b/hw2/hw/requirements.txt index 207dcf5c..fae256c2 100644 --- a/hw2/hw/requirements.txt +++ b/hw2/hw/requirements.txt @@ -1,9 +1,10 @@ # Основные зависимости для ASGI приложения fastapi>=0.117.1 uvicorn>=0.24.0 +websockets>=1.8.0 # Зависимости для тестирования pytest>=7.4.0 pytest-asyncio>=0.21.0 httpx>=0.27.2 -Faker>=37.8.0 +Faker>=37.8.0 \ No newline at end of file diff --git a/hw2/hw/shop_api/chat_manager/__init__.py b/hw2/hw/shop_api/chat_manager/__init__.py new file mode 100644 index 00000000..5fcbda19 --- /dev/null +++ b/hw2/hw/shop_api/chat_manager/__init__.py @@ -0,0 +1 @@ +from shop_api.chat_manager.chat_manager import chat_manager \ No newline at end of file diff --git a/hw2/hw/shop_api/chat_manager/chat_manager.py b/hw2/hw/shop_api/chat_manager/chat_manager.py new file mode 100644 index 00000000..a73bba87 --- /dev/null +++ b/hw2/hw/shop_api/chat_manager/chat_manager.py @@ -0,0 +1,73 @@ +import uuid +from typing import List, Dict +from fastapi import APIRouter, WebSocket, WebSocketDisconnect + + + +class ChatManager: + def __init__(self): + + self.rooms: Dict[str, List[WebSocket]] = {} + self.user_data: Dict[WebSocket, tuple[str, str]] = {} + + def generate_username(self) -> str: + return f"user-{uuid.uuid4().hex}" + + async def connect(self, websocket: WebSocket, chat_name: str) -> str: + + await websocket.accept() + username = self.generate_username() + + if chat_name not in self.rooms: + self.rooms[chat_name] = [] + + self.rooms[chat_name].append(websocket) + self.user_data[websocket] = (username, chat_name) + + await websocket.send_text(f"Вы подключены как: {username}") + return username + + + async def disconnect(self, websocket: WebSocket): + + if websocket not in self.user_data: + return + + username, chat_name = self.user_data[websocket] + del self.user_data[websocket] + + if chat_name in self.rooms: + if websocket in self.rooms[chat_name]: + self.rooms[chat_name].remove(websocket) + # Удаляем пустую комнату + if not self.rooms[chat_name]: + del self.rooms[chat_name] + + + async def publish(self, websocket: WebSocket, message: str): + + if websocket not in self.user_data: + return + + username, chat_name = self.user_data[websocket] + full_message = f"{username} :: {message}" + + if chat_name not in self.rooms: + return + + disconnected = [] + for client in self.rooms[chat_name]: + try: + await client.send_text(full_message) + except Exception: + disconnected.append(client) + + for client in disconnected: + if client in self.rooms[chat_name]: + self.rooms[chat_name].remove(client) + if client in self.user_data: + del self.user_data[client] + + + +chat_manager = ChatManager() \ No newline at end of file diff --git a/hw2/hw/shop_api/main.py b/hw2/hw/shop_api/main.py index 7c801dc5..53442985 100644 --- a/hw2/hw/shop_api/main.py +++ b/hw2/hw/shop_api/main.py @@ -1,7 +1,8 @@ from fastapi import FastAPI -from shop_api.routers import cart_router, item_router +from shop_api.routers import cart_router, item_router, chat_router app = FastAPI(title="Shop API") app.include_router(cart_router) -app.include_router(item_router) \ No newline at end of file +app.include_router(item_router) +app.include_router(chat_router) \ No newline at end of file diff --git a/hw2/hw/shop_api/routers/__init__.py b/hw2/hw/shop_api/routers/__init__.py index ddae3a8b..5bc6a52f 100644 --- a/hw2/hw/shop_api/routers/__init__.py +++ b/hw2/hw/shop_api/routers/__init__.py @@ -1,2 +1,3 @@ from shop_api.routers.cart import router as cart_router -from shop_api.routers.item import router as item_router \ No newline at end of file +from shop_api.routers.item import router as item_router +from shop_api.routers.chat import router as chat_router \ No newline at end of file diff --git a/hw2/hw/shop_api/routers/chat.py b/hw2/hw/shop_api/routers/chat.py new file mode 100644 index 00000000..796cff91 --- /dev/null +++ b/hw2/hw/shop_api/routers/chat.py @@ -0,0 +1,22 @@ +from fastapi import APIRouter, WebSocket, WebSocketDisconnect +from shop_api.chat_manager import chat_manager + + + +router = APIRouter( + prefix="/chat", + tags=["Chat"] +) + + +@router.websocket("/{chat_name}") +async def websocket_endpoint(websocket: WebSocket, chat_name: str): + try: + await chat_manager.connect(websocket, chat_name) + while True: + data = await websocket.receive_text() + await chat_manager.publish(websocket, data) + except WebSocketDisconnect: + pass + finally: + await chat_manager.disconnect(websocket) \ No newline at end of file diff --git a/hw2/hw/test_websockets.py b/hw2/hw/test_websockets.py new file mode 100644 index 00000000..b2a9c66f --- /dev/null +++ b/hw2/hw/test_websockets.py @@ -0,0 +1,36 @@ +import pytest +from fastapi.testclient import TestClient +from shop_api.main import app + + +client = TestClient(app) + + +def test_websocket_connection(): + + with client.websocket_connect("/chat/testroom") as websocket: + + data = websocket.receive_text() + assert data.startswith("Вы подключены как: user-") + + websocket.send_text("Привет из теста!") + + response = websocket.receive_text() + assert " :: Привет из теста!" in response + assert response.startswith("user-") + + +def test_message_format_with_extracted_username(): + with client.websocket_connect("/chat/test") as ws: + + welcome = ws.receive_text() + assert welcome.startswith("Вы подключены как: ") + username = welcome.replace("Вы подключены как: ", "").strip() + + test_msg = "Привет, это тест!" + ws.send_text(test_msg) + + response = ws.receive_text() + + expected = f"{username} :: {test_msg}" + assert response == expected, f"Ожидалось '{expected}', получено '{response}'" \ No newline at end of file