Improve Airzone diagnostics

homeassistant/components/airzone/diagnostics.py

@@ -3,16 +3,25 @@
 
 from typing import Any
 
-from aioairzone.const import AZD_MAC
+from aioairzone.const import API_MAC, AZD_MAC
 
 from homeassistant.components.diagnostics.util import async_redact_data
 from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_UNIQUE_ID
 from homeassistant.core import HomeAssistant
 
 from .const import DOMAIN
 from .coordinator import AirzoneUpdateCoordinator
 
-TO_REDACT = [
+TO_REDACT_API = [
+    API_MAC,
+]
+
+TO_REDACT_CONFIG = [
+    CONF_UNIQUE_ID,
+]
+
+TO_REDACT_COORD = [
     AZD_MAC,
 ]
 
@@ -24,6 +33,7 @@ async def async_get_config_entry_diagnostics(
     coordinator: AirzoneUpdateCoordinator = hass.data[DOMAIN][config_entry.entry_id]
 
     return {
-        "info": async_redact_data(config_entry.data, TO_REDACT),
-        "data": async_redact_data(coordinator.data, TO_REDACT),
+        "api_data": async_redact_data(coordinator.airzone.raw_data(), TO_REDACT_API),
+        "config_entry": async_redact_data(config_entry.as_dict(), TO_REDACT_CONFIG),
+        "coord_data": async_redact_data(coordinator.data, TO_REDACT_COORD),
     }

tests/components/airzone/test_diagnostics.py

@@ -1,20 +1,30 @@
 """The diagnostics tests for the Airzone platform."""
 
+from unittest.mock import patch
+
 from aioairzone.const import (
+    API_DATA,
+    API_MAC,
+    API_SYSTEM_ID,
+    API_SYSTEMS,
+    API_WIFI_RSSI,
     AZD_ID,
     AZD_MASTER,
     AZD_SYSTEM,
     AZD_SYSTEMS,
     AZD_ZONES,
     AZD_ZONES_NUM,
+    RAW_HVAC,
+    RAW_WEBSERVER,
 )
 from aiohttp import ClientSession
 
 from homeassistant.components.airzone.const import DOMAIN
+from homeassistant.components.diagnostics.const import REDACTED
 from homeassistant.const import CONF_HOST, CONF_PORT
 from homeassistant.core import HomeAssistant
 
-from .util import CONFIG, async_init_integration
+from .util import CONFIG, HVAC_MOCK, HVAC_WEBSERVER_MOCK, async_init_integration
 
 from tests.components.diagnostics import get_diagnostics_for_config_entry
 
@@ -27,31 +37,55 @@ async def test_config_entry_diagnostics(
     assert hass.data[DOMAIN]
 
     config_entry = hass.config_entries.async_entries(DOMAIN)[0]
+    with patch(
+        "homeassistant.components.airzone.AirzoneLocalApi.raw_data",
+        return_value={
+            RAW_HVAC: HVAC_MOCK,
+            RAW_WEBSERVER: HVAC_WEBSERVER_MOCK,
+        },
+    ):
+        diag = await get_diagnostics_for_config_entry(hass, hass_client, config_entry)
 
-    diag = await get_diagnostics_for_config_entry(hass, hass_client, config_entry)
-
-    assert diag["info"][CONF_HOST] == CONFIG[CONF_HOST]
-    assert diag["info"][CONF_PORT] == CONFIG[CONF_PORT]
-
-    assert diag["data"][AZD_SYSTEMS]["1"][AZD_ID] == 1
-    assert diag["data"][AZD_SYSTEMS]["1"][AZD_ZONES_NUM] == 5
-
-    assert diag["data"][AZD_ZONES]["1:1"][AZD_ID] == 1
-    assert diag["data"][AZD_ZONES]["1:1"][AZD_MASTER] == 1
-    assert diag["data"][AZD_ZONES]["1:1"][AZD_SYSTEM] == 1
+    assert (
+        diag["api_data"][RAW_HVAC][API_SYSTEMS][0][API_DATA][0].items()
+        >= {
+            API_SYSTEM_ID: HVAC_MOCK[API_SYSTEMS][0][API_DATA][0][API_SYSTEM_ID],
+        }.items()
+    )
 
-    assert diag["data"][AZD_ZONES]["1:2"][AZD_ID] == 2
-    assert diag["data"][AZD_ZONES]["1:2"][AZD_MASTER] == 0
-    assert diag["data"][AZD_ZONES]["1:2"][AZD_SYSTEM] == 1
+    assert (
+        diag["api_data"][RAW_WEBSERVER].items()
+        >= {
+            API_MAC: REDACTED,
+            API_WIFI_RSSI: HVAC_WEBSERVER_MOCK[API_WIFI_RSSI],
+        }.items()
+    )
 
-    assert diag["data"][AZD_ZONES]["1:3"][AZD_ID] == 3
-    assert diag["data"][AZD_ZONES]["1:3"][AZD_MASTER] == 0
-    assert diag["data"][AZD_ZONES]["1:3"][AZD_SYSTEM] == 1
+    assert (
+        diag["config_entry"].items()
+        >= {
+            "data": {
+                CONF_HOST: CONFIG[CONF_HOST],
+                CONF_PORT: CONFIG[CONF_PORT],
+            },
+            "domain": DOMAIN,
+            "unique_id": REDACTED,
+        }.items()
+    )
 
-    assert diag["data"][AZD_ZONES]["1:4"][AZD_ID] == 4
-    assert diag["data"][AZD_ZONES]["1:4"][AZD_MASTER] == 0
-    assert diag["data"][AZD_ZONES]["1:4"][AZD_SYSTEM] == 1
+    assert (
+        diag["coord_data"][AZD_SYSTEMS]["1"].items()
+        >= {
+            AZD_ID: 1,
+            AZD_ZONES_NUM: 5,
+        }.items()
+    )
 
-    assert diag["data"][AZD_ZONES]["1:5"][AZD_ID] == 5
-    assert diag["data"][AZD_ZONES]["1:5"][AZD_MASTER] == 0
-    assert diag["data"][AZD_ZONES]["1:5"][AZD_SYSTEM] == 1
+    assert (
+        diag["coord_data"][AZD_ZONES]["1:1"].items()
+        >= {
+            AZD_ID: 1,
+            AZD_MASTER: True,
+            AZD_SYSTEM: 1,
+        }.items()
+    )