Тест long polling

app/chat/websocket.py

@@ -11,7 +11,7 @@ from app.unit_of_work import UnitOfWork
 from app.utils.auth import AuthService
 from app.chat.router import router
 from app.chat.shemas import SSendMessage, SDeleteMessage, SEditMessage, SPinMessage, SUnpinMessage
-from app.users.dependencies import get_current_user_ws, get_token, get_subprotocol_ws
+from app.users.dependencies import get_current_user_ws, get_token, get_subprotocol_ws, get_verificated_user
 from app.users.schemas import SUser
 
 
@@ -42,6 +42,7 @@ class ConnectionManager:
 			new_message = await self.message_methods[message["flag"]](uow, user_id, chat_id, message)
 			for websocket in self.active_connections[chat_id]:
 				await websocket.send_json(new_message)
+			await polling_manager.send(new_message)
 		except KeyError:
 			raise IncorrectDataException
 
@@ -155,38 +156,68 @@ async def chat_ws(
 		print(await websocket.recv())
 
 
-class LongPollingManager:
+class PollingManager:
 	def __init__(self):
 		self.waiters: list[asyncio.Future] = []
-		self.messages: list[str] = []
+		self.messages: list[dict] = []
+
+	async def poll(self):
+		future = asyncio.Future()
+		self.waiters.append(future)
+		try:
+			await future
+			message = self.messages[0]
+			return message
+		except asyncio.CancelledError:
+			self.waiters.remove(future)
+			raise HTTPException(status_code=status.HTTP_408_REQUEST_TIMEOUT, detail="Client disconnected")
+
+	async def prepare(
+			self,
+			uow: UnitOfWork,
+			user_id: int,
+			chat_id: int,
+			message: SSendMessage | SDeleteMessage | SEditMessage | SPinMessage | SUnpinMessage,
+	):
+		message = message.model_dump()
+		await manager.broadcast(uow=uow, user_id=user_id, chat_id=chat_id, message=message)
+
+		await self.send(message)
+
+	async def send(
+			self,
+			message: dict,
+	):
+		self.messages.append(message)
+		while self.waiters:
+			waiter = self.waiters.pop(0)
+			if not waiter.done():
+				waiter.set_result(None)
+		await asyncio.sleep(0)
+		self.messages.pop(0)
 
 
-long_polling_manager = LongPollingManager()
+polling_manager = PollingManager()
 
 
 @router.get(
 	"/poll",
 	status_code=status.HTTP_200_OK,
+	response_model=SSendMessage | SDeleteMessage | SEditMessage | SPinMessage | SUnpinMessage,
 )
 async def poll():
-	future = asyncio.Future()
-	long_polling_manager.waiters.append(future)
-	try:
-		await future
-		return {"message": long_polling_manager.messages.pop(0)}
-	except asyncio.CancelledError:
-		long_polling_manager.waiters.remove(future)
-		raise HTTPException(status_code=500, detail="Client disconnected")
+	return await polling_manager.poll()
 
 
 @router.post(
-	"/send",
+	"/send/{chat_id}",
 	status_code=status.HTTP_201_CREATED,
 	response_model=None,
 )
-async def send(message: str):
-	long_polling_manager.messages.append(message)
-	while long_polling_manager.waiters:
-		waiter = long_polling_manager.waiters.pop(0)
-		if not waiter.done():
-			waiter.set_result(None)
+async def send(
+		chat_id: int,
+		message: SSendMessage | SDeleteMessage | SEditMessage | SPinMessage | SUnpinMessage,
+		user: SUser = Depends(get_verificated_user),
+		uow=Depends(UnitOfWork),
+):
+	await polling_manager.prepare(uow=uow, user_id=user.id, chat_id=chat_id, message=message)

app/users/dependencies.py

@@ -43,7 +43,7 @@ async def get_current_user(token: str = Depends(get_token), uow=Depends(UnitOfWo
 
 
 async def get_verificated_user(user: SUser = Depends(get_current_user)) -> SUser:
-	if not user.role >= settings.VERIFICATED_USER:
+	if user.role < settings.VERIFICATED_USER:
 		raise UserMustConfirmEmailException
 	return user