From cc6b0cc84394f78a326624eb503bf0517f1bdbf4 Mon Sep 17 00:00:00 2001
From: Aaron Bach <bachya1208@gmail.com>
Date: Sat, 29 Jan 2022 13:30:15 -0700
Subject: [PATCH] Ensure diagnostics redaction can handle lists of lists
 (#65170)

* Ensure diagnostics redaction can handle lists of lists

* Code review

* Update homeassistant/components/diagnostics/util.py

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

* Code review

* Typing

* Revert "Typing"

This reverts commit 8a57f772caa5180b609175591d81dfc473769f70.

* New typing attempt

* Revert "New typing attempt"

This reverts commit e26e4aae69f62325fdd6af4d80c8fd1f74846e54.

* Fix typing

* Fix typing again

* Add tests

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>
---
 homeassistant/components/diagnostics/util.py | 11 +++++--
 tests/components/diagnostics/test_util.py    | 33 ++++++++++++++++++++
 2 files changed, 41 insertions(+), 3 deletions(-)
 create mode 100644 tests/components/diagnostics/test_util.py

diff --git a/homeassistant/components/diagnostics/util.py b/homeassistant/components/diagnostics/util.py
index d849dc052e4..6154dd14bd2 100644
--- a/homeassistant/components/diagnostics/util.py
+++ b/homeassistant/components/diagnostics/util.py
@@ -2,19 +2,24 @@
 from __future__ import annotations
 
 from collections.abc import Iterable, Mapping
-from typing import Any
+from typing import Any, TypeVar, cast
 
 from homeassistant.core import callback
 
 from .const import REDACTED
 
+T = TypeVar("T")
+
 
 @callback
-def async_redact_data(data: Mapping, to_redact: Iterable[Any]) -> dict[str, Any]:
+def async_redact_data(data: T, to_redact: Iterable[Any]) -> T:
     """Redact sensitive data in a dict."""
     if not isinstance(data, (Mapping, list)):
         return data
 
+    if isinstance(data, list):
+        return cast(T, [async_redact_data(val, to_redact) for val in data])
+
     redacted = {**data}
 
     for key, value in redacted.items():
@@ -25,4 +30,4 @@ def async_redact_data(data: Mapping, to_redact: Iterable[Any]) -> dict[str, Any]
         elif isinstance(value, list):
             redacted[key] = [async_redact_data(item, to_redact) for item in value]
 
-    return redacted
+    return cast(T, redacted)
diff --git a/tests/components/diagnostics/test_util.py b/tests/components/diagnostics/test_util.py
new file mode 100644
index 00000000000..702b838334f
--- /dev/null
+++ b/tests/components/diagnostics/test_util.py
@@ -0,0 +1,33 @@
+"""Test Diagnostics utils."""
+from homeassistant.components.diagnostics import REDACTED, async_redact_data
+
+
+def test_redact():
+    """Test the async_redact_data helper."""
+    data = {
+        "key1": "value1",
+        "key2": ["value2_a", "value2_b"],
+        "key3": [["value_3a", "value_3b"], ["value_3c", "value_3d"]],
+        "key4": {
+            "key4_1": "value4_1",
+            "key4_2": ["value4_2a", "value4_2b"],
+            "key4_3": [["value4_3a", "value4_3b"], ["value4_3c", "value4_3d"]],
+        },
+    }
+
+    to_redact = {
+        "key1",
+        "key3",
+        "key4_1",
+    }
+
+    assert async_redact_data(data, to_redact) == {
+        "key1": REDACTED,
+        "key2": ["value2_a", "value2_b"],
+        "key3": REDACTED,
+        "key4": {
+            "key4_1": REDACTED,
+            "key4_2": ["value4_2a", "value4_2b"],
+            "key4_3": [["value4_3a", "value4_3b"], ["value4_3c", "value4_3d"]],
+        },
+    }