Add appropriate guards to websocket actions (#292)

Author: bachya

simplipy/websocket.py

@@ -116,8 +116,9 @@ def __init__(
 
     def _on_expire(self) -> None:
         """Log and act when the watchdog expires."""
-        LOGGER.info("Websocket watchdog expired")
-        schedule_callback(self._action)
+        if self._timer_task and not self._timer_task.cancelled():
+            LOGGER.info("Websocket watchdog expired")
+            schedule_callback(self._action)
 
     def cancel(self) -> None:
         """Cancel the watchdog."""
@@ -314,6 +315,9 @@ def add_event_callback(self, callback: Callable[..., Any]) -> Callable[..., None
 
     async def async_connect(self) -> None:
         """Connect to the websocket server."""
+        if self.connected:
+            return
+
         try:
             self._client = await self._api.session.ws_connect(
                 WEBSOCKET_SERVER_URL, heartbeat=55
@@ -330,15 +334,15 @@ async def async_connect(self) -> None:
 
     async def async_disconnect(self) -> None:
         """Disconnect from the websocket server."""
+        if not self.connected:
+            return
+
         assert self._client
 
         await self._client.close()
 
         LOGGER.info("Disconnected from websocket server")
 
-        for callback in self._disconnect_callbacks:
-            schedule_callback(callback)
-
     async def async_listen(self) -> None:
         """Start listening to the websocket server."""
         assert self._client

tests/conftest.py

@@ -226,8 +226,6 @@ async def receive():
 
     ws_client.receive.side_effect = receive
 
-    ws_client.send_json.side_effect = Mock()
-
     async def reset_close():
         """Reset the websocket client close method."""
         ws_client.closed = True

tests/test_websocket.py

@@ -33,6 +33,46 @@
 from .common import create_ws_message
 
 
+@pytest.mark.asyncio
+async def test_callbacks(caplog, mock_api, ws_message_event, ws_messages):
+    """Test that callbacks are executed correctly."""
+    caplog.set_level(logging.INFO)
+
+    mock_connect_callback = Mock()
+    mock_disconnect_callback = Mock()
+    mock_event_callback = Mock()
+
+    async def async_mock_connect_callback():
+        """Define a mock async conenct callback."""
+        LOGGER.info("We are connected!")
+
+    client = WebsocketClient(mock_api)
+    client.add_connect_callback(mock_connect_callback)
+    client.add_connect_callback(async_mock_connect_callback)
+    client.add_disconnect_callback(mock_disconnect_callback)
+    client.add_event_callback(mock_event_callback)
+
+    assert mock_connect_callback.call_count == 0
+    assert mock_disconnect_callback.call_count == 0
+    assert mock_event_callback.call_count == 0
+
+    await client.async_connect()
+    assert client.connected
+    await asyncio.sleep(1)
+    assert mock_connect_callback.call_count == 1
+    assert any("We are connected!" in e.message for e in caplog.records)
+
+    ws_messages.append(create_ws_message(ws_message_event))
+    await client.async_listen()
+    await asyncio.sleep(1)
+    expected_event = websocket_event_from_payload(ws_message_event)
+    mock_event_callback.assert_called_once_with(expected_event)
+
+    await client.async_disconnect()
+    assert not client.connected
+    assert mock_disconnect_callback.call_count == 1
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "error",
@@ -61,6 +101,9 @@ async def test_connect_disconnect(mock_api):
     await client.async_connect()
     assert client.connected
 
+    # Attempt to connect again, which should just return:
+    await client.async_connect()
+
     await client.async_disconnect()
     assert not client.connected
 
@@ -159,48 +202,6 @@ async def test_listen_error_message_types(
         await client.async_listen()
 
 
-@pytest.mark.asyncio
-async def test_callbacks(caplog, mock_api, ws_message_event, ws_messages):
-    """Test that callbacks are executed correctly."""
-    caplog.set_level(logging.INFO)
-
-    mock_connect_callback = Mock()
-    mock_disconnect_callback = Mock()
-    mock_event_callback = Mock()
-
-    async def async_mock_connect_callback():
-        """Define a mock async conenct callback."""
-        LOGGER.info("We are connected!")
-
-    client = WebsocketClient(mock_api)
-    client.add_connect_callback(mock_connect_callback)
-    client.add_connect_callback(async_mock_connect_callback)
-    client.add_disconnect_callback(mock_disconnect_callback)
-    client.add_event_callback(mock_event_callback)
-
-    assert mock_connect_callback.call_count == 0
-    assert mock_disconnect_callback.call_count == 0
-    assert mock_event_callback.call_count == 0
-
-    await client.async_connect()
-    assert client.connected
-    await asyncio.sleep(1)
-    assert mock_connect_callback.call_count == 1
-    assert any("We are connected!" in e.message for e in caplog.records)
-
-    ws_messages.append(create_ws_message(ws_message_event))
-    await client.async_listen()
-    await asyncio.sleep(1)
-    expected_event = websocket_event_from_payload(ws_message_event)
-    mock_event_callback.assert_called_once_with(expected_event)
-
-    # Add another message to the queue to keep the websocket open while we disconnect
-    # from it:
-    await client.async_disconnect()
-    assert not client.connected
-    assert mock_disconnect_callback.call_count == 1
-
-
 @pytest.mark.asyncio
 async def test_reconnect(mock_api):
     """Test reconnecting to the websocket."""