From 687987f05b9a456f2ca0799c1b4c256b462a7252 Mon Sep 17 00:00:00 2001
From: Aaron Bach <bachya1208@gmail.com>
Date: Tue, 11 Oct 2022 12:15:07 -0600
Subject: [PATCH] Use `entry.as_dict()` in RainMachine diagnostics (#80118)

* Use `entry.as_dict()` in RainMachine diagnostics

* One call
---
 .../components/rainmachine/diagnostics.py     | 25 ++++++++---------
 .../rainmachine/test_diagnostics.py           | 28 ++++++++++++-------
 2 files changed, 30 insertions(+), 23 deletions(-)

diff --git a/homeassistant/components/rainmachine/diagnostics.py b/homeassistant/components/rainmachine/diagnostics.py
index 47ded7990c6..e4835d514e6 100644
--- a/homeassistant/components/rainmachine/diagnostics.py
+++ b/homeassistant/components/rainmachine/diagnostics.py
@@ -12,6 +12,7 @@ from homeassistant.const import (
     CONF_LATITUDE,
     CONF_LONGITUDE,
     CONF_PASSWORD,
+    CONF_UNIQUE_ID,
 )
 from homeassistant.core import HomeAssistant
 
@@ -32,6 +33,8 @@ TO_REDACT = {
     CONF_STATION_NAME,
     CONF_STATION_SOURCE,
     CONF_TIMEZONE,
+    # Config entry unique ID may contain sensitive data:
+    CONF_UNIQUE_ID,
 }
 
 
@@ -47,20 +50,16 @@ async def async_get_config_entry_diagnostics(
         LOGGER.warning("Unable to download controller-specific diagnostics")
         controller_diagnostics = None
 
-    return {
-        "entry": {
-            "title": entry.title,
-            "data": async_redact_data(entry.data, TO_REDACT),
-            "options": dict(entry.options),
-        },
-        "data": {
-            "coordinator": async_redact_data(
-                {
+    return async_redact_data(
+        {
+            "entry": entry.as_dict(),
+            "data": {
+                "coordinator": {
                     api_category: controller.data
                     for api_category, controller in data.coordinators.items()
                 },
-                TO_REDACT,
-            ),
-            "controller_diagnostics": controller_diagnostics,
+                "controller_diagnostics": controller_diagnostics,
+            },
         },
-    }
+        TO_REDACT,
+    )
diff --git a/tests/components/rainmachine/test_diagnostics.py b/tests/components/rainmachine/test_diagnostics.py
index 7cf8406d2ae..a3c03c956a4 100644
--- a/tests/components/rainmachine/test_diagnostics.py
+++ b/tests/components/rainmachine/test_diagnostics.py
@@ -10,6 +10,9 @@ async def test_entry_diagnostics(hass, config_entry, hass_client, setup_rainmach
     """Test config entry diagnostics."""
     assert await get_diagnostics_for_config_entry(hass, hass_client, config_entry) == {
         "entry": {
+            "entry_id": config_entry.entry_id,
+            "version": 2,
+            "domain": "rainmachine",
             "title": "Mock Title",
             "data": {
                 "ip_address": "192.168.1.100",
@@ -18,14 +21,15 @@ async def test_entry_diagnostics(hass, config_entry, hass_client, setup_rainmach
                 "ssl": True,
             },
             "options": {},
+            "pref_disable_new_entities": False,
+            "pref_disable_polling": False,
+            "source": "user",
+            "unique_id": REDACTED,
+            "disabled_by": None,
         },
         "data": {
             "coordinator": {
-                "api.versions": {
-                    "apiVer": "4.6.1",
-                    "hwVer": "3",
-                    "swVer": "4.0.1144",
-                },
+                "api.versions": {"apiVer": "4.6.1", "hwVer": "3", "swVer": "4.0.1144"},
                 "machine.firmware_update_status": {
                     "lastUpdateCheckTimestamp": 1657825288,
                     "packageDetails": [],
@@ -628,6 +632,9 @@ async def test_entry_diagnostics_failed_controller_diagnostics(
     controller.diagnostics.current.side_effect = RainMachineError
     assert await get_diagnostics_for_config_entry(hass, hass_client, config_entry) == {
         "entry": {
+            "entry_id": config_entry.entry_id,
+            "version": 2,
+            "domain": "rainmachine",
             "title": "Mock Title",
             "data": {
                 "ip_address": "192.168.1.100",
@@ -636,14 +643,15 @@ async def test_entry_diagnostics_failed_controller_diagnostics(
                 "ssl": True,
             },
             "options": {},
+            "pref_disable_new_entities": False,
+            "pref_disable_polling": False,
+            "source": "user",
+            "unique_id": REDACTED,
+            "disabled_by": None,
         },
         "data": {
             "coordinator": {
-                "api.versions": {
-                    "apiVer": "4.6.1",
-                    "hwVer": "3",
-                    "swVer": "4.0.1144",
-                },
+                "api.versions": {"apiVer": "4.6.1", "hwVer": "3", "swVer": "4.0.1144"},
                 "machine.firmware_update_status": {
                     "lastUpdateCheckTimestamp": 1657825288,
                     "packageDetails": [],