From 94779dddaae020f856db052f082de9a9fc0b7714 Mon Sep 17 00:00:00 2001
From: epenet <6771947+epenet@users.noreply.github.com>
Date: Thu, 9 Feb 2023 12:15:18 +0100
Subject: [PATCH] Add type hints to diagnostics test helper (#85494)

* Add type hints to diagnostics test helper

* Move type alias to the top

* Use Any

* Move TestClientGenerator to typing helper

* Use `dict[str, Any]` again

* Use JsonObjectType

* Add cast
---
 tests/components/diagnostics/__init__.py | 41 +++++++++++++++++++-----
 1 file changed, 33 insertions(+), 8 deletions(-)

diff --git a/tests/components/diagnostics/__init__.py b/tests/components/diagnostics/__init__.py
index 3dbbb741a9f..e576fbd974f 100644
--- a/tests/components/diagnostics/__init__.py
+++ b/tests/components/diagnostics/__init__.py
@@ -1,10 +1,21 @@
 """Tests for the Diagnostics integration."""
 from http import HTTPStatus
+from typing import cast
 
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.device_registry import DeviceEntry
+from homeassistant.helpers.json import JsonObjectType
 from homeassistant.setup import async_setup_component
 
+from tests.typing import ClientSessionGenerator
 
-async def _get_diagnostics_for_config_entry(hass, hass_client, config_entry):
+
+async def _get_diagnostics_for_config_entry(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    config_entry: ConfigEntry,
+) -> JsonObjectType:
     """Return the diagnostics config entry for the specified domain."""
     assert await async_setup_component(hass, "diagnostics", {})
 
@@ -13,16 +24,25 @@ async def _get_diagnostics_for_config_entry(hass, hass_client, config_entry):
         f"/api/diagnostics/config_entry/{config_entry.entry_id}"
     )
     assert response.status == HTTPStatus.OK
-    return await response.json()
+    return cast(JsonObjectType, await response.json())
 
 
-async def get_diagnostics_for_config_entry(hass, hass_client, config_entry):
+async def get_diagnostics_for_config_entry(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    config_entry: ConfigEntry,
+) -> JsonObjectType:
     """Return the diagnostics config entry for the specified domain."""
     data = await _get_diagnostics_for_config_entry(hass, hass_client, config_entry)
-    return data["data"]
+    return cast(JsonObjectType, data["data"])
 
 
-async def _get_diagnostics_for_device(hass, hass_client, config_entry, device):
+async def _get_diagnostics_for_device(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    config_entry: ConfigEntry,
+    device: DeviceEntry,
+) -> JsonObjectType:
     """Return the diagnostics for the specified device."""
     assert await async_setup_component(hass, "diagnostics", {})
 
@@ -31,10 +51,15 @@ async def _get_diagnostics_for_device(hass, hass_client, config_entry, device):
         f"/api/diagnostics/config_entry/{config_entry.entry_id}/device/{device.id}"
     )
     assert response.status == HTTPStatus.OK
-    return await response.json()
+    return cast(JsonObjectType, await response.json())
 
 
-async def get_diagnostics_for_device(hass, hass_client, config_entry, device):
+async def get_diagnostics_for_device(
+    hass: HomeAssistant,
+    hass_client: ClientSessionGenerator,
+    config_entry: ConfigEntry,
+    device: DeviceEntry,
+) -> JsonObjectType:
     """Return the diagnostics for the specified device."""
     data = await _get_diagnostics_for_device(hass, hass_client, config_entry, device)
-    return data["data"]
+    return cast(JsonObjectType, data["data"])