Implement base for MQTT device registry integration

homeassistant/components/mqtt/__init__.py

@@ -21,7 +21,7 @@
 from homeassistant import config_entries
 from homeassistant.const import (
     CONF_PASSWORD, CONF_PAYLOAD, CONF_PORT, CONF_PROTOCOL, CONF_USERNAME,
-    CONF_VALUE_TEMPLATE, EVENT_HOMEASSISTANT_STOP)
+    CONF_VALUE_TEMPLATE, EVENT_HOMEASSISTANT_STOP, CONF_NAME)
 from homeassistant.core import Event, ServiceCall, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import config_validation as cv
@@ -73,6 +73,12 @@
 CONF_QOS = 'qos'
 CONF_RETAIN = 'retain'
 
+CONF_IDENTIFIERS = 'identifiers'
+CONF_CONNECTIONS = 'connections'
+CONF_MANUFACTURER = 'manufacturer'
+CONF_MODEL = 'model'
+CONF_SW_VERSION = 'sw_version'
+
 PROTOCOL_31 = '3.1'
 PROTOCOL_311 = '3.1.1'
 
@@ -144,6 +150,15 @@ def valid_publish_topic(value: Any) -> str:
     return value
 
 
+def validate_device_has_at_least_one_identifier(value: ConfigType) -> \
+        ConfigType:
+    """Validate that a device info entry has at least one identifying value."""
+    if not value.get(CONF_IDENTIFIERS) and not value.get(CONF_CONNECTIONS):
+        raise vol.Invalid("Device must have at least one identifying value in "
+                          "'identifiers' and/or 'connections'")
+    return value
+
+
 _VALID_QOS_SCHEMA = vol.All(vol.Coerce(int), vol.In([0, 1, 2]))
 
 CLIENT_KEY_AUTH_MSG = 'client_key and client_cert must both be present in ' \
@@ -198,6 +213,17 @@ def valid_publish_topic(value: Any) -> str:
                  default=DEFAULT_PAYLOAD_NOT_AVAILABLE): cv.string,
 })
 
+MQTT_ENTITY_DEVICE_INFO_SCHEMA = vol.All(vol.Schema({
+    vol.Optional(CONF_IDENTIFIERS, default=list):
+        vol.All(cv.ensure_list, [cv.string]),
+    vol.Optional(CONF_CONNECTIONS, default=list):
+        vol.All(cv.ensure_list, [vol.All(vol.Length(2), [cv.string])]),
+    vol.Optional(CONF_MANUFACTURER): cv.string,
+    vol.Optional(CONF_MODEL): cv.string,
+    vol.Optional(CONF_NAME): cv.string,
+    vol.Optional(CONF_SW_VERSION): cv.string,
+}), validate_device_has_at_least_one_identifier)
+
 MQTT_BASE_PLATFORM_SCHEMA = cv.PLATFORM_SCHEMA.extend(SCHEMA_BASE)
 
 # Sensor type platforms subscribe to MQTT events
@@ -868,3 +894,41 @@ def discovery_callback(payload):
                 self.hass,
                 MQTT_DISCOVERY_UPDATED.format(self._discovery_hash),
                 discovery_callback)
+
+
+class MqttEntityDeviceInfo(Entity):
+    """Mixin used for mqtt platforms that support the device registry."""
+
+    def __init__(self, device_config: Optional[ConfigType]) -> None:
+        """Initialize the device mixin."""
+        self._device_config = device_config
+
+    @property
+    def device_info(self):
+        """Return a device description for device registry."""
+        if not self._device_config:
+            return None
+
+        info = {
+            'identifiers': {
+                (DOMAIN, id_)
+                for id_ in self._device_config[CONF_IDENTIFIERS]
+            },
+            'connections': {
+                tuple(x) for x in self._device_config[CONF_CONNECTIONS]
+            }
+        }
+
+        if CONF_MANUFACTURER in self._device_config:
+            info['manufacturer'] = self._device_config[CONF_MANUFACTURER]
+
+        if CONF_MODEL in self._device_config:
+            info['model'] = self._device_config[CONF_MODEL]
+
+        if CONF_NAME in self._device_config:
+            info['name'] = self._device_config[CONF_NAME]
+
+        if CONF_SW_VERSION in self._device_config:
+            info['sw_version'] = self._device_config[CONF_SW_VERSION]
+
+        return info

homeassistant/components/sensor/mqtt.py

@@ -16,12 +16,12 @@
 from homeassistant.components.mqtt import (
     ATTR_DISCOVERY_HASH, CONF_AVAILABILITY_TOPIC, CONF_STATE_TOPIC,
     CONF_PAYLOAD_AVAILABLE, CONF_PAYLOAD_NOT_AVAILABLE, CONF_QOS,
-    MqttAvailability, MqttDiscoveryUpdate)
+    MqttAvailability, MqttDiscoveryUpdate, MqttEntityDeviceInfo)
 from homeassistant.components.mqtt.discovery import MQTT_DISCOVERY_NEW
 from homeassistant.components.sensor import DEVICE_CLASSES_SCHEMA
 from homeassistant.const import (
     CONF_FORCE_UPDATE, CONF_NAME, CONF_VALUE_TEMPLATE, STATE_UNKNOWN,
-    CONF_UNIT_OF_MEASUREMENT, CONF_ICON, CONF_DEVICE_CLASS)
+    CONF_UNIT_OF_MEASUREMENT, CONF_ICON, CONF_DEVICE_CLASS, CONF_DEVICE)
 from homeassistant.helpers.entity import Entity
 from homeassistant.components import mqtt
 import homeassistant.helpers.config_validation as cv
