From dd21bf73fc5af18e0d2302f8015f5dc797ac6669 Mon Sep 17 00:00:00 2001
From: Matthew Garrett <mjg59@srcf.ucam.org>
Date: Mon, 30 Aug 2021 20:33:06 -0700
Subject: [PATCH] Assistant sensors (#55480)

---
 .../components/google_assistant/const.py      |  1 +
 .../components/google_assistant/trait.py      | 59 +++++++++++++++++++
 .../components/google_assistant/test_trait.py | 53 +++++++++++++++++
 3 files changed, 113 insertions(+)

diff --git a/homeassistant/components/google_assistant/const.py b/homeassistant/components/google_assistant/const.py
index 2e43e20f124..d23560b85c1 100644
--- a/homeassistant/components/google_assistant/const.py
+++ b/homeassistant/components/google_assistant/const.py
@@ -133,6 +133,7 @@ DOMAIN_TO_GOOGLE_TYPES = {
     media_player.DOMAIN: TYPE_SETTOP,
     scene.DOMAIN: TYPE_SCENE,
     script.DOMAIN: TYPE_SCENE,
+    sensor.DOMAIN: TYPE_SENSOR,
     select.DOMAIN: TYPE_SENSOR,
     switch.DOMAIN: TYPE_SWITCH,
     vacuum.DOMAIN: TYPE_VACUUM,
diff --git a/homeassistant/components/google_assistant/trait.py b/homeassistant/components/google_assistant/trait.py
index dda8a04c2ed..d1ed328703e 100644
--- a/homeassistant/components/google_assistant/trait.py
+++ b/homeassistant/components/google_assistant/trait.py
@@ -108,6 +108,7 @@ TRAIT_MEDIA_STATE = f"{PREFIX_TRAITS}MediaState"
 TRAIT_CHANNEL = f"{PREFIX_TRAITS}Channel"
 TRAIT_LOCATOR = f"{PREFIX_TRAITS}Locator"
 TRAIT_ENERGYSTORAGE = f"{PREFIX_TRAITS}EnergyStorage"
+TRAIT_SENSOR_STATE = f"{PREFIX_TRAITS}SensorState"
 
 PREFIX_COMMANDS = "action.devices.commands."
 COMMAND_ONOFF = f"{PREFIX_COMMANDS}OnOff"
@@ -2286,3 +2287,61 @@ class ChannelTrait(_Trait):
             blocking=True,
             context=data.context,
         )
+
+
+@register_trait
+class SensorStateTrait(_Trait):
+    """Trait to get sensor state.
+
+    https://developers.google.com/actions/smarthome/traits/sensorstate
+    """
+
+    sensor_types = {
+        sensor.DEVICE_CLASS_AQI: ("AirQuality", "AQI"),
+        sensor.DEVICE_CLASS_CO: ("CarbonDioxideLevel", "PARTS_PER_MILLION"),
+        sensor.DEVICE_CLASS_CO2: ("CarbonMonoxideLevel", "PARTS_PER_MILLION"),
+        sensor.DEVICE_CLASS_PM25: ("PM2.5", "MICROGRAMS_PER_CUBIC_METER"),
+        sensor.DEVICE_CLASS_PM10: ("PM10", "MICROGRAMS_PER_CUBIC_METER"),
+        sensor.DEVICE_CLASS_VOLATILE_ORGANIC_COMPOUNDS: (
+            "VolatileOrganicCompounds",
+            "PARTS_PER_MILLION",
+        ),
+    }
+
+    name = TRAIT_SENSOR_STATE
+    commands = []
+
+    @staticmethod
+    def supported(domain, features, device_class, _):
+        """Test if state is supported."""
+        return domain == sensor.DOMAIN and device_class in (
+            sensor.DEVICE_CLASS_AQI,
+            sensor.DEVICE_CLASS_CO,
+            sensor.DEVICE_CLASS_CO2,
+            sensor.DEVICE_CLASS_PM25,
+            sensor.DEVICE_CLASS_PM10,
+            sensor.DEVICE_CLASS_VOLATILE_ORGANIC_COMPOUNDS,
+        )
+
+    def sync_attributes(self):
+        """Return attributes for a sync request."""
+        device_class = self.state.attributes.get(ATTR_DEVICE_CLASS)
+        data = self.sensor_types.get(device_class)
+        if data is not None:
+            return {
+                "sensorStatesSupported": {
+                    "name": data[0],
+                    "numericCapabilities": {"rawValueUnit": data[1]},
+                }
+            }
+
+    def query_attributes(self):
+        """Return the attributes of this trait for this entity."""
+        device_class = self.state.attributes.get(ATTR_DEVICE_CLASS)
+        data = self.sensor_types.get(device_class)
+        if data is not None:
+            return {
+                "currentSensorStateData": [
+                    {"name": data[0], "rawValue": self.state.state}
+                ]
+            }
diff --git a/tests/components/google_assistant/test_trait.py b/tests/components/google_assistant/test_trait.py
index 50006060f51..290aa00bb47 100644
--- a/tests/components/google_assistant/test_trait.py
+++ b/tests/components/google_assistant/test_trait.py
@@ -3003,3 +3003,56 @@ async def test_channel(hass):
     with pytest.raises(SmartHomeError, match="Unsupported command"):
         await trt.execute("Unknown command", BASIC_DATA, {"channelNumber": "1"}, {})
     assert len(media_player_calls) == 1
+
+
+async def test_sensorstate(hass):
+    """Test SensorState trait support for sensor domain."""
+    sensor_types = {
+        sensor.DEVICE_CLASS_AQI: ("AirQuality", "AQI"),
+        sensor.DEVICE_CLASS_CO: ("CarbonDioxideLevel", "PARTS_PER_MILLION"),
+        sensor.DEVICE_CLASS_CO2: ("CarbonMonoxideLevel", "PARTS_PER_MILLION"),
+        sensor.DEVICE_CLASS_PM25: ("PM2.5", "MICROGRAMS_PER_CUBIC_METER"),
+        sensor.DEVICE_CLASS_PM10: ("PM10", "MICROGRAMS_PER_CUBIC_METER"),
+        sensor.DEVICE_CLASS_VOLATILE_ORGANIC_COMPOUNDS: (
+            "VolatileOrganicCompounds",
+            "PARTS_PER_MILLION",
+        ),
+    }
+
+    for sensor_type in sensor_types:
+        assert helpers.get_google_type(sensor.DOMAIN, None) is not None
+        assert trait.SensorStateTrait.supported(sensor.DOMAIN, None, sensor_type, None)
+
+        trt = trait.SensorStateTrait(
+            hass,
+            State(
+                "sensor.test",
+                100.0,
+                {
+                    "device_class": sensor_type,
+                },
+            ),
+            BASIC_CONFIG,
+        )
+
+        name = sensor_types[sensor_type][0]
+        unit = sensor_types[sensor_type][1]
+
+        assert trt.sync_attributes() == {
+            "sensorStatesSupported": {
+                "name": name,
+                "numericCapabilities": {"rawValueUnit": unit},
+            }
+        }
+
+        assert trt.query_attributes() == {
+            "currentSensorStateData": [{"name": name, "rawValue": "100.0"}]
+        }
+
+    assert helpers.get_google_type(sensor.DOMAIN, None) is not None
+    assert (
+        trait.SensorStateTrait.supported(
+            sensor.DOMAIN, None, sensor.DEVICE_CLASS_MONETARY, None
+        )
+        is False
+    )