Migrate discovery apis to aiohasupervisor

Author: mdegat01

homeassistant/components/hassio/__init__.py

@@ -101,7 +101,6 @@
     HassIO,
     HassioAPIError,
     async_create_backup,
-    async_get_addon_discovery_info,
     async_get_green_settings,
     async_get_yellow_settings,
     async_reboot_host,

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/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(
@@ -332,22 +323,6 @@ def refresh_updates(self) -> 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.

tests/components/conftest.py

@@ -8,7 +8,7 @@
 from typing import TYPE_CHECKING, Any
 from unittest.mock import AsyncMock, MagicMock, patch
 
-from aiohasupervisor.models import Repository, StoreAddon, StoreInfo
+from aiohasupervisor.models import Discovery, Repository, StoreAddon, StoreInfo
 import pytest
 
 from homeassistant.config_entries import (
@@ -205,12 +205,9 @@ def addon_manager_fixture(
 
 
 @pytest.fixture(name="discovery_info")
-def discovery_info_fixture() -> Any:
+def discovery_info_fixture() -> list[Discovery]:
     """Return the discovery info from the supervisor."""
-    # pylint: disable-next=import-outside-toplevel
-    from .hassio.common import mock_discovery_info
-
-    return mock_discovery_info()
+    return []
 
 
 @pytest.fixture(name="discovery_info_side_effect")
@@ -221,13 +218,29 @@ def discovery_info_side_effect_fixture() -> Any | None:
 
 @pytest.fixture(name="get_addon_discovery_info")
 def get_addon_discovery_info_fixture(
-    discovery_info: dict[str, Any], discovery_info_side_effect: Any | None
-) -> Generator[AsyncMock]:
+    supervisor_client: AsyncMock,
+    discovery_info: list[Discovery],
+    discovery_info_side_effect: Any | None,
+) -> AsyncMock:
     """Mock get add-on discovery info."""
-    # pylint: disable-next=import-outside-toplevel
-    from .hassio.common import mock_get_addon_discovery_info
+    supervisor_client.discovery.list.return_value = discovery_info
+    supervisor_client.discovery.list.side_effect = discovery_info_side_effect
+    return supervisor_client.discovery.list
 
-    yield from mock_get_addon_discovery_info(discovery_info, discovery_info_side_effect)
+
+@pytest.fixture(name="get_discovery_message_side_effect")
+def get_discovery_message_side_effect_fixture() -> Any | None:
+    """Side effect for getting a discovery message by uuid."""
+    return None
+
+
+@pytest.fixture(name="get_discovery_message")
+def get_discovery_message_fixture(
+    supervisor_client: AsyncMock, get_discovery_message_side_effect: Any | None
+) -> AsyncMock:
+    """Mock getting a discovery message by uuid."""
+    supervisor_client.discovery.get.side_effect = get_discovery_message_side_effect
+    return supervisor_client.discovery.get
 
 
 @pytest.fixture(name="addon_store_info_side_effect")
@@ -456,6 +469,7 @@ def supervisor_client() -> Generator[AsyncMock]:
     """Mock the supervisor client."""
     supervisor_client = AsyncMock()
     supervisor_client.addons = AsyncMock()
+    supervisor_client.discovery = AsyncMock()
     with (
         patch(
             "homeassistant.components.hassio.get_supervisor_client",

tests/components/hassio/common.py

@@ -7,7 +7,7 @@
 import logging
 from types import MethodType
 from typing import Any
-from unittest.mock import DEFAULT, AsyncMock, Mock, patch
+from unittest.mock import AsyncMock, Mock, patch
 
 from aiohasupervisor.models import (
     AddonsOptions,
@@ -75,23 +75,6 @@ def mock_addon_manager(hass: HomeAssistant) -> AddonManager:
     return AddonManager(hass, LOGGER, "Test", "test_addon")
 
 
-def mock_discovery_info() -> Any:
-    """Return the discovery info from the supervisor."""
-    return DEFAULT
-
-
-def mock_get_addon_discovery_info(
-    discovery_info: dict[str, Any], discovery_info_side_effect: Any | None
-) -> Generator[AsyncMock]:
-    """Mock get add-on discovery info."""
-    with patch(
-        "homeassistant.components.hassio.addon_manager.async_get_addon_discovery_info",
-        side_effect=discovery_info_side_effect,
-        return_value=discovery_info,
-    ) as get_addon_discovery_info:
-        yield get_addon_discovery_info
-
-
 def mock_addon_store_info(
     supervisor_client: AsyncMock,
     addon_store_info_side_effect: Any | None,

tests/components/hassio/test_addon_manager.py

@@ -5,9 +5,10 @@
 import asyncio
 from typing import Any
 from unittest.mock import AsyncMock, call
+from uuid import uuid4
 
 from aiohasupervisor import SupervisorError
-from aiohasupervisor.models import AddonsOptions
+from aiohasupervisor.models import AddonsOptions, Discovery
 import pytest
 
 from homeassistant.components.hassio.addon_manager import (
@@ -62,7 +63,11 @@ async def test_get_addon_discovery_info(
     addon_manager: AddonManager, get_addon_discovery_info: AsyncMock
 ) -> None:
     """Test get addon discovery info."""
-    get_addon_discovery_info.return_value = {"config": {"test_key": "test"}}
+    get_addon_discovery_info.return_value = [
+        Discovery(
+            addon="test_addon", service="", uuid=uuid4(), config={"test_key": "test"}
+        )
+    ]
 
     assert await addon_manager.async_get_addon_discovery_info() == {"test_key": "test"}
 
@@ -73,8 +78,6 @@ async def test_missing_addon_discovery_info(
     addon_manager: AddonManager, get_addon_discovery_info: AsyncMock
 ) -> None:
     """Test missing addon discovery info."""
-    get_addon_discovery_info.return_value = None
-
     with pytest.raises(AddonError):
         await addon_manager.async_get_addon_discovery_info()
 
@@ -85,7 +88,7 @@ async def test_get_addon_discovery_info_error(
     addon_manager: AddonManager, get_addon_discovery_info: AsyncMock
 ) -> None:
     """Test get addon discovery info raises error."""
-    get_addon_discovery_info.side_effect = HassioAPIError("Boom")
+    get_addon_discovery_info.side_effect = SupervisorError("Boom")
 
     with pytest.raises(AddonError) as err:
         assert await addon_manager.async_get_addon_discovery_info()