Add support for Hass.io discovery feature for Add-ons

homeassistant/components/hassio/__init__.py

@@ -19,7 +19,8 @@
 from homeassistant.loader import bind_hass
 from homeassistant.util.dt import utcnow
 
-from .handler import HassIO
+from .handler import HassIO, HassioAPIError
+from .discovery import async_setup_discovery
 from .http import HassIOView
 
 _LOGGER = logging.getLogger(__name__)
@@ -136,29 +137,27 @@ def is_hassio(hass):
 async def async_check_config(hass):
     """Check configuration over Hass.io API."""
     hassio = hass.data[DOMAIN]
-    result = await hassio.check_homeassistant_config()
 
-    if not result:
-        return "Hass.io config check API error"
+    try:
+        result = await hassio.check_homeassistant_config()
+    except HassioAPIError as err:
+        _LOGGER.error("Error on Hass.io API: %s", err)
+
     if result['result'] == "error":
         return result['message']
     return None
 
 
 async def async_setup(hass, config):
     """Set up the Hass.io component."""
-    try:
-        host = os.environ['HASSIO']
-    except KeyError:
-        _LOGGER.error("Missing HASSIO environment variable.")
-        return False
-
-    try:
-        os.environ['HASSIO_TOKEN']
-    except KeyError:
-        _LOGGER.error("Missing HASSIO_TOKEN environment variable.")
+    # Check local setup
+    for env in ('HASSIO', 'HASSIO_TOKEN'):
+        if os.environ.get(env):
+            continue
+        _LOGGER.error("Missing %s environment variable.", env)
         return False
 
+    host = os.environ['HASSIO']
     websession = hass.helpers.aiohttp_client.async_get_clientsession()
     hass.data[DOMAIN] = hassio = HassIO(hass.loop, websession, host)
 
@@ -229,23 +228,25 @@ async def async_service_handler(service):
             payload = data
 
         # Call API
-        ret = await hassio.send_command(
-            api_command.format(addon=addon, snapshot=snapshot),
-            payload=payload, timeout=MAP_SERVICE_API[service.service][2]
-        )
-
-        if not ret or ret['result'] != "ok":
-            _LOGGER.error("Error on Hass.io API: %s", ret['message'])
+        try:
+            await hassio.send_command(
+                api_command.format(addon=addon, snapshot=snapshot),
+                payload=payload, timeout=MAP_SERVICE_API[service.service][2]
+            )
+        except HassioAPIError as err:
+            _LOGGER.error("Error on Hass.io API: %s", err)
 
     for service, settings in MAP_SERVICE_API.items():
         hass.services.async_register(
             DOMAIN, service, async_service_handler, schema=settings[1])
 
     async def update_homeassistant_version(now):
         """Update last available Home Assistant version."""
-        data = await hassio.get_homeassistant_info()
-        if data:
+        try:
+            data = await hassio.get_homeassistant_info()
             hass.data[DATA_HOMEASSISTANT_VERSION] = data['last_version']
+        except HassioAPIError as err:
+            _LOGGER.warning("Can't read last version: %s", err)
 
         hass.helpers.event.async_track_point_in_utc_time(
             update_homeassistant_version, utcnow() + HASSIO_UPDATE_INTERVAL)
@@ -276,4 +277,7 @@ async def async_handle_core_service(call):
         hass.services.async_register(
             HASS_DOMAIN, service, async_handle_core_service)
 
+    # Init discovery Hass.io feature
+    async_setup_discovery(hass, hassio, config)
+
     return True

homeassistant/components/hassio/discovery.py

@@ -0,0 +1,115 @@
+"""Implement the serivces discovery feature from Hass.io for Add-ons."""
+import asyncio
+import logging
+
+from aiohttp import web
+from aiohttp.web_exceptions import HTTPServiceUnavailable
+
+from homeassistant.core import callback
+from homeassistant.const import EVENT_HOMEASSISTANT_START
+from homeassistant.components.http import HomeAssistantView
+
+from .handler import HassioAPIError
+
+_LOGGER = logging.getLogger(__name__)
+
+ATTR_DISCOVERY = 'discovery'
+ATTR_ADDON = 'addon'
+ATTR_NAME = 'name'
+ATTR_SERVICE = 'service'
+ATTR_CONFIG = 'config'
+ATTR_UUID = 'uuid'
+
+
+@callback
+def async_setup_discovery(hass, hassio, config):
+    """Discovery setup."""
+    hassio_discovery = HassIODiscovery(hass, hassio, config)
+
+    # Handle exists discovery messages
+    async def async_discovery_start_handler(event):
+        """Process all exists discovery on startup."""
+        try:
+            data = await hassio.retrieve_discovery_messages()
+        except HassioAPIError as err:
+            _LOGGER.error("Can't read discover info: %s", err)
+            return
+
+        jobs = [hassio_discovery.async_process_new(discovery)
+                for discovery in data[ATTR_DISCOVERY]]
+        if jobs:
+            await asyncio.wait(jobs)
+
+    hass.bus.async_listen_once(
+        EVENT_HOMEASSISTANT_START, async_discovery_start_handler)
+
+    hass.http.register_view(hassio_discovery)
+
+
+class HassIODiscovery(HomeAssistantView):
+    """Hass.io view to handle base part."""
+
+    name = "api:hassio_push:discovery"
+    url = "/api/hassio_push/discovery/{uuid}"
+
+    def __init__(self, hass, hassio, config):
+        """Initialize WebView."""
+        self.hass = hass
+        self.hassio = hassio
+        self.config = config
+
+    async def post(self, request, uuid):
+        """Handle new discovery requests."""
+        # Fetch discovery data and prevent injections
+        try:
+            data = await self.hassio.get_discovery_message(uuid)
+        except HassioAPIError as err:
+            _LOGGER.error("Can't read discovey data: %s", err)
+            raise HTTPServiceUnavailable() from None
+
+        await self.async_process_new(data)
+        return web.Response()
+
+    async def delete(self, request, uuid):
+        """Handle remove discovery requests."""
+        data = request.json()
+
+        await self.async_process_del(data)
+        return web.Response()
+
+    async def async_process_new(self, data):
+        """Process add discovery entry."""
+        service = data[ATTR_SERVICE]
+        config_data = data[ATTR_CONFIG]
+
+        # Read addinional Add-on info
+        try:
+            addon_info = await self.hassio.get_addon_info(data[ATTR_ADDON])
+        except HassioAPIError as err:
+            _LOGGER.error("Can't read add-on info: %s", err)
+            return
+        config_data[ATTR_ADDON] = addon_info[ATTR_NAME]
+
+        # Use config flow
+        await self.hass.config_entries.flow.async_init(
+            service, context={'source': 'hassio'}, data=config_data)
+
+    async def async_process_del(self, data):
+        """Process remove discovery entry."""
+        service = data[ATTR_SERVICE]
+        uuid = data[ATTR_UUID]
+
+        # Check if realy deletet / prevent injections
+        try:
+            data = await self.hassio.get_discovery_message(uuid)
+        except HassioAPIError:
+            pass
+        else:
+            _LOGGER.warning("Retrieve wrong unload for %s", service)
+            return
+
+        # Use config flow
+        for entry in self.hass.config_entries.async_entries(service):
+            if entry.source != 'hassio':
+                continue
+            await self.hass.config_entries.async_remove(entry)

