diff --git a/flanabot/bots/btc_offers_bot.py b/flanabot/bots/btc_offers_bot.py index 95c3e55..a46a16c 100644 --- a/flanabot/bots/btc_offers_bot.py +++ b/flanabot/bots/btc_offers_bot.py @@ -11,8 +11,8 @@ from collections.abc import Callable import aiohttp import flanautils +import websockets from multibot import MultiBot, constants as multibot_constants -from websockets.asyncio import client import constants from flanabot.models import Chat, Message @@ -97,7 +97,7 @@ def preprocess_btc_offers(func: Callable) -> Callable: class BtcOffersBot(MultiBot, ABC): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self._websocket: client.ClientConnection | None = None + self._websocket: websockets.ClientConnection | None = None self._notification_task: asyncio.Task[None] | None = None self._api_endpoint = f"{os.environ['BTC_OFFERS_API_HOST']}:{os.environ['BTC_OFFERS_API_PORT']}/offers" @@ -135,7 +135,9 @@ class BtcOffersBot(MultiBot, ABC): if offer['author']: offer_parts.append(f"Autor: {offer['author']}") - payment_methods_text = ''.join(f'\n {payment_method}' for payment_method in offer['payment_methods']) + payment_methods_text = ''.join( + f'\n {payment_method}' for payment_method in offer['payment_methods'] + ) offer_parts.extend( ( @@ -148,7 +150,9 @@ class BtcOffersBot(MultiBot, ABC): ) if offer['description']: - offer_parts.append(f"Descripción:\n{offer['description']}") + offer_parts.append( + f"Descripción:\n{offer['description']}" + ) offers_parts.append('\n'.join(offer_parts)) @@ -175,7 +179,7 @@ class BtcOffersBot(MultiBot, ABC): ( '', '-' * 70, - 'ℹ️ Los avisos de ofertas de BTC se han desactivado. Si quieres volver a recibirlos, no dudes en pedírmelo.' + 'ℹ️ Los avisos de ofertas de BTC se han eliminado. Si quieres volver a recibirlos, no dudes en pedírmelo.' ) ) @@ -184,7 +188,14 @@ class BtcOffersBot(MultiBot, ABC): async def _wait_btc_offers_notification(self): while True: - data = json.loads(await self._websocket.recv()) + while True: + try: + data = json.loads(await self._websocket.recv()) + except websockets.ConnectionClosed: + await self.start_saved_btc_offers_notification() + else: + break + chat = await self.get_chat(data['chat_id']) chat.btc_offers_query = {} chat.save(pull_exclude_fields=('btc_offers_query',)) @@ -227,14 +238,7 @@ class BtcOffersBot(MultiBot, ABC): async def _on_ready(self): await super()._on_ready() - - for chat in self.Chat.find({ - 'platform': self.platform.value, - 'btc_offers_max_eur': {'$exists': True, '$ne': None} - }): - chat = await self.get_chat(chat.id) - chat.pull_from_database(overwrite_fields=('_id', 'btc_offers_max_eur')) - await self.start_btc_offers_notification(chat, chat.btc_offers_max_eur) + await self.start_saved_btc_offers_notification() async def _on_stop_btc_offers_notification(self, message: Message): previous_btc_offers_query = message.chat.btc_offers_query @@ -249,14 +253,31 @@ class BtcOffersBot(MultiBot, ABC): # -------------------------------------------------------- # # -------------------- PUBLIC METHODS -------------------- # # -------------------------------------------------------- # - async def start_btc_offers_notification(self, chat: Chat, max_price_eur: float): - if not self._websocket: - self._websocket = await client.connect(f'ws://{self._api_endpoint}') + async def start_btc_offers_notification(self, chat: Chat, query: dict[str, float]): + if not self._websocket or self._websocket.state in {websockets.State.CLOSING, websockets.State.CLOSED}: + while True: + try: + self._websocket = await websockets.connect(f'ws://{self._api_endpoint}') + except ConnectionRefusedError: + await asyncio.sleep(constants.BTC_OFFERS_WEBSOCKET_RETRY_DELAY_SECONDS) + else: + break + + if not self._notification_task: self._notification_task = asyncio.create_task(self._wait_btc_offers_notification()) chat.btc_offers_query = query chat.save() - await self._websocket.send(json.dumps({'action': 'start', 'chat_id': chat.id, 'max_price_eur': max_price_eur})) + await self._websocket.send(json.dumps({'action': 'start', 'chat_id': chat.id, 'query': query})) + + async def start_saved_btc_offers_notification(self): + for chat in self.Chat.find({ + 'platform': self.platform.value, + 'btc_offers_query': {'$exists': True, '$ne': {}} + }): + chat = await self.get_chat(chat.id) + chat.pull_from_database(overwrite_fields=('_id', 'btc_offers_query')) + await self.start_btc_offers_notification(chat, chat.btc_offers_query) async def stop_btc_offers_notification(self, chat: Chat): if not self._websocket: diff --git a/flanabot/constants.py b/flanabot/constants.py index 903ec96..25a9b5f 100644 --- a/flanabot/constants.py +++ b/flanabot/constants.py @@ -6,6 +6,7 @@ from multibot import Platform AUDIT_LOG_AGE = datetime.timedelta(hours=1) AUDIT_LOG_LIMIT = 5 AUTO_WEATHER_EVERY = datetime.timedelta(hours=6) +BTC_OFFERS_WEBSOCKET_RETRY_DELAY_SECONDS = datetime.timedelta(hours=1).total_seconds() CHECK_PUNISHMENTS_EVERY_SECONDS = datetime.timedelta(hours=1).total_seconds() CONNECT_4_AI_DELAY_SECONDS = 1 CONNECT_4_CENTER_COLUMN_POINTS = 2