Continue migration of methods from handler to aiohasupervisor (#129183)

Author: mdegat01

homeassistant/components/hassio/__init__.py

@@ -10,6 +10,7 @@
 import re
 from typing import Any, NamedTuple
 
+from aiohasupervisor import SupervisorError
 import voluptuous as vol
 
 from homeassistant.auth.const import GROUP_ID_ADMIN
@@ -101,16 +102,12 @@
     HassIO,
     HassioAPIError,
     async_create_backup,
-    async_get_addon_discovery_info,
     async_get_green_settings,
     async_get_yellow_settings,
     async_reboot_host,
     async_set_green_settings,
     async_set_yellow_settings,
-    async_update_core,
     async_update_diagnostics,
-    async_update_os,
-    async_update_supervisor,
     get_supervisor_client,
 )
 from .http import HassIOView
@@ -310,8 +307,11 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:  # noqa:
     host = os.environ["SUPERVISOR"]
     websession = async_get_clientsession(hass)
     hass.data[DOMAIN] = hassio = HassIO(hass.loop, websession, host)
+    supervisor_client = get_supervisor_client(hass)
 
-    if not await hassio.is_connected():
+    try:
+        await supervisor_client.supervisor.ping()
+    except SupervisorError:
         _LOGGER.warning("Not connected with the supervisor / system too busy!")
 
     store = Store[dict[str, str]](hass, STORAGE_VERSION, STORAGE_KEY)
@@ -468,9 +468,9 @@ async def update_info_data(_: datetime | None = None) -> None:
     async def _async_stop(hass: HomeAssistant, restart: bool) -> None:
         """Stop or restart home assistant."""
         if restart:
-            await hassio.restart_homeassistant()
+            await supervisor_client.homeassistant.restart()
         else:
-            await hassio.stop_homeassistant()
+            await supervisor_client.homeassistant.stop()
 
     # Set a custom handler for the homeassistant.restart and homeassistant.stop services
     async_set_stop_handler(hass, _async_stop)

homeassistant/components/hassio/addon_manager.py

@@ -21,12 +21,7 @@
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 
-from .handler import (
-    HassioAPIError,
-    async_create_backup,
-    async_get_addon_discovery_info,
-    get_supervisor_client,
-)
+from .handler import HassioAPIError, async_create_backup, get_supervisor_client
 
 type _FuncType[_T, **_P, _R] = Callable[Concatenate[_T, _P], Awaitable[_R]]
 type _ReturnFuncType[_T, **_P, _R] = Callable[
@@ -128,18 +123,25 @@ def task_in_progress(self) -> bool:
             )
         )
 
-    @api_error("Failed to get the {addon_name} add-on discovery info")
+    @api_error(
+        "Failed to get the {addon_name} add-on discovery info",
+        expected_error_type=SupervisorError,
+    )
     async def async_get_addon_discovery_info(self) -> dict:
         """Return add-on discovery info."""
