Add support for Hass.io discovery feature for Add-ons (#17035)

* Update handler.py * Update __init__.py * Update handler.py * Update __init__.py * Create discovery.py * Update handler.py * Update discovery.py * Update __init__.py * Update discovery.py * Update discovery.py * Update discovery.py * Update struct * Update handler.py * Update discovery.py * Update discovery.py * Update discovery.py * Update __init__.py * Update discovery.py * Update discovery.py * Update discovery.py * Update discovery.py * Update discovery.py * Update discovery.py * Update discovery.py * Update discovery.py * Update __init__.py * Update discovery.py * fix lint * Update discovery.py * cleanup old discovery * Update discovery.py * Update discovery.py * Fix lint * Fix tests * Write more tests with new functions * Update test_handler.py * Create test_discovery.py * Update conftest.py * Update test_discovery.py * Update conftest.py * Update test_discovery.py * Update conftest.py * Update test_discovery.py * Update test_discovery.py * Update test_discovery.py * Update test_discovery.py * Update test_discovery.py * Fix test * Add test * fix lint * Update handler.py * Update discovery.py * Update test_discovery.py * fix lint * Lint
2018-10-03 13:10:38 +02:00 · 2018-10-03 13:10:38 +02:00 · 704c9d8582
commit 704c9d8582
parent 2e5eb4d9dc
6 changed files with 326 additions and 59 deletions
--- a/homeassistant/components/hassio/init.py
+++ b/homeassistant/components/hassio/init.py
@ -19,7 +19,8 @@ import homeassistant.helpers.config_validation as cv
 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,10 +137,12 @@ 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
@ -147,18 +150,14 @@ async def async_check_config(hass):

 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,13 +228,13 @@ async def async_setup(hass, config):
            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(
@ -243,9 +242,11 @@ async def async_setup(hass, config):

    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_setup(hass, config):
        hass.services.async_register(
            HASS_DOMAIN, service, async_handle_core_service)

+    # Init discovery Hass.io feature
+    async_setup_discovery(hass, hassio, config)
+
    return True
--- a/homeassistant/components/hassio/discovery.py
+++ b/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)
--- a/homeassistant/components/hassio/handler.py
+++ b/homeassistant/components/hassio/handler.py
@ -21,12 +21,19 @@ _LOGGER = logging.getLogger(__name__)
 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 @@ class HassIO:
        """
        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 @@ class HassIO:
        """
        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 @@ class HassIO:
                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 @@ class HassIO:
        except aiohttp.ClientError as err:
            _LOGGER.error("Client error on %s request %s", command, err)

-        return None
+        raise HassioAPIError()
--- a/tests/components/hassio/conftest.py
+++ b/tests/components/hassio/conftest.py
@ -5,7 +5,7 @@ from unittest.mock import patch, Mock
 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
--- a/tests/components/hassio/test_discovery.py
+++ b/tests/components/hassio/test_discovery.py
@ -0,0 +1,88 @@
+"""Test config flow."""
+from unittest.mock import patch, Mock
+
+from homeassistant.const import EVENT_HOMEASSISTANT_START, HTTP_HEADER_HA_AUTH
+
+from tests.common import mock_coro
+from . import API_PASSWORD
+
+
+async def test_hassio_discovery_startup(hass, aioclient_mock, hassio_client):
+    """Test startup and discovery after event."""
+    aioclient_mock.get(
+        "http://127.0.0.1/discovery", json={
+            'result': 'ok', 'data': {'discovery': [
+                {
+                    "service": "mqtt", "uuid": "test",
+                    "addon": "mosquitto", "config":
+                    {
+                        'broker': 'mock-broker',
+                        'port': 1883,
+                        'username': 'mock-user',
+                        'password': 'mock-pass',
+                        'protocol': '3.1.1'
+                    }
+                }
+            ]}})
+    aioclient_mock.get(
+        "http://127.0.0.1/addons/mosquitto/info", json={
+            'result': 'ok', 'data': {'name': "Mosquitto Test"}
+        })
+
+    with patch('homeassistant.components.mqtt.'
+               'config_flow.FlowHandler.async_step_hassio',
+               Mock(return_value=mock_coro({"type": "abort"}))) as mock_mqtt:
+        hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
+        await hass.async_block_till_done()
+
+        assert aioclient_mock.call_count == 2
+        assert mock_mqtt.called
+        mock_mqtt.assert_called_with({
+            'broker': 'mock-broker', 'port': 1883, 'username': 'mock-user',
+            'password': 'mock-pass', 'protocol': '3.1.1',
+            'addon': 'Mosquitto Test',
+        })
+
+
+async def test_hassio_discovery_webhook(hass, aioclient_mock, hassio_client):
+    """Test discovery webhook."""
+    aioclient_mock.get(
+        "http://127.0.0.1/discovery/testuuid", json={
+            'result': 'ok', 'data':
+                {
+                    "service": "mqtt", "uuid": "test",
+                    "addon": "mosquitto", "config":
+                    {
+                        'broker': 'mock-broker',
+                        'port': 1883,
+                        'username': 'mock-user',
+                        'password': 'mock-pass',
+                        'protocol': '3.1.1'
+                    }
+                }
+            })
+    aioclient_mock.get(
+        "http://127.0.0.1/addons/mosquitto/info", json={
+            'result': 'ok', 'data': {'name': "Mosquitto Test"}
+        })
+
+    with patch('homeassistant.components.mqtt.'
+               'config_flow.FlowHandler.async_step_hassio',
+               Mock(return_value=mock_coro({"type": "abort"}))) as mock_mqtt:
+        resp = await hassio_client.post(
+            '/api/hassio_push/discovery/testuuid', headers={
+                HTTP_HEADER_HA_AUTH: API_PASSWORD
+            }, json={
+                "addon": "mosquitto", "service": "mqtt", "uuid": "testuuid"
+            }
+        )
+        await hass.async_block_till_done()
+
+        assert resp.status == 200
+        assert aioclient_mock.call_count == 2
+        assert mock_mqtt.called
+        mock_mqtt.assert_called_with({
+            'broker': 'mock-broker', 'port': 1883, 'username': 'mock-user',
+            'password': 'mock-pass', 'protocol': '3.1.1',
+            'addon': 'Mosquitto Test',
+        })
--- a/tests/components/hassio/test_handler.py
+++ b/tests/components/hassio/test_handler.py
@ -1,90 +1,118 @@
 """The tests for the hassio component."""
-import asyncio

 import aiohttp
+import pytest
+
+from homeassistant.components.hassio.handler import HassioAPIError


-@asyncio.coroutine
-def test_api_ping(hassio_handler, aioclient_mock):
+async def test_api_ping(hassio_handler, aioclient_mock):
    """Test setup with API ping."""
    aioclient_mock.get(
        "http://127.0.0.1/supervisor/ping", json={'result': 'ok'})

-    assert (yield from hassio_handler.is_connected())
+    assert (await hassio_handler.is_connected())
    assert aioclient_mock.call_count == 1


-@asyncio.coroutine
-def test_api_ping_error(hassio_handler, aioclient_mock):
+async def test_api_ping_error(hassio_handler, aioclient_mock):
    """Test setup with API ping error."""
    aioclient_mock.get(
        "http://127.0.0.1/supervisor/ping", json={'result': 'error'})

-    assert not (yield from hassio_handler.is_connected())
+    assert not (await hassio_handler.is_connected())
    assert aioclient_mock.call_count == 1


-@asyncio.coroutine
-def test_api_ping_exeption(hassio_handler, aioclient_mock):
+async def test_api_ping_exeption(hassio_handler, aioclient_mock):
    """Test setup with API ping exception."""
    aioclient_mock.get(
        "http://127.0.0.1/supervisor/ping", exc=aiohttp.ClientError())

-    assert not (yield from hassio_handler.is_connected())
+    assert not (await hassio_handler.is_connected())
    assert aioclient_mock.call_count == 1


-@asyncio.coroutine
-def test_api_homeassistant_info(hassio_handler, aioclient_mock):
+async def test_api_homeassistant_info(hassio_handler, aioclient_mock):
    """Test setup with API homeassistant info."""
    aioclient_mock.get(
        "http://127.0.0.1/homeassistant/info", json={
            'result': 'ok', 'data': {'last_version': '10.0'}})

-    data = yield from hassio_handler.get_homeassistant_info()
+    data = await hassio_handler.get_homeassistant_info()
    assert aioclient_mock.call_count == 1
    assert data['last_version'] == "10.0"


-@asyncio.coroutine
-def test_api_homeassistant_info_error(hassio_handler, aioclient_mock):
+async def test_api_homeassistant_info_error(hassio_handler, aioclient_mock):
    """Test setup with API homeassistant info error."""
    aioclient_mock.get(
        "http://127.0.0.1/homeassistant/info", json={
            'result': 'error', 'message': None})

-    data = yield from hassio_handler.get_homeassistant_info()
+    with pytest.raises(HassioAPIError):
+        await hassio_handler.get_homeassistant_info()
+
    assert aioclient_mock.call_count == 1
-    assert data is None


-@asyncio.coroutine
-def test_api_homeassistant_stop(hassio_handler, aioclient_mock):
+async def test_api_homeassistant_stop(hassio_handler, aioclient_mock):
    """Test setup with API HomeAssistant stop."""
    aioclient_mock.post(
        "http://127.0.0.1/homeassistant/stop", json={'result': 'ok'})

-    assert (yield from hassio_handler.stop_homeassistant())
+    assert (await hassio_handler.stop_homeassistant())
    assert aioclient_mock.call_count == 1


-@asyncio.coroutine
-def test_api_homeassistant_restart(hassio_handler, aioclient_mock):
+async def test_api_homeassistant_restart(hassio_handler, aioclient_mock):
    """Test setup with API HomeAssistant restart."""
    aioclient_mock.post(
        "http://127.0.0.1/homeassistant/restart", json={'result': 'ok'})

-    assert (yield from hassio_handler.restart_homeassistant())
+    assert (await hassio_handler.restart_homeassistant())
    assert aioclient_mock.call_count == 1


-@asyncio.coroutine
-def test_api_homeassistant_config(hassio_handler, aioclient_mock):
-    """Test setup with API HomeAssistant restart."""
+async def test_api_homeassistant_config(hassio_handler, aioclient_mock):
+    """Test setup with API HomeAssistant config."""
    aioclient_mock.post(
        "http://127.0.0.1/homeassistant/check", json={
            'result': 'ok', 'data': {'test': 'bla'}})

-    data = yield from hassio_handler.check_homeassistant_config()
+    data = await hassio_handler.check_homeassistant_config()
    assert data['data']['test'] == 'bla'
    assert aioclient_mock.call_count == 1
+
+
+async def test_api_addon_info(hassio_handler, aioclient_mock):
+    """Test setup with API Add-on info."""
+    aioclient_mock.get(
+        "http://127.0.0.1/addons/test/info", json={
+            'result': 'ok', 'data': {'name': 'bla'}})
+
+    data = await hassio_handler.get_addon_info("test")
+    assert data['name'] == 'bla'
+    assert aioclient_mock.call_count == 1
+
+
+async def test_api_discovery_message(hassio_handler, aioclient_mock):
+    """Test setup with API discovery message."""
+    aioclient_mock.get(
+        "http://127.0.0.1/discovery/test", json={
+            'result': 'ok', 'data': {"service": "mqtt"}})
+
+    data = await hassio_handler.get_discovery_message("test")
+    assert data['service'] == "mqtt"
+    assert aioclient_mock.call_count == 1
+
+
+async def test_api_retrieve_discovery(hassio_handler, aioclient_mock):
+    """Test setup with API discovery message."""
+    aioclient_mock.get(
+        "http://127.0.0.1/discovery", json={
+            'result': 'ok', 'data': {'discovery': [{"service": "mqtt"}]}})
+
+    data = await hassio_handler.retrieve_discovery_messages()
+    assert data['discovery'][-1]['service'] == "mqtt"
+    assert aioclient_mock.call_count == 1