From 8f98fb2ec44c586c65d6d6e74708512f46fca5cf Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sat, 29 Jun 2024 11:06:56 -0500
Subject: [PATCH] Fix publish cancellation handling in MQTT (#120826)

---
 homeassistant/components/mqtt/client.py |  4 ++--
 tests/components/mqtt/test_client.py    | 17 +++++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/homeassistant/components/mqtt/client.py b/homeassistant/components/mqtt/client.py
index 7788c1db641..f65769badfa 100644
--- a/homeassistant/components/mqtt/client.py
+++ b/homeassistant/components/mqtt/client.py
@@ -1141,8 +1141,8 @@ class MQTT:
         # see https://github.com/eclipse/paho.mqtt.python/issues/687
         # properties and reason codes are not used in Home Assistant
         future = self._async_get_mid_future(mid)
-        if future.done() and future.exception():
-            # Timed out
+        if future.done() and (future.cancelled() or future.exception()):
+            # Timed out or cancelled
             return
         future.set_result(None)
 
diff --git a/tests/components/mqtt/test_client.py b/tests/components/mqtt/test_client.py
index 49b590383d1..cd02d805e1c 100644
--- a/tests/components/mqtt/test_client.py
+++ b/tests/components/mqtt/test_client.py
@@ -1194,6 +1194,23 @@ async def test_handle_mqtt_on_callback(
         assert "No ACK from MQTT server" not in caplog.text
 
 
+async def test_handle_mqtt_on_callback_after_cancellation(
+    hass: HomeAssistant,
+    caplog: pytest.LogCaptureFixture,
+    mqtt_mock_entry: MqttMockHAClientGenerator,
+    mqtt_client_mock: MqttMockPahoClient,
+) -> None:
+    """Test receiving an ACK after a cancellation."""
+    mqtt_mock = await mqtt_mock_entry()
+    # Simulate the mid future getting a cancellation
+    mqtt_mock()._async_get_mid_future(101).cancel()
+    # Simulate an ACK for mid == 101, being received after the cancellation
+    mqtt_client_mock.on_publish(mqtt_client_mock, None, 101)
+    await hass.async_block_till_done()
+    assert "No ACK from MQTT server" not in caplog.text
+    assert "InvalidStateError" not in caplog.text
+
+
 async def test_handle_mqtt_on_callback_after_timeout(
     hass: HomeAssistant,
     caplog: pytest.LogCaptureFixture,