-        discovery_info = await async_get_addon_discovery_info(
-            self._hass, self.addon_slug
+        discovery_info = next(
+            (
+                msg
+                for msg in await self._supervisor_client.discovery.list()
+                if msg.addon == self.addon_slug
+            ),
+            None,
         )
 
         if not discovery_info:
             raise AddonError(f"Failed to get {self.addon_name} add-on discovery info")
 
-        discovery_info_config: dict = discovery_info["config"]
-        return discovery_info_config
+        return discovery_info.config
 
     @api_error(
         "Failed to get the {addon_name} add-on info",

homeassistant/components/hassio/coordinator.py

@@ -563,8 +563,8 @@ async def _async_refresh(
             # updates if this is not a scheduled refresh and
             # we are not doing the first refresh.
             try:
-                await self.hassio.refresh_updates()
-            except HassioAPIError as err:
+                await self.supervisor_client.refresh_updates()
+            except SupervisorError as err:
                 _LOGGER.warning("Error on Supervisor API: %s", err)
 
         await super()._async_refresh(

homeassistant/components/hassio/discovery.py

@@ -8,6 +8,7 @@
 from typing import Any
 
 from aiohasupervisor import SupervisorError
+from aiohasupervisor.models import Discovery
 from aiohttp import web
 from aiohttp.web_exceptions import HTTPServiceUnavailable
 
@@ -19,8 +20,8 @@
 from homeassistant.helpers import discovery_flow
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 
-from .const import ATTR_ADDON, ATTR_CONFIG, ATTR_DISCOVERY, ATTR_UUID, DOMAIN
-from .handler import HassIO, HassioAPIError, get_supervisor_client
+from .const import ATTR_ADDON, ATTR_UUID, DOMAIN
+from .handler import HassIO, get_supervisor_client
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -39,20 +40,21 @@ class HassioServiceInfo(BaseServiceInfo):
 def async_setup_discovery_view(hass: HomeAssistant, hassio: HassIO) -> None:
     """Discovery setup."""
     hassio_discovery = HassIODiscovery(hass, hassio)
+    supervisor_client = get_supervisor_client(hass)
     hass.http.register_view(hassio_discovery)
 
     # Handle exists discovery messages
     async def _async_discovery_start_handler(event: Event) -> None:
         """Process all exists discovery on startup."""
         try:
-            data = await hassio.retrieve_discovery_messages()
-        except HassioAPIError as err:
+            data = await supervisor_client.discovery.list()
+        except SupervisorError as err:
             _LOGGER.error("Can't read discover info: %s", err)
             return
 
         jobs = [
             asyncio.create_task(hassio_discovery.async_process_new(discovery))
-            for discovery in data[ATTR_DISCOVERY]
+            for discovery in data
         ]
         if jobs:
             await asyncio.wait(jobs)
@@ -95,8 +97,8 @@ async def post(self, request: web.Request, uuid: str) -> web.Response:
         """Handle new discovery requests."""
         # Fetch discovery data and prevent injections
         try:
-            data = await self.hassio.get_discovery_message(uuid)
-        except HassioAPIError as err:
+            data = await self._supervisor_client.discovery.get(uuid)
+        except SupervisorError as err:
             _LOGGER.error("Can't read discovery data: %s", err)
             raise HTTPServiceUnavailable from None
 
@@ -113,52 +115,50 @@ async def delete(self, request: web.Request, uuid: str) -> web.Response:
     async def async_rediscover(self, uuid: str) -> None:
         """Rediscover add-on when config entry is removed."""
         try:
-            data = await self.hassio.get_discovery_message(uuid)
-        except HassioAPIError as err:
+            data = await self._supervisor_client.discovery.get(uuid)
+        except SupervisorError as err:
             _LOGGER.debug("Can't read discovery data: %s", err)
         else:
             await self.async_process_new(data)
 
-    async def async_process_new(self, data: dict[str, Any]) -> None:
+    async def async_process_new(self, data: Discovery) -> None:
         """Process add discovery entry."""
-        service: str = data[ATTR_SERVICE]
-        config_data: dict[str, Any] = data[ATTR_CONFIG]
-        slug: str = data[ATTR_ADDON]
-        uuid: str = data[ATTR_UUID]
-
         # Read additional Add-on info
         try:
-            addon_info = await self._supervisor_client.addons.addon_info(slug)
+            addon_info = await self._supervisor_client.addons.addon_info(data.addon)
         except SupervisorError as err:
             _LOGGER.error("Can't read add-on info: %s", err)
             return
 
-        config_data[ATTR_ADDON] = addon_info.name
+        data.config[ATTR_ADDON] = addon_info.name
 
         # Use config flow
         discovery_flow.async_create_flow(
             self.hass,
-            service,
+            data.service,
             context={"source": config_entries.SOURCE_HASSIO},
             data=HassioServiceInfo(
-                config=config_data, name=addon_info.name, slug=slug, uuid=uuid
+                config=data.config,
+                name=addon_info.name,
+                slug=data.addon,
+                uuid=data.uuid,
             ),
             discovery_key=discovery_flow.DiscoveryKey(
                 domain=DOMAIN,
-                key=data[ATTR_UUID],
+                key=data.uuid,
                 version=1,
             ),
         )
 
     async def async_process_del(self, data: dict[str, Any]) -> None:
         """Process remove discovery entry."""
-        service = data[ATTR_SERVICE]
-        uuid = data[ATTR_UUID]
+        service: str = data[ATTR_SERVICE]
+        uuid: str = data[ATTR_UUID]
 
         # Check if really deletet / prevent injections
         try:
-            data = await self.hassio.get_discovery_message(uuid)
-        except HassioAPIError:
+            data = await self._supervisor_client.discovery.get(uuid)
+        except SupervisorError:
             pass
         else:
             _LOGGER.warning("Retrieve wrong unload for %s", service)

homeassistant/components/hassio/handler.py

@@ -24,7 +24,7 @@
 from homeassistant.helpers.singleton import singleton
 from homeassistant.loader import bind_hass
 
-from .const import ATTR_DISCOVERY, ATTR_MESSAGE, ATTR_RESULT, DOMAIN, X_HASS_SOURCE
+from .const import ATTR_MESSAGE, ATTR_RESULT, DOMAIN, X_HASS_SOURCE
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -76,15 +76,6 @@ async def async_update_diagnostics(hass: HomeAssistant, diagnostics: bool) -> bo
     return await hassio.update_diagnostics(diagnostics)
 
 
-@bind_hass
-async def async_get_addon_discovery_info(hass: HomeAssistant, slug: str) -> dict | None:
-    """Return discovery data for an add-on."""
-    hassio: HassIO = hass.data[DOMAIN]
-    data = await hassio.retrieve_discovery_messages()
-    discovered_addons = data[ATTR_DISCOVERY]
-    return next((addon for addon in discovered_addons if addon["addon"] == slug), None)
-
-
 @bind_hass
 @api_data
 async def async_create_backup(
@@ -100,52 +91,6 @@ async def async_create_backup(
     return await hassio.send_command(command, payload=payload, timeout=None)
 
 
-@bind_hass
-@api_data
-async def async_update_os(hass: HomeAssistant, version: str | None = None) -> dict:
-    """Update Home Assistant Operating System.
-
-    The caller of the function should handle HassioAPIError.
-    """
-    hassio: HassIO = hass.data[DOMAIN]
-    command = "/os/update"
-    return await hassio.send_command(
-        command,
-        payload={"version": version},
-        timeout=None,
-    )
-
-
-@bind_hass
-@api_data
-async def async_update_supervisor(hass: HomeAssistant) -> dict:
-    """Update Home Assistant Supervisor.
-
-    The caller of the function should handle HassioAPIError.
-    """
-    hassio: HassIO = hass.data[DOMAIN]
-    command = "/supervisor/update"
-    return await hassio.send_command(command, timeout=None)
-
-
-@bind_hass
-@api_data
-async def async_update_core(
-    hass: HomeAssistant, version: str | None = None, backup: bool = False
-) -> dict:
-    """Update Home Assistant Core.
-
-    The caller of the function should handle HassioAPIError.
-    """
-    hassio: HassIO = hass.data[DOMAIN]
-    command = "/core/update"
-    return await hassio.send_command(
-        command,
-        payload={"version": version, "backup": backup},
-        timeout=None,
-    )
-
-
 @bind_hass
 @_api_bool
 async def async_apply_suggestion(hass: HomeAssistant, suggestion_uuid: str) -> dict:
@@ -228,14 +173,6 @@ def base_url(self) -> URL:
         """Return base url for Supervisor."""
         return self._base_url
 
-    @_api_bool
-    def is_connected(self) -> Coroutine:
-        """Return true if it connected to Hass.io supervisor.
-
-        This method returns a coroutine.
-        """
-        return self.send_command("/supervisor/ping", method="get", timeout=15)
-
     @api_data
     def get_info(self) -> Coroutine:
         """Return generic Supervisor information.
@@ -308,46 +245,6 @@ def get_ingress_panels(self) -> Coroutine:
         """
         return self.send_command("/ingress/panels", method="get")
 
-    @_api_bool
-    def restart_homeassistant(self) -> Coroutine:
-        """Restart Home-Assistant container.
-
-        This method returns a coroutine.
-        """
-        return self.send_command("/homeassistant/restart")
-
-    @_api_bool
-    def stop_homeassistant(self) -> Coroutine:
-        """Stop Home-Assistant container.
-
-        This method returns a coroutine.
-        """
-        return self.send_command("/homeassistant/stop")
-
-    @_api_bool
-    def refresh_updates(self) -> Coroutine:
-        """Refresh available updates.
-
-        This method returns a coroutine.
-        """
-        return self.send_command("/refresh_updates", timeout=300)
-
-    @api_data
-    def retrieve_discovery_messages(self) -> Coroutine:
-        """Return all discovery data from Hass.io API.
-
-        This method returns a coroutine.
-        """
-        return self.send_command("/discovery", method="get", timeout=60)
-
-    @api_data
-    def get_discovery_message(self, uuid: str) -> Coroutine:
-        """Return a single discovery data message.
-
-        This method returns a coroutine.
-        """
-        return self.send_command(f"/discovery/{uuid}", method="get")
-
     @api_data
     def get_resolution_info(self) -> Coroutine:
         """Return data for Supervisor resolution center.