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"])