test improvements

realtime/connection.py

@@ -18,7 +18,6 @@
 from realtime.exceptions import NotConnectedError
 from realtime.message import HEARTBEAT_PAYLOAD, PHOENIX_CHANNEL, ChannelEvents, Message
 
-
 T_Retval = TypeVar("T_Retval")
 T_ParamSpec = ParamSpec("T_ParamSpec")
 
@@ -40,12 +39,12 @@ def wrapper(*args: T_ParamSpec.args, **kwargs: T_ParamSpec.kwargs) -> T_Retval:
 
 class Socket:
     def __init__(
-        self,
-        url: str,
-        auto_reconnect: bool = False,
-        params: Dict[str, Any] = {},
-        hb_interval: int = 30,
-        version: int = 2,
+            self,
+            url: str,
+            auto_reconnect: bool = False,
+            params: Dict[str, Any] = {},
+            hb_interval: int = 30,
+            version: int = 2,
     ) -> None:
         """
         `Socket` is the abstraction for an actual socket connection that receives and 'reroutes' `Message` according to its `topic` and `event`.
@@ -68,6 +67,51 @@ def __init__(
 
         self.channels: DefaultDict[str, List[Channel]] = defaultdict(list)
 
+    async def _process_message(self, msg: str):
+        try:
+            if self.version == 1:
+                msg = Message(**json.loads(msg))
+            elif self.version == 2:
+                msg_array = json.loads(msg)
+                msg = Message(
+                    join_ref=msg_array[0],
+                    ref=msg_array[1],
+                    topic=msg_array[2],
+                    event=msg_array[3],
+                    payload=msg_array[4],
+                )
+            if msg.event == ChannelEvents.reply:
+                for channel in self.channels.get(msg.topic, []):
+                    if msg.ref == channel.control_msg_ref:
+                        if msg.payload["status"] == "error":
+                            logging.info(
+                                f"Error joining channel: {msg.topic} - {msg.payload['response']['reason']}"
+                            )
+                            break
+                        elif msg.payload["status"] == "ok":
+                            logging.info(f"Successfully joined {msg.topic}")
+                            continue
+                    else:
+                        for cl in channel.listeners:
+                            if cl.ref in ["*", msg.ref]:
+                                cl.callback(msg.payload)
+
+            if msg.event == ChannelEvents.close:
+                for channel in self.channels.get(msg.topic, []):
+                    if msg.join_ref == channel.join_ref:
+                        logging.info(f"Successfully left {msg.topic}")
+                        continue
+
+            for channel in self.channels.get(msg.topic, []):
+                for cl in channel.listeners:
+                    if cl.event in ["*", msg.event]:
+                        if asyncio.iscoroutinefunction(cl.callback):
+                            await cl.callback(msg.payload)
+                        else:
+                            cl.callback(msg.payload)
+        except Exception as e:
+            logging.error(f"Error processing message: {e}", exc_info=True)
+
     @ensure_connection
     async def listen(self) -> None:
         """
@@ -79,47 +123,7 @@ async def listen(self) -> None:
         while True:
             try:
                 msg = await self.ws_connection.recv()
-                if self.version == 1:
-                    msg = Message(**json.loads(msg))
-                elif self.version == 2:
-                    msg_array = json.loads(msg)
-                    msg = Message(
-                        join_ref=msg_array[0],
-                        ref=msg_array[1],
-                        topic=msg_array[2],
-                        event=msg_array[3],
-                        payload=msg_array[4],
-                    )
-                if msg.event == ChannelEvents.reply:
-                    for channel in self.channels.get(msg.topic, []):
-                        if msg.ref == channel.control_msg_ref:
-                            if msg.payload["status"] == "error":
-                                logging.info(
-                                    f"Error joining channel: {msg.topic} - {msg.payload['response']['reason']}"
-                                )
-                                break
-                            elif msg.payload["status"] == "ok":
-                                logging.info(f"Successfully joined {msg.topic}")
-                                continue
-                        else:
-                            for cl in channel.listeners:
-                                if cl.ref in ["*", msg.ref]:
-                                    cl.callback(msg.payload)
-
-                if msg.event == ChannelEvents.close:
-                    for channel in self.channels.get(msg.topic, []):
-                        if msg.join_ref == channel.join_ref:
-                            logging.info(f"Successfully left {msg.topic}")
-                            continue
-
-                for channel in self.channels.get(msg.topic, []):
-                    for cl in channel.listeners:
-                        if cl.event in ["*", msg.event]:
-                            if asyncio.iscoroutinefunction(cl.callback):
-                                asyncio.create_task(cl.callback(msg.payload))
-                            else:
-                                cl.callback(msg.payload)
-
+                asyncio.create_task(self._process_message(msg))
 
             except ConnectionClosedOK:
                 logging.info("Connection was closed normally.")
@@ -145,7 +149,7 @@ async def listen(self) -> None:
                 await self.leave_all()
 
             except (
-                Exception
+                    Exception
             ) as e:  # A general exception handler should be the last resort
                 logging.error(f"Unexpected error in listen: {e}")
                 await self._handle_reconnection()
@@ -209,7 +213,7 @@ async def keep_alive(self) -> None:
                 await self._handle_reconnection()
 
             except (
-                Exception
+                    Exception
             ) as e:  # A general exception handler should be the last resort
                 logging.error(f"Unexpected error in keep_alive: {e}")