From 155db2f2e122365448265ea002beca7e9b1ea2a9 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Fri, 2 Dec 2022 15:33:55 +0100
Subject: [PATCH] Simplify diagnostics.DownloadDiagnosticsView (#83116)

* Refactor diagnostics.DownloadDiagnosticsView

* Simplify
---
 .../components/diagnostics/__init__.py        | 39 +++++++++++--------
 1 file changed, 22 insertions(+), 17 deletions(-)

diff --git a/homeassistant/components/diagnostics/__init__.py b/homeassistant/components/diagnostics/__init__.py
index 932bef56241..b54a710e807 100644
--- a/homeassistant/components/diagnostics/__init__.py
+++ b/homeassistant/components/diagnostics/__init__.py
@@ -151,9 +151,7 @@ async def _async_get_json_file_response(
     data: Any,
     filename: str,
     domain: str,
-    d_type: DiagnosticsType,
     d_id: str,
-    sub_type: DiagnosticsSubType | None = None,
     sub_id: str | None = None,
 ) -> web.Response:
     """Return JSON file from dictionary."""
@@ -183,9 +181,11 @@ async def _async_get_json_file_response(
     except TypeError:
         _LOGGER.error(
             "Failed to serialize to JSON: %s/%s%s. Bad data at %s",
-            d_type.value,
+            DiagnosticsType.CONFIG_ENTRY.value,
             d_id,
-            f"/{sub_type.value}/{sub_id}" if sub_type is not None else "",
+            f"/{DiagnosticsSubType.DEVICE.value}/{sub_id}"
+            if sub_id is not None
+            else "",
             format_unserializable_data(find_paths_unserializable_data(data)),
         )
         return web.Response(status=HTTPStatus.INTERNAL_SERVER_ERROR)
@@ -213,12 +213,20 @@ class DownloadDiagnosticsView(http.HomeAssistantView):
         sub_id: str | None = None,
     ) -> web.Response:
         """Download diagnostics."""
-        # t_type handling
+        # Validate d_type and sub_type
         try:
-            d_type = DiagnosticsType(d_type)
+            DiagnosticsType(d_type)
         except ValueError:
             return web.Response(status=HTTPStatus.BAD_REQUEST)
 
+        if sub_type is not None:
+            try:
+                DiagnosticsSubType(sub_type)
+            except ValueError:
+                return web.Response(status=HTTPStatus.BAD_REQUEST)
+
+        device_diagnostics = sub_type is not None
+
         hass: HomeAssistant = request.app["hass"]
 
         if (config_entry := hass.config_entries.async_get_entry(d_id)) is None:
@@ -230,23 +238,20 @@ class DownloadDiagnosticsView(http.HomeAssistantView):
 
         filename = f"{config_entry.domain}-{config_entry.entry_id}"
 
-        if sub_type is None:
+        if not device_diagnostics:
+            # Config entry diagnostics
             if info.config_entry_diagnostics is None:
                 return web.Response(status=HTTPStatus.NOT_FOUND)
             data = await info.config_entry_diagnostics(hass, config_entry)
-            filename = f"{d_type}-{filename}"
+            filename = f"{DiagnosticsType.CONFIG_ENTRY}-{filename}"
             return await _async_get_json_file_response(
-                hass, data, filename, config_entry.domain, d_type.value, d_id
+                hass, data, filename, config_entry.domain, d_id
             )
 
-        # sub_type handling
-        try:
-            sub_type = DiagnosticsSubType(sub_type)
-        except ValueError:
-            return web.Response(status=HTTPStatus.BAD_REQUEST)
-
+        # Device diagnostics
         dev_reg = async_get(hass)
-        assert sub_id
+        if sub_id is None:
+            return web.Response(status=HTTPStatus.BAD_REQUEST)
 
         if (device := dev_reg.async_get(sub_id)) is None:
             return web.Response(status=HTTPStatus.NOT_FOUND)
@@ -258,5 +263,5 @@ class DownloadDiagnosticsView(http.HomeAssistantView):
 
         data = await info.device_diagnostics(hass, config_entry, device)
         return await _async_get_json_file_response(
-            hass, data, filename, config_entry.domain, d_type, d_id, sub_type, sub_id
+            hass, data, filename, config_entry.domain, d_id, sub_id
         )