From 4f56f4e7e940a6c90f4e5877a551b73712037d91 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ville=20Skytt=C3=A4?= <ville.skytta@iki.fi>
Date: Fri, 8 Nov 2019 10:19:59 +0200
Subject: [PATCH] Add Huawei LTE device registry support (#28594)

---
 .../components/huawei_lte/__init__.py         | 37 ++++++++++++++++++-
 1 file changed, 36 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/huawei_lte/__init__.py b/homeassistant/components/huawei_lte/__init__.py
index 112fcb6ec52..fa1423edcca 100644
--- a/homeassistant/components/huawei_lte/__init__.py
+++ b/homeassistant/components/huawei_lte/__init__.py
@@ -6,7 +6,7 @@ from functools import partial
 from urllib.parse import urlparse
 import ipaddress
 import logging
-from typing import Any, Callable, Dict, List, Set
+from typing import Any, Callable, Dict, List, Set, Tuple
 
 import voluptuous as vol
 import attr
@@ -36,6 +36,7 @@ from homeassistant.const import (
 from homeassistant.core import CALLBACK_TYPE
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import config_validation as cv, discovery
+from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.dispatcher import (
     async_dispatcher_connect,
     async_dispatcher_send,
@@ -135,6 +136,11 @@ class Router:
                 pass
         return DEFAULT_DEVICE_NAME
 
+    @property
+    def device_connections(self) -> Set[Tuple[str, str]]:
+        """Get router connections for device registry."""
+        return {(dr.CONNECTION_NETWORK_MAC, self.mac)}
+
     def update(self) -> None:
         """Update router data."""
 
@@ -283,6 +289,30 @@ async def async_setup_entry(hass: HomeAssistantType, config_entry: ConfigEntry)
     # Clear all subscriptions, enabled entities will push back theirs
     router.subscriptions.clear()
 
+    # Set up device registry
+    device_data = {}
+    sw_version = None
+    if router.data.get(KEY_DEVICE_INFORMATION):
+        device_info = router.data[KEY_DEVICE_INFORMATION]
+        serial_number = device_info.get("SerialNumber")
+        if serial_number:
+            device_data["identifiers"] = {(DOMAIN, serial_number)}
+        sw_version = device_info.get("SoftwareVersion")
+        if device_info.get("DeviceName"):
+            device_data["model"] = device_info["DeviceName"]
+    if not sw_version and router.data.get(KEY_DEVICE_BASIC_INFORMATION):
+        sw_version = router.data[KEY_DEVICE_BASIC_INFORMATION].get("SoftwareVersion")
+    if sw_version:
+        device_data["sw_version"] = sw_version
+    device_registry = await dr.async_get_registry(hass)
+    device_registry.async_get_or_create(
+        config_entry_id=config_entry.entry_id,
+        connections=router.device_connections,
+        name=router.device_name,
+        manufacturer="Huawei",
+        **device_data,
+    )
+
     # Forward config entry setup to platforms
     for domain in (DEVICE_TRACKER_DOMAIN, SENSOR_DOMAIN, SWITCH_DOMAIN):
         hass.async_create_task(
@@ -408,6 +438,11 @@ class HuaweiLteBaseEntity(Entity):
         """Huawei LTE entities report their state without polling."""
         return False
 
+    @property
+    def device_info(self) -> Dict[str, Any]:
+        """Get info for matching with parent router."""
+        return {"connections": self.router.device_connections}
+
     async def async_update(self) -> None:
         """Update state."""
         raise NotImplementedError