From 4e316429d31df34d7e1c7eb69a368a4522fff315 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sat, 21 Dec 2024 00:18:47 -1000
Subject: [PATCH] Handle WebsocketConnectionError during mqtt auto reconnect
 (#133697)

followup to #133610 to handle the exception in the auto reconnect
path as well

fixes #132985
---
 homeassistant/components/mqtt/client.py |  5 ++++-
 tests/components/mqtt/test_client.py    | 10 +++++++++-
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/mqtt/client.py b/homeassistant/components/mqtt/client.py
index 73c6b80cb14..6500c9f91c9 100644
--- a/homeassistant/components/mqtt/client.py
+++ b/homeassistant/components/mqtt/client.py
@@ -695,12 +695,15 @@ class MQTT:
 
     async def _reconnect_loop(self) -> None:
         """Reconnect to the MQTT server."""
+        # pylint: disable-next=import-outside-toplevel
+        import paho.mqtt.client as mqtt
+
         while True:
             if not self.connected:
                 try:
                     async with self._connection_lock, self._async_connect_in_executor():
                         await self.hass.async_add_executor_job(self._mqttc.reconnect)
-                except OSError as err:
+                except (OSError, mqtt.WebsocketConnectionError) as err:
                     _LOGGER.debug(
                         "Error re-connecting to MQTT server due to exception: %s", err
                     )
diff --git a/tests/components/mqtt/test_client.py b/tests/components/mqtt/test_client.py
index 1878045a9b9..1daad0e3914 100644
--- a/tests/components/mqtt/test_client.py
+++ b/tests/components/mqtt/test_client.py
@@ -1888,10 +1888,18 @@ async def test_mqtt_subscribes_and_unsubscribes_in_chunks(
     assert len(mqtt_client_mock.unsubscribe.mock_calls[1][1][0]) == 2
 
 
+@pytest.mark.parametrize(
+    "exception",
+    [
+        OSError,
+        paho_mqtt.WebsocketConnectionError,
+    ],
+)
 async def test_auto_reconnect(
     hass: HomeAssistant,
     setup_with_birth_msg_client_mock: MqttMockPahoClient,
     caplog: pytest.LogCaptureFixture,
+    exception: Exception,
 ) -> None:
     """Test reconnection is automatically done."""
     mqtt_client_mock = setup_with_birth_msg_client_mock
@@ -1902,7 +1910,7 @@ async def test_auto_reconnect(
     mqtt_client_mock.on_disconnect(None, None, 0)
     await hass.async_block_till_done()
 
-    mqtt_client_mock.reconnect.side_effect = OSError("foo")
+    mqtt_client_mock.reconnect.side_effect = exception("foo")
     async_fire_time_changed(
         hass, utcnow() + timedelta(seconds=RECONNECT_INTERVAL_SECONDS)
     )