Fix MQTT last will on graceful disconnect

homeassistant/components/mqtt/__init__.py

@@ -450,6 +450,7 @@ def __init__(self, hass, broker, port, client_id, keepalive, username,
         self.subscribed_topics = {}
         self.progress = {}
         self.birth_message = birth_message
+        self.will_message = will_message
         self._mqttc = None
         self._paho_lock = asyncio.Lock(loop=hass.loop)
 
@@ -490,7 +491,7 @@ def __init__(self, hass, broker, port, client_id, keepalive, username,
     def async_publish(self, topic, payload, qos, retain):
         """Publish a MQTT message.
 
-        This method must be run in the event loop and returns a coroutine.
+        This method is a coroutine.
         """
         with (yield from self._paho_lock):
             yield from self.hass.async_add_job(
@@ -513,18 +514,23 @@ def async_connect(self):
 
         return not result
 
+    @asyncio.coroutine
     def async_disconnect(self):
         """Stop the MQTT client.
 
-        This method must be run in the event loop and returns a coroutine.
+        This method is a coroutine.
         """
-        def stop():
-            """Stop the MQTT client."""
+        if self.will_message:
+            yield from self.async_publish(
+                self.will_message.get(ATTR_TOPIC),
+                self.will_message.get(ATTR_PAYLOAD),
+                self.will_message.get(ATTR_QOS),
+                self.will_message.get(ATTR_RETAIN))
+
+        with (yield from self._paho_lock):
             self._mqttc.disconnect()
             self._mqttc.loop_stop()
 
-        return self.hass.async_add_job(stop)
-
     @asyncio.coroutine
     def async_subscribe(self, topic, qos):
         """Subscribe to a topic.

tests/components/mqtt/test_init.py

@@ -553,9 +553,26 @@ def test_birth_message(hass):
     mqtt_client.publish = lambda *args: calls.append(args)
     hass.data['mqtt']._mqtt_on_connect(None, None, 0, 0)
     yield from hass.async_block_till_done()
+    assert len(calls) == 1
     assert calls[-1] == ('birth', 'birth', 0, False)
 
 
+@asyncio.coroutine
+def test_last_will_message_on_graceful_disconnect(hass):
+    """Test sending last will message on graceful disconnect."""
+    mqtt_client = yield from mock_mqtt_client(hass, {
+        mqtt.CONF_BROKER: 'mock-broker',
+        mqtt.CONF_WILL_MESSAGE: {mqtt.ATTR_TOPIC: 'will',
+                                 mqtt.ATTR_PAYLOAD: 'will'}
+    })
+    calls = []
+    mqtt_client.publish = lambda *args: calls.append(args)
+    yield from hass.data['mqtt'].async_disconnect()
+    yield from hass.async_block_till_done()
+    assert len(calls) == 1
+    assert calls[-1] == ('will', 'will', 0, False)
+
+
 @asyncio.coroutine
 def test_mqtt_subscribes_topics_on_connect(hass):
     """Test subscription to topic on connect."""