From d369d679c717a1929571fbccc26a14a56259c3e4 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Wed, 14 Jun 2023 21:16:39 -1000 Subject: [PATCH] Fix ESPHome entries reloading after startup when dashboard is in use (#94362) --- homeassistant/components/esphome/__init__.py | 11 +- homeassistant/components/esphome/dashboard.py | 171 +++++++++++++----- tests/components/esphome/test_dashboard.py | 81 +++++++++ 3 files changed, 220 insertions(+), 43 deletions(-) diff --git a/homeassistant/components/esphome/__init__.py b/homeassistant/components/esphome/__init__.py index 41b1b780a1a..5e113aff86f 100644 --- a/homeassistant/components/esphome/__init__.py +++ b/homeassistant/components/esphome/__init__.py @@ -56,10 +56,11 @@ from homeassistant.helpers.issue_registry import ( ) from homeassistant.helpers.service import async_set_service_schema from homeassistant.helpers.template import Template +from homeassistant.helpers.typing import ConfigType from .bluetooth import async_connect_scanner from .const import DOMAIN -from .dashboard import async_get_dashboard +from .dashboard import async_get_dashboard, async_setup as async_setup_dashboard from .domain_data import DomainData # Import config flow so that it's added to the registry @@ -79,6 +80,8 @@ PROJECT_URLS = { } DEFAULT_URL = f"https://esphome.io/changelog/{STABLE_BLE_VERSION_STR}.html" +CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN) + @callback def _async_check_firmware_version( @@ -135,6 +138,12 @@ def _async_check_using_api_password( ) +async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: + """Set up the esphome component.""" + await async_setup_dashboard(hass) + return True + + async def async_setup_entry( # noqa: C901 hass: HomeAssistant, entry: ConfigEntry ) -> bool: diff --git a/homeassistant/components/esphome/dashboard.py b/homeassistant/components/esphome/dashboard.py index a8332f8d040..35e9cf74555 100644 --- a/homeassistant/components/esphome/dashboard.py +++ b/homeassistant/components/esphome/dashboard.py @@ -4,6 +4,7 @@ from __future__ import annotations import asyncio from datetime import timedelta import logging +from typing import Any import aiohttp from awesomeversion import AwesomeVersion @@ -11,62 +12,148 @@ from esphome_dashboard_api import ConfiguredDevice, ESPHomeDashboardAPI from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntryState from homeassistant.const import EVENT_HOMEASSISTANT_STOP -from homeassistant.core import Event, HomeAssistant, callback +from homeassistant.core import CALLBACK_TYPE, Event, HomeAssistant, callback from homeassistant.helpers.aiohttp_client import async_get_clientsession -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed +from homeassistant.helpers.singleton import singleton +from homeassistant.helpers.storage import Store +from homeassistant.helpers.update_coordinator import DataUpdateCoordinator from .const import DOMAIN -KEY_DASHBOARD = "esphome_dashboard" +_LOGGER = logging.getLogger(__name__) + + +KEY_DASHBOARD_MANAGER = "esphome_dashboard_manager" + +STORAGE_KEY = "esphome.dashboard" +STORAGE_VERSION = 1 + + +async def async_setup(hass: HomeAssistant) -> None: + """Set up the ESPHome dashboard.""" + # Try to restore the dashboard manager from storage + # to avoid reloading every ESPHome config entry after + # Home Assistant starts and the dashboard is discovered. + await async_get_or_create_dashboard_manager(hass) + + +@singleton(KEY_DASHBOARD_MANAGER) +async def async_get_or_create_dashboard_manager( + hass: HomeAssistant, +) -> ESPHomeDashboardManager: + """Get the dashboard manager or create it.""" + manager = ESPHomeDashboardManager(hass) + await manager.async_setup() + return manager + + +class ESPHomeDashboardManager: + """Class to manage the dashboard and restore it from storage.""" + + def __init__(self, hass: HomeAssistant) -> None: + """Initialize the dashboard manager.""" + self._hass = hass + self._store: Store[dict[str, Any]] = Store(hass, STORAGE_VERSION, STORAGE_KEY) + self._data: dict[str, Any] | None = None + self._current_dashboard: ESPHomeDashboard | None = None + self._cancel_shutdown: CALLBACK_TYPE | None = None + + async def async_setup(self) -> None: + """Restore the dashboard from storage.""" + self._data = await self._store.async_load() + if (data := self._data) and (info := data.get("info")): + await self.async_set_dashboard_info( + info["addon_slug"], info["host"], info["port"] + ) + + @callback + def async_get(self) -> ESPHomeDashboard | None: + """Get the current dashboard.""" + return self._current_dashboard + + async def async_set_dashboard_info( + self, addon_slug: str, host: str, port: int + ) -> None: + """Set the dashboard info.""" + url = f"http://{host}:{port}" + hass = self._hass + + if cur_dashboard := self._current_dashboard: + if cur_dashboard.addon_slug == addon_slug and cur_dashboard.url == url: + # Do nothing if we already have this data. + return + # Clear and make way for new dashboard + await cur_dashboard.async_shutdown() + if self._cancel_shutdown is not None: + self._cancel_shutdown() + self._cancel_shutdown = None + self._current_dashboard = None + + dashboard = ESPHomeDashboard( + hass, addon_slug, url, async_get_clientsession(hass) + ) + await dashboard.async_request_refresh() + if not cur_dashboard and not dashboard.last_update_success: + # If there was no previous dashboard and the new one is not available, + # we skip setup and wait for discovery. + _LOGGER.error( + "Dashboard unavailable; skipping setup: %s", dashboard.last_exception + ) + return + + self._current_dashboard = dashboard + + async def on_hass_stop(_: Event) -> None: + await dashboard.async_shutdown() + + self._cancel_shutdown = hass.bus.async_listen_once( + EVENT_HOMEASSISTANT_STOP, on_hass_stop + ) + + new_data = {"info": {"addon_slug": addon_slug, "host": host, "port": port}} + if self._data != new_data: + await self._store.async_save(new_data) + + reloads = [ + hass.config_entries.async_reload(entry.entry_id) + for entry in hass.config_entries.async_entries(DOMAIN) + if entry.state == ConfigEntryState.LOADED + ] + # Re-auth flows will check the dashboard for encryption key when the form is requested + # but we only trigger reauth if the dashboard is available. + if dashboard.last_update_success: + reauths = [ + hass.config_entries.flow.async_configure(flow["flow_id"]) + for flow in hass.config_entries.flow.async_progress() + if flow["handler"] == DOMAIN + and flow["context"]["source"] == SOURCE_REAUTH + ] + else: + reauths = [] + _LOGGER.error( + "Dashboard unavailable; skipping reauth: %s", dashboard.last_exception + ) + + _LOGGER.debug( + "Reloading %d and re-authenticating %d", len(reloads), len(reauths) + ) + if reloads or reauths: + await asyncio.gather(*reloads, *reauths) @callback def async_get_dashboard(hass: HomeAssistant) -> ESPHomeDashboard | None: """Get an instance of the dashboard if set.""" - return hass.data.get(KEY_DASHBOARD) + manager: ESPHomeDashboardManager | None = hass.data.get(KEY_DASHBOARD_MANAGER) + return manager.async_get() if manager else None async def async_set_dashboard_info( hass: HomeAssistant, addon_slug: str, host: str, port: int ) -> None: """Set the dashboard info.""" - url = f"http://{host}:{port}" - - if cur_dashboard := async_get_dashboard(hass): - if cur_dashboard.addon_slug == addon_slug and cur_dashboard.url == url: - # Do nothing if we already have this data. - return - # Clear and make way for new dashboard - await cur_dashboard.async_shutdown() - del hass.data[KEY_DASHBOARD] - - dashboard = ESPHomeDashboard(hass, addon_slug, url, async_get_clientsession(hass)) - try: - await dashboard.async_request_refresh() - except UpdateFailed as err: - logging.getLogger(__name__).error("Ignoring dashboard info: %s", err) - return - - hass.data[KEY_DASHBOARD] = dashboard - - async def on_hass_stop(_: Event) -> None: - await dashboard.async_shutdown() - - hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, on_hass_stop) - - reloads = [ - hass.config_entries.async_reload(entry.entry_id) - for entry in hass.config_entries.async_entries(DOMAIN) - if entry.state == ConfigEntryState.LOADED - ] - # Re-auth flows will check the dashboard for encryption key when the form is requested - reauths = [ - hass.config_entries.flow.async_configure(flow["flow_id"]) - for flow in hass.config_entries.flow.async_progress() - if flow["handler"] == DOMAIN and flow["context"]["source"] == SOURCE_REAUTH - ] - if reloads or reauths: - await asyncio.gather(*reloads, *reauths) + manager = await async_get_or_create_dashboard_manager(hass) + await manager.async_set_dashboard_info(addon_slug, host, port) class ESPHomeDashboard(DataUpdateCoordinator[dict[str, ConfiguredDevice]]): @@ -82,7 +169,7 @@ class ESPHomeDashboard(DataUpdateCoordinator[dict[str, ConfiguredDevice]]): """Initialize.""" super().__init__( hass, - logging.getLogger(__name__), + _LOGGER, name="ESPHome Dashboard", update_interval=timedelta(minutes=5), ) diff --git a/tests/components/esphome/test_dashboard.py b/tests/components/esphome/test_dashboard.py index 580e741e03f..d16bf7c4d00 100644 --- a/tests/components/esphome/test_dashboard.py +++ b/tests/components/esphome/test_dashboard.py @@ -1,4 +1,5 @@ """Test ESPHome dashboard features.""" +import asyncio from unittest.mock import patch from aioesphomeapi import DeviceInfo, InvalidAuthAPIError @@ -10,6 +11,86 @@ from homeassistant.data_entry_flow import FlowResultType from . import VALID_NOISE_PSK +from tests.common import MockConfigEntry + + +async def test_dashboard_storage( + hass: HomeAssistant, init_integration, mock_dashboard, hass_storage +) -> None: + """Test dashboard storage.""" + assert hass_storage[dashboard.STORAGE_KEY]["data"] == { + "info": {"addon_slug": "mock-slug", "host": "mock-host", "port": 1234} + } + await dashboard.async_set_dashboard_info(hass, "test-slug", "new-host", 6052) + assert hass_storage[dashboard.STORAGE_KEY]["data"] == { + "info": {"addon_slug": "test-slug", "host": "new-host", "port": 6052} + } + + +async def test_restore_dashboard_storage( + hass: HomeAssistant, mock_config_entry: MockConfigEntry, hass_storage +) -> MockConfigEntry: + """Restore dashboard url and slug from storage.""" + hass_storage[dashboard.STORAGE_KEY] = { + "version": dashboard.STORAGE_VERSION, + "minor_version": dashboard.STORAGE_VERSION, + "key": dashboard.STORAGE_KEY, + "data": {"info": {"addon_slug": "test-slug", "host": "new-host", "port": 6052}}, + } + with patch.object( + dashboard, "async_get_or_create_dashboard_manager" + ) as mock_get_or_create: + await hass.config_entries.async_setup(mock_config_entry.entry_id) + await hass.async_block_till_done() + assert mock_get_or_create.call_count == 1 + + +async def test_setup_dashboard_fails( + hass: HomeAssistant, mock_config_entry: MockConfigEntry, hass_storage +) -> MockConfigEntry: + """Test that nothing is stored on failed dashboard setup when there was no dashboard before.""" + with patch.object( + dashboard.ESPHomeDashboardAPI, "get_devices", side_effect=asyncio.TimeoutError + ) as mock_get_devices: + await hass.config_entries.async_setup(mock_config_entry.entry_id) + await hass.async_block_till_done() + await dashboard.async_set_dashboard_info(hass, "test-slug", "test-host", 6052) + assert mock_config_entry.state == ConfigEntryState.LOADED + assert mock_get_devices.call_count == 1 + + assert dashboard.STORAGE_KEY not in hass_storage + + +async def test_setup_dashboard_fails_when_already_setup( + hass: HomeAssistant, mock_config_entry: MockConfigEntry, hass_storage +) -> MockConfigEntry: + """Test failed dashboard setup still reloads entries if one existed before.""" + with patch.object(dashboard.ESPHomeDashboardAPI, "get_devices") as mock_get_devices: + await dashboard.async_set_dashboard_info( + hass, "test-slug", "working-host", 6052 + ) + await hass.async_block_till_done() + + assert mock_get_devices.call_count == 1 + assert dashboard.STORAGE_KEY in hass_storage + + await hass.config_entries.async_setup(mock_config_entry.entry_id) + await hass.async_block_till_done() + + with patch.object( + dashboard.ESPHomeDashboardAPI, "get_devices", side_effect=asyncio.TimeoutError + ) as mock_get_devices, patch( + "homeassistant.components.esphome.async_setup_entry", return_value=True + ) as mock_setup: + await dashboard.async_set_dashboard_info(hass, "test-slug", "test-host", 6052) + await hass.async_block_till_done() + + assert mock_config_entry.state == ConfigEntryState.LOADED + assert mock_get_devices.call_count == 1 + # We still setup, and reload, but we do not do the reauths + assert dashboard.STORAGE_KEY in hass_storage + assert len(mock_setup.mock_calls) == 1 + async def test_new_info_reload_config_entries( hass: HomeAssistant, init_integration, mock_dashboard