Add zwave_js add-on info dataclass (#50776)

homeassistant/components/zwave_js/__init__.py

@@ -29,7 +29,7 @@ from homeassistant.helpers import device_registry, entity_registry
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 
-from .addon import AddonError, AddonManager, get_addon_manager
+from .addon import AddonError, AddonManager, AddonState, get_addon_manager
 from .api import async_register_api
 from .const import (
     ATTR_COMMAND_CLASS,
@@ -559,22 +559,22 @@ async def async_ensure_addon_running(hass: HomeAssistant, entry: ConfigEntry) ->
     if addon_manager.task_in_progress():
         raise ConfigEntryNotReady
     try:
-        addon_is_installed = await addon_manager.async_is_addon_installed()
+        addon_info = await addon_manager.async_get_addon_info()
-        addon_is_running = await addon_manager.async_is_addon_running()
     except AddonError as err:
-        LOGGER.error("Failed to get the Z-Wave JS add-on info")
+        LOGGER.error(err)
         raise ConfigEntryNotReady from err
 
     usb_path: str = entry.data[CONF_USB_PATH]
     network_key: str = entry.data[CONF_NETWORK_KEY]
+    addon_state = addon_info.state
 
-    if not addon_is_installed:
+    if addon_state == AddonState.NOT_INSTALLED:
         addon_manager.async_schedule_install_setup_addon(
             usb_path, network_key, catch_error=True
         )
         raise ConfigEntryNotReady
 
-    if not addon_is_running:
+    if addon_state == AddonState.NOT_RUNNING:
         addon_manager.async_schedule_setup_addon(
             usb_path, network_key, catch_error=True
         )

homeassistant/components/zwave_js/addon.py

@@ -2,6 +2,8 @@
 from __future__ import annotations
 
 import asyncio
+from dataclasses import dataclass
+from enum import Enum
 from functools import partial
 from typing import Any, Callable, TypeVar, cast
 
@@ -55,6 +57,26 @@ def api_error(error_message: str) -> Callable[[F], F]:
     return handle_hassio_api_error
 
 
+@dataclass
+class AddonInfo:
+    """Represent the current add-on info state."""
+
+    options: dict[str, Any]
+    state: AddonState
+    update_available: bool
+    version: str | None
+
+
+class AddonState(Enum):
+    """Represent the current state of the add-on."""
+
+    NOT_INSTALLED = "not_installed"
+    INSTALLING = "installing"
+    UPDATING = "updating"
+    NOT_RUNNING = "not_running"
+    RUNNING = "running"
+
+
 class AddonManager:
     """Manage the add-on.
 
@@ -93,25 +115,32 @@ class AddonManager:
         return discovery_info_config
 
     @api_error("Failed to get the Z-Wave JS add-on info")
-    async def async_get_addon_info(self) -> dict:
+    async def async_get_addon_info(self) -> AddonInfo:
         """Return and cache Z-Wave JS add-on info."""
         addon_info: dict = await async_get_addon_info(self._hass, ADDON_SLUG)
-        return addon_info
+        addon_state = self.async_get_addon_state(addon_info)
+        return AddonInfo(
+            options=addon_info["options"],
+            state=addon_state,
+            update_available=addon_info["update_available"],
+            version=addon_info["version"],
+        )
 
-    async def async_is_addon_running(self) -> bool:
+    @callback
-        """Return True if Z-Wave JS add-on is running."""
+    def async_get_addon_state(self, addon_info: dict[str, Any]) -> AddonState:
-        addon_info = await self.async_get_addon_info()
+        """Return the current state of the Z-Wave JS add-on."""
-        return bool(addon_info["state"] == "started")
+        addon_state = AddonState.NOT_INSTALLED
 
-    async def async_is_addon_installed(self) -> bool:
+        if addon_info["version"] is not None:
-        """Return True if Z-Wave JS add-on is installed."""
+            addon_state = AddonState.NOT_RUNNING
-        addon_info = await self.async_get_addon_info()
+        if addon_info["state"] == "started":
-        return addon_info["version"] is not None
+            addon_state = AddonState.RUNNING
+        if self._install_task and not self._install_task.done():
+            addon_state = AddonState.INSTALLING
+        if self._update_task and not self._update_task.done():
+            addon_state = AddonState.UPDATING
 
-    async def async_get_addon_options(self) -> dict:
+        return addon_state
-        """Get Z-Wave JS add-on options."""
-        addon_info = await self.async_get_addon_info()
-        return cast(dict, addon_info["options"])
 
     @api_error("Failed to set the Z-Wave JS add-on options")
     async def async_set_addon_options(self, config: dict) -> None:
@@ -164,13 +193,11 @@ class AddonManager:
     async def async_update_addon(self) -> None:
         """Update the Z-Wave JS add-on if needed."""
         addon_info = await self.async_get_addon_info()
-        addon_version = addon_info["version"]
-        update_available = addon_info["update_available"]
 
-        if addon_version is None:
+        if addon_info.version is None:
             raise AddonError("Z-Wave JS add-on is not installed")
 
-        if not update_available:
+        if not addon_info.update_available:
             return
 
         await self.async_create_snapshot()
@@ -215,14 +242,14 @@ class AddonManager:
 
     async def async_configure_addon(self, usb_path: str, network_key: str) -> None:
         """Configure and start Z-Wave JS add-on."""
-        addon_options = await self.async_get_addon_options()
+        addon_info = await self.async_get_addon_info()
 
         new_addon_options = {
             CONF_ADDON_DEVICE: usb_path,
             CONF_ADDON_NETWORK_KEY: network_key,
         }
 
-        if new_addon_options != addon_options:
+        if new_addon_options != addon_info.options:
             await self.async_set_addon_options(new_addon_options)
 
     @callback
@@ -246,8 +273,7 @@ class AddonManager:
     async def async_create_snapshot(self) -> None:
         """Create a partial snapshot of the Z-Wave JS add-on."""
         addon_info = await self.async_get_addon_info()
-        addon_version = addon_info["version"]
+        name = f"addon_{ADDON_SLUG}_{addon_info.version}"
-        name = f"addon_{ADDON_SLUG}_{addon_version}"
 
         LOGGER.debug("Creating snapshot: %s", name)
         await async_create_snapshot(

homeassistant/components/zwave_js/config_flow.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 
 import asyncio
 import logging
-from typing import Any, cast
+from typing import Any
 
 import aiohttp
 from async_timeout import timeout
@@ -17,7 +17,7 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.data_entry_flow import AbortFlow, FlowResult
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 
-from .addon import AddonError, AddonManager, get_addon_manager
+from .addon import AddonError, AddonInfo, AddonManager, AddonState, get_addon_manager
 from .const import (
     CONF_ADDON_DEVICE,
     CONF_ADDON_NETWORK_KEY,
@@ -187,13 +187,15 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
         self.use_addon = True
 
-        if await self._async_is_addon_running():
+        addon_info = await self._async_get_addon_info()
-            addon_config = await self._async_get_addon_config()
+
+        if addon_info.state == AddonState.RUNNING:
+            addon_config = addon_info.options
             self.usb_path = addon_config[CONF_ADDON_DEVICE]
             self.network_key = addon_config.get(CONF_ADDON_NETWORK_KEY, "")
             return await self.async_step_finish_addon_setup()
 
-        if await self._async_is_addon_installed():
+        if addon_info.state == AddonState.NOT_RUNNING:
             return await self.async_step_configure_addon()
 
         return await self.async_step_install_addon()
@@ -228,7 +230,8 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         self, user_input: dict[str, Any] | None = None
     ) -> FlowResult:
         """Ask for config for Z-Wave JS add-on."""
-        addon_config = await self._async_get_addon_config()
+        addon_info = await self._async_get_addon_info()
+        addon_config = addon_info.options
 
         errors: dict[str, str] = {}
 
@@ -345,32 +348,17 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
         )
         return self._async_create_entry_from_vars()
 
-    async def _async_get_addon_info(self) -> dict:
+    async def _async_get_addon_info(self) -> AddonInfo:
         """Return and cache Z-Wave JS add-on info."""
         addon_manager: AddonManager = get_addon_manager(self.hass)
         try:
-            addon_info: dict = await addon_manager.async_get_addon_info()
+            addon_info: AddonInfo = await addon_manager.async_get_addon_info()
         except AddonError as err:
             _LOGGER.error(err)
             raise AbortFlow("addon_info_failed") from err
 
         return addon_info
 
-    async def _async_is_addon_running(self) -> bool:
-        """Return True if Z-Wave JS add-on is running."""
-        addon_info = await self._async_get_addon_info()
-        return bool(addon_info["state"] == "started")
-
-    async def _async_is_addon_installed(self) -> bool:
-        """Return True if Z-Wave JS add-on is installed."""
-        addon_info = await self._async_get_addon_info()
-        return addon_info["version"] is not None
-
-    async def _async_get_addon_config(self) -> dict:
-        """Get Z-Wave JS add-on config."""
-        addon_info = await self._async_get_addon_info()
-        return cast(dict, addon_info["options"])
-
     async def _async_set_addon_config(self, config: dict) -> None:
         """Set Z-Wave JS add-on config."""
         addon_manager: AddonManager = get_addon_manager(self.hass)

tests/components/zwave_js/conftest.py

@@ -30,7 +30,12 @@ def mock_addon_info(addon_info_side_effect):
         "homeassistant.components.zwave_js.addon.async_get_addon_info",
         side_effect=addon_info_side_effect,
     ) as addon_info:
-        addon_info.return_value = {}
+        addon_info.return_value = {
+            "options": {},
+            "state": None,
+            "update_available": False,
+            "version": None,
+        }
         yield addon_info
 
 
@@ -52,7 +57,6 @@ def mock_addon_installed(addon_info):
 @pytest.fixture(name="addon_options")
 def mock_addon_options(addon_info):
     """Mock add-on options."""
-    addon_info.return_value["options"] = {}
     return addon_info.return_value["options"]