Add JSON attribute topic to MQTT vacuum

homeassistant/components/mqtt/vacuum.py

@@ -10,7 +10,7 @@ import voluptuous as vol
 
 from homeassistant.components import mqtt
 from homeassistant.components.mqtt import (
-    ATTR_DISCOVERY_HASH, MqttAvailability, MqttDiscoveryUpdate,
+    ATTR_DISCOVERY_HASH, MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
     MqttEntityDeviceInfo, subscription)
 from homeassistant.components.mqtt.discovery import MQTT_DISCOVERY_NEW
 from homeassistant.components.vacuum import (
@@ -147,7 +147,8 @@ PLATFORM_SCHEMA = mqtt.MQTT_BASE_PLATFORM_SCHEMA.extend({
     vol.Optional(CONF_SEND_COMMAND_TOPIC): mqtt.valid_publish_topic,
     vol.Optional(CONF_UNIQUE_ID): cv.string,
     vol.Optional(CONF_DEVICE): mqtt.MQTT_ENTITY_DEVICE_INFO_SCHEMA,
-}).extend(mqtt.MQTT_AVAILABILITY_SCHEMA.schema)
+}).extend(mqtt.MQTT_AVAILABILITY_SCHEMA.schema).extend(
+    mqtt.MQTT_JSON_ATTRS_SCHEMA.schema)
 
 
 async def async_setup_platform(hass, config, async_add_entities,
@@ -176,8 +177,8 @@ async def _async_setup_entity(config, async_add_entities,
 
 
 # pylint: disable=too-many-ancestors
-class MqttVacuum(MqttAvailability, MqttDiscoveryUpdate, MqttEntityDeviceInfo,
+class MqttVacuum(MqttAttributes, MqttAvailability, MqttDiscoveryUpdate,
-                 VacuumDevice):
+                 MqttEntityDeviceInfo, VacuumDevice):
     """Representation of a MQTT-controlled vacuum."""
 
     def __init__(self, config, discovery_info):
@@ -198,6 +199,7 @@ class MqttVacuum(MqttAvailability, MqttDiscoveryUpdate, MqttEntityDeviceInfo,
 
         device_config = config.get(CONF_DEVICE)
 
+        MqttAttributes.__init__(self, config)
         MqttAvailability.__init__(self, config)
         MqttDiscoveryUpdate.__init__(self, discovery_info,
                                      self.discovery_update)
@@ -253,6 +255,7 @@ class MqttVacuum(MqttAvailability, MqttDiscoveryUpdate, MqttEntityDeviceInfo,
         """Handle updated discovery message."""
         config = PLATFORM_SCHEMA(discovery_payload)
         self._setup_from_config(config)
+        await self.attributes_discovery_update(config)
         await self.availability_discovery_update(config)
         await self._subscribe_topics()
         self.async_schedule_update_ha_state()
@@ -265,6 +268,7 @@ class MqttVacuum(MqttAvailability, MqttDiscoveryUpdate, MqttEntityDeviceInfo,
     async def async_will_remove_from_hass(self):
         """Unsubscribe when removed."""
         await subscription.async_unsubscribe_topics(self.hass, self._sub_state)
+        await MqttAttributes.async_will_remove_from_hass(self)
         await MqttAvailability.async_will_remove_from_hass(self)
 
     async def _subscribe_topics(self):

tests/components/mqtt/test_vacuum.py

@@ -333,6 +333,106 @@ async def test_discovery_update_vacuum(hass, mock_publish):
     assert state is None
 
 
+async def test_setting_attribute_via_mqtt_json_message(hass, mqtt_mock):
+    """Test the setting of attribute via MQTT with JSON payload."""
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: {
+            'platform': 'mqtt',
+            'name': 'test',
+            'state_topic': 'test-topic',
+            'json_attributes_topic': 'attr-topic'
+        }
+    })
+
+    async_fire_mqtt_message(hass, 'attr-topic', '{ "val": "100" }')
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.test')
+
+    assert '100' == state.attributes.get('val')
+
+
+async def test_update_with_json_attrs_not_dict(hass, mqtt_mock, caplog):
+    """Test attributes get extracted from a JSON result."""
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: {
+            'platform': 'mqtt',
+            'name': 'test',
+            'state_topic': 'test-topic',
+            'json_attributes_topic': 'attr-topic'
+        }
+    })
+
+    async_fire_mqtt_message(hass, 'attr-topic', '[ "list", "of", "things"]')
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.test')
+
+    assert state.attributes.get('val') is None
+    assert 'JSON result was not a dictionary' in caplog.text
+
+
+async def test_update_with_json_attrs_bad_JSON(hass, mqtt_mock, caplog):
+    """Test attributes get extracted from a JSON result."""
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: {
+            'platform': 'mqtt',
+            'name': 'test',
+            'state_topic': 'test-topic',
+            'json_attributes_topic': 'attr-topic'
+        }
+    })
+
+    async_fire_mqtt_message(hass, 'attr-topic', 'This is not JSON')
+    await hass.async_block_till_done()
+
+    state = hass.states.get('vacuum.test')
+    assert state.attributes.get('val') is None
+    assert 'Erroneous JSON: This is not JSON' in caplog.text
+
+
+async def test_discovery_update_attr(hass, mqtt_mock, caplog):
+    """Test update of discovered MQTTAttributes."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    await async_start(hass, 'homeassistant', {}, entry)
+    data1 = (
+        '{ "name": "Beer",'
+        '  "command_topic": "test_topic",'
+        '  "json_attributes_topic": "attr-topic1" }'
+    )
+    data2 = (
+        '{ "name": "Beer",'
+        '  "command_topic": "test_topic",'
+        '  "json_attributes_topic": "attr-topic2" }'
+    )
+    async_fire_mqtt_message(hass, 'homeassistant/vacuum/bla/config',
+                            data1)
+    await hass.async_block_till_done()
+    async_fire_mqtt_message(hass, 'attr-topic1', '{ "val": "100" }')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.beer')
+    assert '100' == state.attributes.get('val')
+
+    # Change json_attributes_topic
+    async_fire_mqtt_message(hass, 'homeassistant/vacuum/bla/config',
+                            data2)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    # Verify we are no longer subscribing to the old topic
+    async_fire_mqtt_message(hass, 'attr-topic1', '{ "val": "50" }')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.beer')
+    assert '100' == state.attributes.get('val')
+
+    # Verify we are subscribing to the new topic
+    async_fire_mqtt_message(hass, 'attr-topic2', '{ "val": "75" }')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    state = hass.states.get('vacuum.beer')
+    assert '75' == state.attributes.get('val')
+
+
 async def test_unique_id(hass, mock_publish):
     """Test unique id option only creates one vacuum per unique_id."""
     await async_mock_mqtt_component(hass)