homeassistant/components/hassio/handler.py

@@ -21,12 +21,19 @@
 X_HASSIO = 'X-HASSIO-KEY'
 
 
+class HassioAPIError(RuntimeError):
+    """Return if a API trow a error."""
+
+
 def _api_bool(funct):
     """Return a boolean."""
     async def _wrapper(*argv, **kwargs):
         """Wrap function."""
-        data = await funct(*argv, **kwargs)
-        return data and data['result'] == "ok"
+        try:
+            data = await funct(*argv, **kwargs)
+            return data['result'] == "ok"
+        except HassioAPIError:
+            return False
 
     return _wrapper
 
@@ -36,9 +43,9 @@ def _api_data(funct):
     async def _wrapper(*argv, **kwargs):
         """Wrap function."""
         data = await funct(*argv, **kwargs)
-        if data and data['result'] == "ok":
+        if data['result'] == "ok":
             return data['data']
-        return None
+        raise HassioAPIError(data['message'])
 
     return _wrapper
 
@@ -68,6 +75,15 @@ def get_homeassistant_info(self):
         """
         return self.send_command("/homeassistant/info", method="get")
 
+    @_api_data
+    def get_addon_info(self, addon):
+        """Return data for a Add-on.
+
+        This method return a coroutine.
+        """
+        return self.send_command(
+            "/addons/{}/info".format(addon), method="get")
+
     @_api_bool
     def restart_homeassistant(self):
         """Restart Home-Assistant container.
@@ -91,6 +107,22 @@ def check_homeassistant_config(self):
         """
         return self.send_command("/homeassistant/check", timeout=300)
 
+    @_api_data
+    def retrieve_discovery_messages(self):
+        """Return all discovery data from Hass.io API.
+
+        This method return a coroutine.
+        """
+        return self.send_command("/discovery", method="get")
+
+    @_api_data
+    def get_discovery_message(self, uuid):
+        """Return a single discovery data message.
+
+        This method return a coroutine.
+        """
+        return self.send_command("/discovery/{}".format(uuid), method="get")
+
     @_api_bool
     async def update_hass_api(self, http_config, refresh_token):
         """Update Home Assistant API data on Hass.io."""
@@ -137,7 +169,7 @@ async def send_command(self, command, method="post", payload=None,
                 if request.status not in (200, 400):
                     _LOGGER.error(
                         "%s return code %d.", command, request.status)
-                    return None
+                    raise HassioAPIError()
 
                 answer = await request.json()
                 return answer
@@ -148,4 +180,4 @@ async def send_command(self, command, method="post", payload=None,
         except aiohttp.ClientError as err:
             _LOGGER.error("Client error on %s request %s", command, err)
 
-        return None
+        raise HassioAPIError()

tests/components/hassio/conftest.py

@@ -5,7 +5,7 @@
 import pytest
 
 from homeassistant.setup import async_setup_component
-from homeassistant.components.hassio.handler import HassIO
+from homeassistant.components.hassio.handler import HassIO, HassioAPIError
 
 from tests.common import mock_coro
 from . import API_PASSWORD, HASSIO_TOKEN
@@ -21,7 +21,7 @@ def hassio_env():
             patch.dict(os.environ, {'HASSIO_TOKEN': "123456"}), \
             patch('homeassistant.components.hassio.HassIO.'
                   'get_homeassistant_info',
-                  Mock(return_value=mock_coro(None))):
+                  Mock(side_effect=HassioAPIError())):
         yield
 
 
@@ -32,7 +32,7 @@ def hassio_client(hassio_env, hass, aiohttp_client):
                Mock(return_value=mock_coro({"result": "ok"}))), \
             patch('homeassistant.components.hassio.HassIO.'
                   'get_homeassistant_info',
-                  Mock(return_value=mock_coro(None))):
+                  Mock(side_effect=HassioAPIError())):
         hass.loop.run_until_complete(async_setup_component(hass, 'hassio', {
             'http': {
                 'api_password': API_PASSWORD