@@ -51,6 +51,7 @@
     # Integrations shouldn't never expose unique_id through configuration
     # this here is an exception because MQTT is a msg transport, not a protocol
     vol.Optional(CONF_UNIQUE_ID): cv.string,
+    vol.Optional(CONF_DEVICE): mqtt.MQTT_ENTITY_DEVICE_INFO_SCHEMA,
 }).extend(mqtt.MQTT_AVAILABILITY_SCHEMA.schema)
 
 
@@ -95,22 +96,25 @@ async def _async_setup_entity(hass: HomeAssistantType, config: ConfigType,
         config.get(CONF_AVAILABILITY_TOPIC),
         config.get(CONF_PAYLOAD_AVAILABLE),
         config.get(CONF_PAYLOAD_NOT_AVAILABLE),
+        config.get(CONF_DEVICE),
         discovery_hash,
     )])
 
 
-class MqttSensor(MqttAvailability, MqttDiscoveryUpdate, Entity):
+class MqttSensor(MqttAvailability, MqttDiscoveryUpdate, MqttEntityDeviceInfo,
+                 Entity):
     """Representation of a sensor that can be updated using MQTT."""
 
     def __init__(self, name, state_topic, qos, unit_of_measurement,
                  force_update, expire_after, icon, device_class: Optional[str],
                  value_template, json_attributes, unique_id: Optional[str],
                  availability_topic, payload_available, payload_not_available,
-                 discovery_hash):
+                 device_config: Optional[ConfigType], discovery_hash):
         """Initialize the sensor."""
         MqttAvailability.__init__(self, availability_topic, qos,
                                   payload_available, payload_not_available)
         MqttDiscoveryUpdate.__init__(self, discovery_hash)
+        MqttEntityDeviceInfo.__init__(self, device_config)
         self._state = STATE_UNKNOWN
         self._name = name
         self._state_topic = state_topic

tests/components/mqtt/test_init.py

@@ -190,6 +190,47 @@ def test_validate_publish_topic(self):
         # Topic names beginning with $ SHOULD NOT be used, but can
         mqtt.valid_publish_topic('$SYS/')
 
+    def test_entity_device_info_schema(self):
+        """Test MQTT entity device info validation."""
+        # just identifier
+        mqtt.MQTT_ENTITY_DEVICE_INFO_SCHEMA({
+            'identifiers': ['abcd']
+        })
+        mqtt.MQTT_ENTITY_DEVICE_INFO_SCHEMA({
+            'identifiers': 'abcd'
+        })
+        # just connection
+        mqtt.MQTT_ENTITY_DEVICE_INFO_SCHEMA({
+            'connections': [
+                ['mac', '02:5b:26:a8:dc:12'],
+            ]
+        })
+        # full device info
+        mqtt.MQTT_ENTITY_DEVICE_INFO_SCHEMA({
+            'identifiers': ['helloworld', 'hello'],
+            'connections': [
+                ["mac", "02:5b:26:a8:dc:12"],
+                ["zigbee", "zigbee_id"],
+            ],
+            'manufacturer': 'Whatever',
+            'name': 'Beer',
+            'model': 'Glass',
+            'sw_version': '0.1-beta',
+        })
+        # no identifiers
+        self.assertRaises(vol.Invalid, mqtt.MQTT_ENTITY_DEVICE_INFO_SCHEMA, {
+            'manufacturer': 'Whatever',
+            'name': 'Beer',
+            'model': 'Glass',
+            'sw_version': '0.1-beta',
+        })
+        # empty identifiers
+        self.assertRaises(vol.Invalid, mqtt.MQTT_ENTITY_DEVICE_INFO_SCHEMA, {
+            'identifiers': [],
+            'connections': [],
+            'name': 'Beer',
+        })
+
 
 # pylint: disable=invalid-name
 class TestMQTTCallbacks(unittest.TestCase):

tests/components/sensor/test_mqtt.py

@@ -1,4 +1,5 @@
 """The tests for the MQTT sensor platform."""
+import json
 import unittest
 
 from datetime import timedelta, datetime
@@ -411,3 +412,41 @@ async def test_discovery_removal_sensor(hass, mqtt_mock, caplog):
     await hass.async_block_till_done()
     state = hass.states.get('sensor.beer')
     assert state is None
+
+
+async def test_entity_device_info_with_identifier(hass, mqtt_mock):
+    """Test MQTT sensor 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',
+        'state_topic': 'test-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/sensor/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'