From 74efe78d0ac12f4e15b1c9c7eec89e347e528859 Mon Sep 17 00:00:00 2001
From: Dror Eiger <45061021+deiger@users.noreply.github.com>
Date: Tue, 26 Jan 2021 22:59:43 +0200
Subject: [PATCH] Add device metadata for Google Assistant (#45507)

* Add device metadata for Google Assistant

* Increase test coverage

* Refactor fetching the device and entity entries.
---
 .../components/google_assistant/helpers.py    | 58 ++++++++++++++-----
 .../google_assistant/test_smart_home.py       |  8 +++
 2 files changed, 52 insertions(+), 14 deletions(-)

diff --git a/homeassistant/components/google_assistant/helpers.py b/homeassistant/components/google_assistant/helpers.py
index 00633422939..b4900d83b64 100644
--- a/homeassistant/components/google_assistant/helpers.py
+++ b/homeassistant/components/google_assistant/helpers.py
@@ -4,7 +4,7 @@ from asyncio import gather
 from collections.abc import Mapping
 import logging
 import pprint
-from typing import List, Optional
+from typing import Dict, List, Optional, Tuple
 
 from aiohttp.web import json_response
 
@@ -18,6 +18,8 @@ from homeassistant.const import (
 )
 from homeassistant.core import Context, HomeAssistant, State, callback
 from homeassistant.helpers.area_registry import AreaEntry
+from homeassistant.helpers.device_registry import DeviceEntry
+from homeassistant.helpers.entity_registry import RegistryEntry
 from homeassistant.helpers.event import async_call_later
 from homeassistant.helpers.network import get_url
 from homeassistant.helpers.storage import Store
@@ -40,29 +42,50 @@ SYNC_DELAY = 15
 _LOGGER = logging.getLogger(__name__)
 
 
-async def _get_area(hass, entity_id) -> Optional[AreaEntry]:
-    """Calculate the area for a entity_id."""
-    dev_reg, ent_reg, area_reg = await gather(
+async def _get_entity_and_device(
+    hass, entity_id
+) -> Optional[Tuple[RegistryEntry, DeviceEntry]]:
+    """Fetch the entity and device entries for a entity_id."""
+    dev_reg, ent_reg = await gather(
         hass.helpers.device_registry.async_get_registry(),
         hass.helpers.entity_registry.async_get_registry(),
-        hass.helpers.area_registry.async_get_registry(),
     )
 
     entity_entry = ent_reg.async_get(entity_id)
     if not entity_entry:
+        return None, None
+    device_entry = dev_reg.devices.get(entity_entry.device_id)
+    return entity_entry, device_entry
+
+
+async def _get_area(hass, entity_entry, device_entry) -> Optional[AreaEntry]:
+    """Calculate the area for an entity."""
+    if entity_entry and entity_entry.area_id:
+        area_id = entity_entry.area_id
+    elif device_entry and device_entry.area_id:
+        area_id = device_entry.area_id
+    else:
         return None
 
-    if entity_entry.area_id:
-        area_id = entity_entry.area_id
-    else:
-        device_entry = dev_reg.devices.get(entity_entry.device_id)
-        if not (device_entry and device_entry.area_id):
-            return None
-        area_id = device_entry.area_id
-
+    area_reg = await hass.helpers.area_registry.async_get_registry()
     return area_reg.areas.get(area_id)
 
 
+async def _get_device_info(device_entry) -> Optional[Dict[str, str]]:
+    """Retrieve the device info for a device."""
+    if not device_entry:
+        return None
+
+    device_info = {}
+    if device_entry.manufacturer:
+        device_info["manufacturer"] = device_entry.manufacturer
+    if device_entry.model:
+        device_info["model"] = device_entry.model
+    if device_entry.sw_version:
+        device_info["swVersion"] = device_entry.sw_version
+    return device_info
+
+
 class AbstractConfig(ABC):
     """Hold the configuration for Google Assistant."""
 
@@ -438,6 +461,9 @@ class GoogleEntity:
         name = (entity_config.get(CONF_NAME) or state.name).strip()
         domain = state.domain
         device_class = state.attributes.get(ATTR_DEVICE_CLASS)
+        entity_entry, device_entry = await _get_entity_and_device(
+            self.hass, state.entity_id
+        )
 
         traits = self.traits()
 
@@ -475,10 +501,14 @@ class GoogleEntity:
         if room:
             device["roomHint"] = room
         else:
-            area = await _get_area(self.hass, state.entity_id)
+            area = await _get_area(self.hass, entity_entry, device_entry)
             if area and area.name:
                 device["roomHint"] = area.name
 
+        device_info = await _get_device_info(device_entry)
+        if device_info:
+            device["deviceInfo"] = device_info
+
         return device
 
     @callback
diff --git a/tests/components/google_assistant/test_smart_home.py b/tests/components/google_assistant/test_smart_home.py
index 5739370b187..9c8f9a48338 100644
--- a/tests/components/google_assistant/test_smart_home.py
+++ b/tests/components/google_assistant/test_smart_home.py
@@ -160,6 +160,9 @@ async def test_sync_in_area(area_on_device, hass, registries):
 
     device = registries.device.async_get_or_create(
         config_entry_id="1234",
+        manufacturer="Someone",
+        model="Some model",
+        sw_version="Some Version",
         connections={(device_registry.CONNECTION_NETWORK_MAC, "12:34:56:AB:CD:EF")},
     )
     registries.device.async_update_device(
@@ -249,6 +252,11 @@ async def test_sync_in_area(area_on_device, hass, registries):
                             "temperatureMaxK": 6535,
                         },
                     },
+                    "deviceInfo": {
+                        "manufacturer": "Someone",
+                        "model": "Some model",
+                        "swVersion": "Some Version",
+                    },
                     "roomHint": "Living Room",
                 }
             ],