From 6eb9d1e01d0cd7283ed3de22af623b5356f26488 Mon Sep 17 00:00:00 2001
From: Jan Bouwhuis <jbouwh@users.noreply.github.com>
Date: Tue, 18 Jun 2024 09:29:22 +0200
Subject: [PATCH] Gracefully disconnect MQTT entry if entry is reloaded
 (#119753)

---
 homeassistant/components/mqtt/__init__.py |  4 ++--
 homeassistant/components/mqtt/client.py   | 15 +++++++++++----
 tests/components/mqtt/test_init.py        |  4 ++++
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/mqtt/__init__.py b/homeassistant/components/mqtt/__init__.py
index ea520e88366..f057dab8bc4 100644
--- a/homeassistant/components/mqtt/__init__.py
+++ b/homeassistant/components/mqtt/__init__.py
@@ -535,8 +535,8 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     registry_hooks = mqtt_data.discovery_registry_hooks
     while registry_hooks:
         registry_hooks.popitem()[1]()
-    # Wait for all ACKs and stop the loop
-    await mqtt_client.async_disconnect()
+    # Wait for all ACKs, stop the loop and disconnect the client
+    await mqtt_client.async_disconnect(disconnect_paho_client=True)
 
     # Cleanup MQTT client availability
     hass.data.pop(DATA_MQTT_AVAILABLE, None)
diff --git a/homeassistant/components/mqtt/client.py b/homeassistant/components/mqtt/client.py
index ace2293e7a6..562fa230bca 100644
--- a/homeassistant/components/mqtt/client.py
+++ b/homeassistant/components/mqtt/client.py
@@ -803,8 +803,12 @@ class MQTT:
 
             await asyncio.sleep(RECONNECT_INTERVAL_SECONDS)
 
-    async def async_disconnect(self) -> None:
-        """Stop the MQTT client."""
+    async def async_disconnect(self, disconnect_paho_client: bool = False) -> None:
+        """Stop the MQTT client.
+
+        We only disconnect grafully if disconnect_paho_client is set, but not
+        when Home Assistant is shut down.
+        """
 
         # stop waiting for any pending subscriptions
         await self._subscribe_debouncer.async_cleanup()
@@ -824,7 +828,9 @@ class MQTT:
             self._should_reconnect = False
             self._async_cancel_reconnect()
             # We do not gracefully disconnect to ensure
-            # the broker publishes the will message
+            # the broker publishes the will message unless the entry is reloaded
+            if disconnect_paho_client:
+                self._mqttc.disconnect()
 
     @callback
     def async_restore_tracked_subscriptions(
@@ -1274,7 +1280,8 @@ class MQTT:
         self._async_connection_result(False)
         self.connected = False
         async_dispatcher_send(self.hass, MQTT_CONNECTION_STATE, False)
-        _LOGGER.warning(
+        _LOGGER.log(
+            logging.INFO if result_code == 0 else logging.DEBUG,
             "Disconnected from MQTT server %s:%s (%s)",
             self.conf[CONF_BROKER],
             self.conf.get(CONF_PORT, DEFAULT_PORT),
diff --git a/tests/components/mqtt/test_init.py b/tests/components/mqtt/test_init.py
index 144b2f9cf45..18310750558 100644
--- a/tests/components/mqtt/test_init.py
+++ b/tests/components/mqtt/test_init.py
@@ -4087,6 +4087,7 @@ async def test_link_config_entry(
 async def test_reload_config_entry(
     hass: HomeAssistant,
     mqtt_mock_entry: MqttMockHAClientGenerator,
+    caplog: pytest.LogCaptureFixture,
 ) -> None:
     """Test manual entities reloaded and set up correctly."""
     await mqtt_mock_entry()
@@ -4153,6 +4154,9 @@ async def test_reload_config_entry(
         assert await hass.config_entries.async_reload(entry.entry_id)
         assert entry.state is ConfigEntryState.LOADED
         await hass.async_block_till_done()
+    # Assert the MQTT client was connected gracefully
+    with caplog.at_level(logging.INFO):
+        assert "Disconnected from MQTT server mock-broker:1883" in caplog.text
 
     assert (state := hass.states.get("sensor.test_manual1")) is not None
     assert state.attributes["friendly_name"] == "test_manual1_updated"