Fix MQTT reconnecting

homeassistant/components/mqtt/__init__.py

@@ -51,7 +51,7 @@ MAX_RECONNECT_WAIT = 300  # seconds
 
 
 def publish(hass, topic, payload, qos=None, retain=None):
-    """ Send an MQTT message. """
+    """Publish message to an MQTT topic."""
     data = {
         ATTR_TOPIC: topic,
         ATTR_PAYLOAD: payload,
@@ -66,7 +66,7 @@ def publish(hass, topic, payload, qos=None, retain=None):
 
 
 def subscribe(hass, topic, callback, qos=DEFAULT_QOS):
-    """ Subscribe to a topic. """
+    """Subscribe to an MQTT topic."""
     def mqtt_topic_subscriber(event):
         """Match subscribed MQTT topic."""
         if _match_topic(topic, event.data[ATTR_TOPIC]):
@@ -78,8 +78,7 @@ def subscribe(hass, topic, callback, qos=DEFAULT_QOS):
 
 
 def setup(hass, config):
-    """ Get the MQTT protocol service. """
+    """Start the MQTT protocol service."""
-
     if not validate_config(config, {DOMAIN: ['broker']}, _LOGGER):
         return False
 
@@ -137,34 +136,32 @@ def setup(hass, config):
 
 # pylint: disable=too-many-arguments
 class MQTT(object):
-    """ Implements messaging service for MQTT. """
+    """Home Assistant MQTT client."""
+
     def __init__(self, hass, broker, port, client_id, keepalive, username,
                  password, certificate):
+        """Initialize Home Assistant MQTT client."""
         import paho.mqtt.client as mqtt
 
-        self.userdata = {
+        self.hass = hass
-            'hass': hass,
+        self.topics = {}
-            'topics': {},
+        self.progress = {}
-            'progress': {},
-        }
 
         if client_id is None:
             self._mqttc = mqtt.Client(protocol=mqtt.MQTTv311)
         else:
             self._mqttc = mqtt.Client(client_id, protocol=mqtt.MQTTv311)
 
-        self._mqttc.user_data_set(self.userdata)
-
         if username is not None:
             self._mqttc.username_pw_set(username, password)
         if certificate is not None:
             self._mqttc.tls_set(certificate)
 
-        self._mqttc.on_subscribe = _mqtt_on_subscribe
+        self._mqttc.on_subscribe = self._mqtt_on_subscribe
-        self._mqttc.on_unsubscribe = _mqtt_on_unsubscribe
+        self._mqttc.on_unsubscribe = self._mqtt_on_unsubscribe
-        self._mqttc.on_connect = _mqtt_on_connect
+        self._mqttc.on_connect = self._mqtt_on_connect
-        self._mqttc.on_disconnect = _mqtt_on_disconnect
+        self._mqttc.on_disconnect = self._mqtt_on_disconnect
-        self._mqttc.on_message = _mqtt_on_message
+        self._mqttc.on_message = self._mqtt_on_message
 
         self._mqttc.connect(broker, port, keepalive)
 
@@ -178,35 +175,31 @@ class MQTT(object):
 
     def stop(self):
         """Stop the MQTT client."""
+        self._mqttc.disconnect()
         self._mqttc.loop_stop()
 
     def subscribe(self, topic, qos):
         """Subscribe to a topic."""
-        if topic in self.userdata['topics']:
+        assert isinstance(topic, str)
+
+        if topic in self.topics:
             return
         result, mid = self._mqttc.subscribe(topic, qos)
         _raise_on_error(result)
-        self.userdata['progress'][mid] = topic
+        self.progress[mid] = topic
-        self.userdata['topics'][topic] = None
+        self.topics[topic] = None
 
     def unsubscribe(self, topic):
         """Unsubscribe from topic."""
         result, mid = self._mqttc.unsubscribe(topic)
         _raise_on_error(result)
-        self.userdata['progress'][mid] = topic
+        self.progress[mid] = topic
 
+    def _mqtt_on_connect(self, _mqttc, _userdata, _flags, result_code):
+        """On connect callback.
 
-def _mqtt_on_message(mqttc, userdata, msg):
+        Resubscribe to all topics we were subscribed to.
-    """ Message callback """
+        """
-    userdata['hass'].bus.fire(EVENT_MQTT_MESSAGE_RECEIVED, {
-        ATTR_TOPIC: msg.topic,
-        ATTR_QOS: msg.qos,
-        ATTR_PAYLOAD: msg.payload.decode('utf-8'),
-    })
-
-
-def _mqtt_on_connect(mqttc, userdata, flags, result_code):
-    """ On connect, resubscribe to all topics we were subscribed to. """
         if result_code != 0:
             _LOGGER.error('Unable to connect to the MQTT broker: %s', {
                 1: 'Incorrect protocol version',
@@ -215,38 +208,52 @@ def _mqtt_on_connect(mqttc, userdata, flags, result_code):
                 4: 'Bad username or password',
                 5: 'Not authorised'
             }.get(result_code, 'Unknown reason'))
-        mqttc.disconnect()
+            self._mqttc.disconnect()
             return
 
-    old_topics = userdata['topics']
+        old_topics = self.topics
 
-    userdata['topics'] = {}
+        self.topics = {key: value for key, value in self.topics.items()
-    userdata['progress'] = {}
+                       if value is None}
 
         for topic, qos in old_topics.items():
             # qos is None if we were in process of subscribing
             if qos is not None:
-            mqttc.subscribe(topic, qos)
+                self.subscribe(topic, qos)
 
-
+    def _mqtt_on_subscribe(self, _mqttc, _userdata, mid, granted_qos):
-def _mqtt_on_subscribe(mqttc, userdata, mid, granted_qos):
+        """Subscribe successful callback."""
-    """ Called when subscribe successful. """
+        topic = self.progress.pop(mid, None)
-    topic = userdata['progress'].pop(mid, None)
         if topic is None:
             return
-    userdata['topics'][topic] = granted_qos
+        self.topics[topic] = granted_qos[0]
 
+    def _mqtt_on_message(self, _mqttc, _userdata, msg):
+        """Message received callback."""
+        self.hass.bus.fire(EVENT_MQTT_MESSAGE_RECEIVED, {
+            ATTR_TOPIC: msg.topic,
+            ATTR_QOS: msg.qos,
+            ATTR_PAYLOAD: msg.payload.decode('utf-8'),
+        })
 
-def _mqtt_on_unsubscribe(mqttc, userdata, mid, granted_qos):
+    def _mqtt_on_unsubscribe(self, _mqttc, _userdata, mid, granted_qos):
-    """ Called when subscribe successful. """
+        """Unsubscribe successful callback."""
-    topic = userdata['progress'].pop(mid, None)
+        topic = self.progress.pop(mid, None)
         if topic is None:
             return
-    userdata['topics'].pop(topic, None)
+        self.topics.pop(topic, None)
 
+    def _mqtt_on_disconnect(self, _mqttc, _userdata, result_code):
+        """Disconnected callback."""
+        self.progress = {}
+        self.topics = {key: value for key, value in self.topics.items()
+                       if value is not None}
+
+        # Remove None values from topic list
+        for key in list(self.topics):
+            if self.topics[key] is None:
+                self.topics.pop(key)
 
-def _mqtt_on_disconnect(mqttc, userdata, result_code):
-    """ Called when being disconnected. """
         # When disconnected because of calling disconnect()
         if result_code == 0:
             return
@@ -256,7 +263,7 @@ def _mqtt_on_disconnect(mqttc, userdata, result_code):
 
         while True:
             try:
-            if mqttc.reconnect() == 0:
+                if self._mqttc.reconnect() == 0:
                     _LOGGER.info('Successfully reconnected to the MQTT server')
                     break
             except socket.error:
@@ -278,7 +285,7 @@ def _raise_on_error(result):
 
 
 def _match_topic(subscription, topic):
-    """ Returns if topic matches subscription. """
+    """Test if topic matches subscription."""
     if subscription.endswith('#'):
         return (subscription[:-2] == topic or
                 topic.startswith(subscription[:-1]))

tests/components/test_mqtt.py

@@ -144,8 +144,15 @@ class TestMQTTCallbacks(unittest.TestCase):
 
     def setUp(self):  # pylint: disable=invalid-name
         self.hass = get_test_home_assistant(1)
-        mock_mqtt_component(self.hass)
+        # mock_mqtt_component(self.hass)
-        self.calls = []
+
+        with mock.patch('paho.mqtt.client.Client'):
+            mqtt.setup(self.hass, {
+                mqtt.DOMAIN: {
+                    mqtt.CONF_BROKER: 'mock-broker',
+                }
+            })
+            self.hass.config.components.append(mqtt.DOMAIN)
 
     def tearDown(self):  # pylint: disable=invalid-name
         """ Stop down stuff we started. """
@@ -162,7 +169,7 @@ class TestMQTTCallbacks(unittest.TestCase):
         MQTTMessage = namedtuple('MQTTMessage', ['topic', 'qos', 'payload'])
         message = MQTTMessage('test_topic', 1, 'Hello World!'.encode('utf-8'))
 
-        mqtt._mqtt_on_message(None, {'hass': self.hass}, message)
+        mqtt.MQTT_CLIENT._mqtt_on_message(None, {'hass': self.hass}, message)
         self.hass.pool.block_till_done()
 
         self.assertEqual(1, len(calls))
@@ -173,36 +180,55 @@ class TestMQTTCallbacks(unittest.TestCase):
 
     def test_mqtt_failed_connection_results_in_disconnect(self):
         for result_code in range(1, 6):
-            mqttc = mock.MagicMock()
+            mqtt.MQTT_CLIENT._mqttc = mock.MagicMock()
-            mqtt._mqtt_on_connect(mqttc, {'topics': {}}, 0, result_code)
+            mqtt.MQTT_CLIENT._mqtt_on_connect(None, {'topics': {}}, 0,
-            self.assertTrue(mqttc.disconnect.called)
+                                              result_code)
+            self.assertTrue(mqtt.MQTT_CLIENT._mqttc.disconnect.called)
 
     def test_mqtt_subscribes_topics_on_connect(self):
-        prev_topics = {
+        from collections import OrderedDict
-            'topic/test': 1,
+        prev_topics = OrderedDict()
-            'home/sensor': 2,
+        prev_topics['topic/test'] = 1,
-            'still/pending': None
+        prev_topics['home/sensor'] = 2,
-        }
+        prev_topics['still/pending'] = None
-        mqttc = mock.MagicMock()
+
-        mqtt._mqtt_on_connect(mqttc, {'topics': prev_topics}, 0, 0)
+        mqtt.MQTT_CLIENT.topics = prev_topics
-        self.assertFalse(mqttc.disconnect.called)
+        mqtt.MQTT_CLIENT.progress = {1: 'still/pending'}
+        # Return values for subscribe calls (rc, mid)
+        mqtt.MQTT_CLIENT._mqttc.subscribe.side_effect = ((0, 2), (0, 3))
+        mqtt.MQTT_CLIENT._mqtt_on_connect(None, None, 0, 0)
+        self.assertFalse(mqtt.MQTT_CLIENT._mqttc.disconnect.called)
 
         expected = [(topic, qos) for topic, qos in prev_topics.items()
                     if qos is not None]
-        self.assertEqual(expected, [call[1] for call
+        self.assertEqual(
-                                    in mqttc.subscribe.mock_calls])
+            expected,
+            [call[1] for call in mqtt.MQTT_CLIENT._mqttc.subscribe.mock_calls])
+        self.assertEqual({
+            1: 'still/pending',
+            2: 'topic/test',
+            3: 'home/sensor',
+        }, mqtt.MQTT_CLIENT.progress)
 
     def test_mqtt_disconnect_tries_no_reconnect_on_stop(self):
-        mqttc = mock.MagicMock()
+        mqtt.MQTT_CLIENT._mqtt_on_disconnect(None, None, 0)
-        mqtt._mqtt_on_disconnect(mqttc, {}, 0)
+        self.assertFalse(mqtt.MQTT_CLIENT._mqttc.reconnect.called)
-        self.assertFalse(mqttc.reconnect.called)
 
     @mock.patch('homeassistant.components.mqtt.time.sleep')
     def test_mqtt_disconnect_tries_reconnect(self, mock_sleep):
-        mqttc = mock.MagicMock()
+        mqtt.MQTT_CLIENT.topics = {
-        mqttc.reconnect.side_effect = [1, 1, 1, 0]
+            'test/topic': 1,
-        mqtt._mqtt_on_disconnect(mqttc, {}, 1)
+            'test/progress': None
-        self.assertTrue(mqttc.reconnect.called)
+        }
-        self.assertEqual(4, len(mqttc.reconnect.mock_calls))
+        mqtt.MQTT_CLIENT.progress = {
+            1: 'test/progress'
+        }
+        mqtt.MQTT_CLIENT._mqttc.reconnect.side_effect = [1, 1, 1, 0]
+        mqtt.MQTT_CLIENT._mqtt_on_disconnect(None, None, 1)
+        self.assertTrue(mqtt.MQTT_CLIENT._mqttc.reconnect.called)
+        self.assertEqual(4, len(mqtt.MQTT_CLIENT._mqttc.reconnect.mock_calls))
         self.assertEqual([1, 2, 4],
                          [call[1][0] for call in mock_sleep.mock_calls])
+
+        self.assertEqual({'test/topic': 1}, mqtt.MQTT_CLIENT.topics)
+        self.assertEqual({}, mqtt.MQTT_CLIENT.progress)