From f2207af1c96b918d0b4bd365b2730c731da3557e Mon Sep 17 00:00:00 2001
From: Aaron Bach <bachya1208@gmail.com>
Date: Tue, 11 Oct 2022 09:26:44 -0600
Subject: [PATCH] Use `entry.as_dict()` in Ambient PWS diagnostics (#80111)

---
 .../components/ambient_station/diagnostics.py | 11 +--
 .../ambient_station/test_diagnostics.py       | 75 +++++++++++--------
 2 files changed, 48 insertions(+), 38 deletions(-)

diff --git a/homeassistant/components/ambient_station/diagnostics.py b/homeassistant/components/ambient_station/diagnostics.py
index 6005b206954..d18047fe8e4 100644
--- a/homeassistant/components/ambient_station/diagnostics.py
+++ b/homeassistant/components/ambient_station/diagnostics.py
@@ -5,7 +5,7 @@ from typing import Any
 
 from homeassistant.components.diagnostics import async_redact_data
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_API_KEY, CONF_LOCATION
+from homeassistant.const import CONF_API_KEY, CONF_LOCATION, CONF_UNIQUE_ID
 from homeassistant.core import HomeAssistant
 
 from . import AmbientStation
@@ -16,6 +16,7 @@ CONF_APP_KEY_CAMEL = "appKey"
 CONF_DEVICE_ID_CAMEL = "deviceId"
 CONF_MAC_ADDRESS = "mac_address"
 CONF_MAC_ADDRESS_CAMEL = "macAddress"
+CONF_TITLE = "title"
 CONF_TZ = "tz"
 
 TO_REDACT = {
@@ -28,6 +29,9 @@ TO_REDACT = {
     CONF_MAC_ADDRESS,
     CONF_MAC_ADDRESS_CAMEL,
     CONF_TZ,
+    # Config entry title and unique ID may contain sensitive data:
+    CONF_TITLE,
+    CONF_UNIQUE_ID,
 }
 
 
@@ -38,9 +42,6 @@ async def async_get_config_entry_diagnostics(
     ambient: AmbientStation = hass.data[DOMAIN][entry.entry_id]
 
     return {
-        "entry": {
-            "title": entry.title,
-            "data": async_redact_data(entry.data, TO_REDACT),
-        },
+        "entry": async_redact_data(entry.as_dict(), TO_REDACT),
         "stations": async_redact_data(ambient.stations, TO_REDACT),
     }
diff --git a/tests/components/ambient_station/test_diagnostics.py b/tests/components/ambient_station/test_diagnostics.py
index 63d5fcff7a1..e6285afa17a 100644
--- a/tests/components/ambient_station/test_diagnostics.py
+++ b/tests/components/ambient_station/test_diagnostics.py
@@ -13,45 +13,54 @@ async def test_entry_diagnostics(
     ambient.stations = station_data
     assert await get_diagnostics_for_config_entry(hass, hass_client, config_entry) == {
         "entry": {
+            "entry_id": config_entry.entry_id,
+            "version": 2,
+            "domain": "ambient_station",
+            "title": REDACTED,
             "data": {"api_key": REDACTED, "app_key": REDACTED},
-            "title": "Mock Title",
+            "options": {},
+            "pref_disable_new_entities": False,
+            "pref_disable_polling": False,
+            "source": "user",
+            "unique_id": REDACTED,
+            "disabled_by": None,
         },
         "stations": {
             "devices": [
                 {
-                    "apiKey": REDACTED,
-                    "info": {"location": REDACTED, "name": "Side Yard"},
-                    "lastData": {
-                        "baromabsin": 25.016,
-                        "baromrelin": 29.953,
-                        "batt_co2": 1,
-                        "dailyrainin": 0,
-                        "date": "2022-01-19T22:38:00.000Z",
-                        "dateutc": 1642631880000,
-                        "deviceId": REDACTED,
-                        "dewPoint": 17.75,
-                        "dewPointin": 37,
-                        "eventrainin": 0,
-                        "feelsLike": 21,
-                        "feelsLikein": 69.1,
-                        "hourlyrainin": 0,
-                        "humidity": 87,
-                        "humidityin": 29,
-                        "lastRain": "2022-01-07T19:45:00.000Z",
-                        "maxdailygust": 9.2,
-                        "monthlyrainin": 0.409,
-                        "solarradiation": 11.62,
-                        "tempf": 21,
-                        "tempinf": 70.9,
-                        "totalrainin": 35.398,
-                        "tz": REDACTED,
-                        "uv": 0,
-                        "weeklyrainin": 0,
-                        "winddir": 25,
-                        "windgustmph": 1.1,
-                        "windspeedmph": 0.2,
-                    },
                     "macAddress": REDACTED,
+                    "lastData": {
+                        "dateutc": 1642631880000,
+                        "tempinf": 70.9,
+                        "humidityin": 29,
+                        "baromrelin": 29.953,
+                        "baromabsin": 25.016,
+                        "tempf": 21,
+                        "humidity": 87,
+                        "winddir": 25,
+                        "windspeedmph": 0.2,
+                        "windgustmph": 1.1,
+                        "maxdailygust": 9.2,
+                        "hourlyrainin": 0,
+                        "eventrainin": 0,
+                        "dailyrainin": 0,
+                        "weeklyrainin": 0,
+                        "monthlyrainin": 0.409,
+                        "totalrainin": 35.398,
+                        "solarradiation": 11.62,
+                        "uv": 0,
+                        "batt_co2": 1,
+                        "feelsLike": 21,
+                        "dewPoint": 17.75,
+                        "feelsLikein": 69.1,
+                        "dewPointin": 37,
+                        "lastRain": "2022-01-07T19:45:00.000Z",
+                        "deviceId": REDACTED,
+                        "tz": REDACTED,
+                        "date": "2022-01-19T22:38:00.000Z",
+                    },
+                    "info": {"name": "Side Yard", "location": REDACTED},
+                    "apiKey": REDACTED,
                 }
             ],
             "method": "subscribe",