Use DeviceInfo in unifi

homeassistant/components/unifi/device_tracker.py

@@ -18,10 +18,12 @@
 from homeassistant.components.device_tracker import DOMAIN
 from homeassistant.components.device_tracker.config_entry import ScannerEntity
 from homeassistant.components.device_tracker.const import SOURCE_TYPE_ROUTER
+from homeassistant.const import ATTR_NAME
 from homeassistant.core import callback
 from homeassistant.helpers import device_registry as dr
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.entity import DeviceInfo
 import homeassistant.util.dt as dt_util
 
 from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN
@@ -405,17 +407,17 @@ def available(self) -> bool:
         return not self.device.disabled and self.controller.available
 
     @property
-    def device_info(self):
+    def device_info(self) -> DeviceInfo:
         """Return a device description for device registry."""
-        info = {
-            "connections": {(CONNECTION_NETWORK_MAC, self.device.mac)},
-            "manufacturer": ATTR_MANUFACTURER,
-            "model": self.device.model,
-            "sw_version": self.device.version,
-        }
+        info = DeviceInfo(
+            connections={(CONNECTION_NETWORK_MAC, self.device.mac)},
+            manufacturer=ATTR_MANUFACTURER,
+            model=self.device.model,
+            sw_version=self.device.version,
+        )
 
         if self.device.name:
-            info["name"] = self.device.name
+            info[ATTR_NAME] = self.device.name
 
         return info
 

homeassistant/components/unifi/switch.py

@@ -369,10 +369,10 @@ async def options_updated(self) -> None:
     @property
     def device_info(self) -> DeviceInfo:
         """Return a service description for device registry."""
-        return {
-            "identifiers": {(DOMAIN, f"unifi_controller_{self._item.site_id}")},
-            "name": "UniFi Controller",
-            "manufacturer": ATTR_MANUFACTURER,
-            "model": "UniFi Controller",
-            "entry_type": "service",
-        }
+        return DeviceInfo(
+            entry_type="service",
+            identifiers={(DOMAIN, f"unifi_controller_{self._item.site_id}")},
+            manufacturer=ATTR_MANUFACTURER,
+            model="UniFi Controller",
+            name="UniFi Controller",
+        )

homeassistant/components/unifi/unifi_client.py

@@ -47,8 +47,8 @@ def available(self) -> bool:
     @property
     def device_info(self) -> DeviceInfo:
         """Return a client description for device registry."""
-        return {
-            "connections": {(CONNECTION_NETWORK_MAC, self.client.mac)},
-            "default_name": self.name,
-            "default_manufacturer": self.client.oui,
-        }
+        return DeviceInfo(
+            connections={(CONNECTION_NETWORK_MAC, self.client.mac)},
+            default_manufacturer=self.client.oui,
+            default_name=self.name,
+        )