Add AirQuality SensorState support for Google Assistant (#80579)

* Add AirQuality SensorState support for Google Assistant

* Code formatting

* Apply suggestions from code review

Co-authored-by: Joakim Plate <elupus@ecce.se>

* Update trait.py

* Update trait.py

* Fix google_assistant tests

* Update trait.py

* Simplify sensor state payload and tests

* Add more tests to fix coverage

* fix tests

* Truncate once

---------

Co-authored-by: Joakim Plate <elupus@ecce.se>
Co-authored-by: Erik <erik@montnemery.com>

homeassistant/components/google_assistant/const.py

@@ -170,6 +170,7 @@ DEVICE_CLASS_TO_GOOGLE_TYPES = {
     (media_player.DOMAIN, media_player.MediaPlayerDeviceClass.RECEIVER): TYPE_RECEIVER,
     (media_player.DOMAIN, media_player.MediaPlayerDeviceClass.SPEAKER): TYPE_SPEAKER,
     (media_player.DOMAIN, media_player.MediaPlayerDeviceClass.TV): TYPE_TV,
+    (sensor.DOMAIN, sensor.SensorDeviceClass.AQI): TYPE_SENSOR,
     (sensor.DOMAIN, sensor.SensorDeviceClass.HUMIDITY): TYPE_SENSOR,
     (sensor.DOMAIN, sensor.SensorDeviceClass.TEMPERATURE): TYPE_SENSOR,
     (switch.DOMAIN, switch.SwitchDeviceClass.OUTLET): TYPE_OUTLET,

homeassistant/components/google_assistant/trait.py

@@ -2413,6 +2413,23 @@ class SensorStateTrait(_Trait):
     name = TRAIT_SENSOR_STATE
     commands: list[str] = []
 
+    def _air_quality_description_for_aqi(self, aqi):
+        if aqi is None or aqi.isnumeric() is False:
+            return "unknown"
+        aqi = int(aqi)
+        if aqi <= 50:
+            return "healthy"
+        if aqi <= 100:
+            return "moderate"
+        if aqi <= 150:
+            return "unhealthy for sensitive groups"
+        if aqi <= 200:
+            return "unhealthy"
+        if aqi <= 300:
+            return "very unhealthy"
+
+        return "hazardous"
+
     @classmethod
     def supported(cls, domain, features, device_class, _):
         """Test if state is supported."""
@@ -2421,20 +2438,44 @@ class SensorStateTrait(_Trait):
     def sync_attributes(self):
         """Return attributes for a sync request."""
         device_class = self.state.attributes.get(ATTR_DEVICE_CLASS)
-        if (data := self.sensor_types.get(device_class)) is not None:
-            return {
-                "sensorStatesSupported": {
+        data = self.sensor_types.get(device_class)
+
+        if device_class is None or data is None:
+            return {}
+
+        sensor_state = {
             "name": data[0],
             "numericCapabilities": {"rawValueUnit": data[1]},
         }
+
+        if device_class == sensor.SensorDeviceClass.AQI:
+            sensor_state["descriptiveCapabilities"] = {
+                "availableStates": [
+                    "healthy",
+                    "moderate",
+                    "unhealthy for sensitive groups",
+                    "unhealthy",
+                    "very unhealthy",
+                    "hazardous",
+                    "unknown",
+                ],
             }
 
+        return {"sensorStatesSupported": [sensor_state]}
+
     def query_attributes(self):
         """Return the attributes of this trait for this entity."""
         device_class = self.state.attributes.get(ATTR_DEVICE_CLASS)
-        if (data := self.sensor_types.get(device_class)) is not None:
-            return {
-                "currentSensorStateData": [
-                    {"name": data[0], "rawValue": self.state.state}
-                ]
-            }
+        data = self.sensor_types.get(device_class)
+
+        if device_class is None or data is None:
+            return {}
+
+        sensor_data = {"name": data[0], "rawValue": self.state.state}
+
+        if device_class == sensor.SensorDeviceClass.AQI:
+            sensor_data["currentSensorState"] = self._air_quality_description_for_aqi(
+                self.state.state
+            )
+
+        return {"currentSensorStateData": [sensor_data]}

tests/components/google_assistant/test_trait.py

@@ -3291,6 +3291,50 @@ async def test_channel(hass: HomeAssistant) -> None:
     assert len(media_player_calls) == 1
 
 
+async def test_air_quality_description_for_aqi(hass: HomeAssistant) -> None:
+    """Test air quality description for a given AQI value."""
+    trt = trait.SensorStateTrait(
+        hass,
+        State(
+            "sensor.test",
+            100.0,
+            {
+                "device_class": sensor.SensorDeviceClass.AQI,
+            },
+        ),
+        BASIC_CONFIG,
+    )
+
+    assert trt._air_quality_description_for_aqi("0") == "healthy"
+    assert trt._air_quality_description_for_aqi("75") == "moderate"
+    assert (
+        trt._air_quality_description_for_aqi("125") == "unhealthy for sensitive groups"
+    )
+    assert trt._air_quality_description_for_aqi("175") == "unhealthy"
+    assert trt._air_quality_description_for_aqi("250") == "very unhealthy"
+    assert trt._air_quality_description_for_aqi("350") == "hazardous"
+    assert trt._air_quality_description_for_aqi("-1") == "unknown"
+    assert trt._air_quality_description_for_aqi("non-numeric") == "unknown"
+
+
+async def test_null_device_class(hass: HomeAssistant) -> None:
+    """Test handling a null device_class in sync_attributes and query_attributes."""
+    trt = trait.SensorStateTrait(
+        hass,
+        State(
+            "sensor.test",
+            100.0,
+            {
+                "device_class": None,
+            },
+        ),
+        BASIC_CONFIG,
+    )
+
+    assert trt.sync_attributes() == {}
+    assert trt.query_attributes() == {}
+
+
 async def test_sensorstate(hass: HomeAssistant) -> None:
     """Test SensorState trait support for sensor domain."""
     sensor_types = {
@@ -3324,15 +3368,51 @@ async def test_sensorstate(hass: HomeAssistant) -> None:
         name = sensor_types[sensor_type][0]
         unit = sensor_types[sensor_type][1]
 
+        if sensor_type == sensor.SensorDeviceClass.AQI:
             assert trt.sync_attributes() == {
-            "sensorStatesSupported": {
+                "sensorStatesSupported": [
+                    {
+                        "name": name,
+                        "numericCapabilities": {"rawValueUnit": unit},
+                        "descriptiveCapabilities": {
+                            "availableStates": [
+                                "healthy",
+                                "moderate",
+                                "unhealthy for sensitive groups",
+                                "unhealthy",
+                                "very unhealthy",
+                                "hazardous",
+                                "unknown",
+                            ],
+                        },
+                    }
+                ]
+            }
+        else:
+            assert trt.sync_attributes() == {
+                "sensorStatesSupported": [
+                    {
                         "name": name,
                         "numericCapabilities": {"rawValueUnit": unit},
                     }
+                ]
             }
 
+        if sensor_type == sensor.SensorDeviceClass.AQI:
             assert trt.query_attributes() == {
-            "currentSensorStateData": [{"name": name, "rawValue": "100.0"}]
+                "currentSensorStateData": [
+                    {
+                        "name": name,
+                        "currentSensorState": trt._air_quality_description_for_aqi(
+                            trt.state.state
+                        ),
+                        "rawValue": trt.state.state,
+                    },
+                ]
+            }
+        else:
+            assert trt.query_attributes() == {
+                "currentSensorStateData": [{"name": name, "rawValue": trt.state.state}]
             }
 
     assert helpers.get_google_type(sensor.DOMAIN, None) is not None