[3/3] mqtt-vacuum device-registry (#19479)

* add device registry to mqtt-vacuum

homeassistant/components/vacuum/mqtt.py

@@ -11,14 +11,15 @@
 from homeassistant.components import mqtt
 from homeassistant.components.mqtt import (
     ATTR_DISCOVERY_HASH, MqttAvailability, MqttDiscoveryUpdate,
-    subscription)
+    MqttEntityDeviceInfo, subscription)
 from homeassistant.components.mqtt.discovery import MQTT_DISCOVERY_NEW
 from homeassistant.components.vacuum import (
     SUPPORT_BATTERY, SUPPORT_CLEAN_SPOT, SUPPORT_FAN_SPEED,
     SUPPORT_LOCATE, SUPPORT_PAUSE, SUPPORT_RETURN_HOME, SUPPORT_SEND_COMMAND,
     SUPPORT_STATUS, SUPPORT_STOP, SUPPORT_TURN_OFF, SUPPORT_TURN_ON,
     VacuumDevice, DOMAIN)
-from homeassistant.const import ATTR_SUPPORTED_FEATURES, CONF_NAME
+from homeassistant.const import (
+    ATTR_SUPPORTED_FEATURES, CONF_NAME, CONF_DEVICE)
 from homeassistant.core import callback
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
@@ -93,6 +94,7 @@ def strings_to_services(strings):
 CONF_SET_FAN_SPEED_TOPIC = 'set_fan_speed_topic'
 CONF_FAN_SPEED_LIST = 'fan_speed_list'
 CONF_SEND_COMMAND_TOPIC = 'send_command_topic'
+CONF_UNIQUE_ID = 'unique_id'
 
 DEFAULT_NAME = 'MQTT Vacuum'
 DEFAULT_RETAIN = False
@@ -143,6 +145,8 @@ def strings_to_services(strings):
     vol.Optional(CONF_FAN_SPEED_LIST, default=[]):
         vol.All(cv.ensure_list, [cv.string]),
     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)
 
 
@@ -171,7 +175,9 @@ async def _async_setup_entity(config, async_add_entities,
     async_add_entities([MqttVacuum(config, discovery_hash)])
 
 
-class MqttVacuum(MqttAvailability, MqttDiscoveryUpdate, VacuumDevice):
+# pylint: disable=too-many-ancestors
+class MqttVacuum(MqttAvailability, MqttDiscoveryUpdate, MqttEntityDeviceInfo,
+                 VacuumDevice):
     """Representation of a MQTT-controlled vacuum."""
 
     def __init__(self, config, discovery_info):
@@ -185,6 +191,7 @@ def __init__(self, config, discovery_info):
         self._fan_speed = 'unknown'
         self._fan_speed_list = []
         self._sub_state = None
+        self._unique_id = config.get(CONF_UNIQUE_ID)
 
         # Load config
         self._setup_from_config(config)
@@ -193,11 +200,13 @@ def __init__(self, config, discovery_info):
         availability_topic = config.get(mqtt.CONF_AVAILABILITY_TOPIC)
         payload_available = config.get(mqtt.CONF_PAYLOAD_AVAILABLE)
         payload_not_available = config.get(mqtt.CONF_PAYLOAD_NOT_AVAILABLE)
+        device_config = config.get(CONF_DEVICE)
 
         MqttAvailability.__init__(self, availability_topic, qos,
                                   payload_available, payload_not_available)
         MqttDiscoveryUpdate.__init__(self, discovery_info,
                                      self.discovery_update)
+        MqttEntityDeviceInfo.__init__(self, device_config)
 
     def _setup_from_config(self, config):
         self._name = config.get(CONF_NAME)
@@ -361,6 +370,11 @@ def is_on(self):
         """Return true if vacuum is on."""
         return self._cleaning
 
+    @property
+    def unique_id(self):
+        """Return a unique ID."""
+        return self._unique_id
+
     @property
     def status(self):
         """Return a status string for the vacuum."""

tests/components/vacuum/test_mqtt.py

@@ -1,4 +1,5 @@
 """The tests for the Mqtt vacuum platform."""
+import json
 import pytest
 
 from homeassistant.setup import async_setup_component
@@ -15,7 +16,6 @@
     async_fire_mqtt_message, MockConfigEntry)
 from tests.components.vacuum import common
 
-
 default_config = {
     CONF_PLATFORM: 'mqtt',
     CONF_NAME: 'mqtttest',
@@ -330,3 +330,66 @@ async def test_discovery_update_vacuum(hass, mock_publish):
     assert state.name == 'Milk'
     state = hass.states.get('vacuum.milk')
     assert state is None
+
+
+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)
+    assert await async_setup_component(hass, vacuum.DOMAIN, {
+        vacuum.DOMAIN: [{
+            'platform': 'mqtt',
+            'name': 'Test 1',
+            'command_topic': 'command-topic',
+            'unique_id': 'TOTALLY_UNIQUE'
+        }, {
+            'platform': 'mqtt',
+            'name': 'Test 2',
+            'command_topic': 'command-topic',
+            'unique_id': 'TOTALLY_UNIQUE'
+        }]
+    })
+
+    async_fire_mqtt_message(hass, 'test-topic', 'payload')
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    assert len(hass.states.async_entity_ids()) == 2
+    # all vacuums group is 1, unique id created is 1
+
+
+async def test_entity_device_info_with_identifier(hass, mock_publish):
+    """Test MQTT vacuum device registry integration."""
+    entry = MockConfigEntry(domain=mqtt.DOMAIN)
+    entry.add_to_hass(hass)
+    await async_start(hass, 'homeassistant', {}, entry)
+    registry = await hass.helpers.device_registry.async_get_registry()
+
+    data = json.dumps({
+        'platform': 'mqtt',
+        'name': 'Test 1',
+        'command_topic': 'test-command-topic',
+        'device': {
+            'identifiers': ['helloworld'],
+            'connections': [
+                ["mac", "02:5b:26:a8:dc:12"],
+            ],
+            'manufacturer': 'Whatever',
+            'name': 'Beer',
+            'model': 'Glass',
+            'sw_version': '0.1-beta',
+        },
+        'unique_id': 'veryunique'
+    })
+    async_fire_mqtt_message(hass, 'homeassistant/vacuum/bla/config',
+                            data)
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+
+    device = registry.async_get_device({('mqtt', 'helloworld')}, set())
+    assert device is not None
+    assert device.identifiers == {('mqtt', 'helloworld')}
+    assert device.connections == {('mac', "02:5b:26:a8:dc:12")}
+    assert device.manufacturer == 'Whatever'
+    assert device.name == 'Beer'
+    assert device.model == 'Glass'
+    assert device.sw_version == '0.1-beta'