UniFi simplify update (#24304)

homeassistant/components/unifi/config_flow.py

@@ -5,8 +5,7 @@ from homeassistant import config_entries
 from homeassistant.const import (
     CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME, CONF_VERIFY_SSL)
 
-from .const import (CONF_CONTROLLER, CONF_POE_CONTROL, CONF_SITE_ID,
+from .const import CONF_CONTROLLER, CONF_SITE_ID, DOMAIN, LOGGER
-                    DOMAIN, LOGGER)
 from .controller import get_controller
 from .errors import (
     AlreadyConfigured, AuthenticationRequired, CannotConnect, UserLevel)
@@ -99,8 +98,7 @@ class UnifiFlowHandler(config_entries.ConfigFlow):
                         raise AlreadyConfigured
 
                 data = {
-                    CONF_CONTROLLER: self.config,
+                    CONF_CONTROLLER: self.config
-                    CONF_POE_CONTROL: True
                 }
 
                 return self.async_create_entry(

homeassistant/components/unifi/const.py

@@ -7,5 +7,4 @@ DOMAIN = 'unifi'
 CONTROLLER_ID = '{host}-{site}'
 
 CONF_CONTROLLER = 'controller'
-CONF_POE_CONTROL = 'poe_control'
 CONF_SITE_ID = 'site'

homeassistant/components/unifi/controller.py

@@ -5,11 +5,14 @@ import async_timeout
 
 from aiohttp import CookieJar
 
+import aiounifi
+
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.const import CONF_HOST
 from homeassistant.helpers import aiohttp_client
+from homeassistant.helpers.dispatcher import async_dispatcher_send
 
-from .const import CONF_CONTROLLER, CONF_POE_CONTROL, LOGGER
+from .const import CONF_CONTROLLER, CONF_SITE_ID, CONTROLLER_ID, LOGGER
 from .errors import AuthenticationRequired, CannotConnect
 
 
@@ -37,7 +40,57 @@ class UniFiController:
                 return client.mac
         return None
 
-    async def async_setup(self, tries=0):
+    @property
+    def event_update(self):
+        """Event specific per UniFi entry to signal new data."""
+        return 'unifi-update-{}'.format(
+            CONTROLLER_ID.format(
+                host=self.host,
+                site=self.config_entry.data[CONF_CONTROLLER][CONF_SITE_ID]))
+
+    async def request_update(self):
+        """Request an update."""
+        if self.progress is not None:
+            return await self.progress
+
+        self.progress = self.hass.async_create_task(self.async_update())
+        await self.progress
+
+        self.progress = None
+
+    async def async_update(self):
+        """Update UniFi controller information."""
+        failed = False
+
+        try:
+            with async_timeout.timeout(4):
+                await self.api.clients.update()
+                await self.api.devices.update()
+
+        except aiounifi.LoginRequired:
+            try:
+                with async_timeout.timeout(5):
+                    await self.api.login()
+
+            except (asyncio.TimeoutError, aiounifi.AiounifiException):
+                failed = True
+                if self.available:
+                    LOGGER.error('Unable to reach controller %s', self.host)
+                    self.available = False
+
+        except (asyncio.TimeoutError, aiounifi.AiounifiException):
+            failed = True
+            if self.available:
+                LOGGER.error('Unable to reach controller %s', self.host)
+                self.available = False
+
+        if not failed and not self.available:
+            LOGGER.info('Reconnected to controller %s', self.host)
+            self.available = True
+
+        async_dispatcher_send(self.hass, self.event_update)
+
+    async def async_setup(self):
         """Set up a UniFi controller."""
         hass = self.hass
 
@@ -54,10 +107,9 @@ class UniFiController:
                 'Unknown error connecting with UniFi controller.')
             return False
 
-        if self.config_entry.data[CONF_POE_CONTROL]:
+        hass.async_create_task(
-            hass.async_create_task(
+            hass.config_entries.async_forward_entry_setup(
-                hass.config_entries.async_forward_entry_setup(
+                self.config_entry, 'switch'))
-                    self.config_entry, 'switch'))
 
         return True
 
@@ -71,17 +123,13 @@ class UniFiController:
         if self.api is None:
             return True
 
-        if self.config_entry.data[CONF_POE_CONTROL]:
+        return await self.hass.config_entries.async_forward_entry_unload(
-            return await self.hass.config_entries.async_forward_entry_unload(
+            self.config_entry, 'switch')
-                self.config_entry, 'switch')
-        return True
 
 
 async def get_controller(
         hass, host, username, password, port, site, verify_ssl):
     """Create a controller object and verify authentication."""
-    import aiounifi
-
     sslcontext = None
 
     if verify_ssl:

homeassistant/components/unifi/switch.py

@@ -1,15 +1,13 @@
 """Support for devices connected to UniFi POE."""
-import asyncio
 from datetime import timedelta
 import logging
 
-import async_timeout
-
 from homeassistant.components import unifi
 from homeassistant.components.switch import SwitchDevice
 from homeassistant.const import CONF_HOST
 from homeassistant.core import callback
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
 
 from .const import CONF_CONTROLLER, CONF_SITE_ID, CONTROLLER_ID
 
@@ -36,79 +34,23 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     controller = hass.data[unifi.DOMAIN][controller_id]
     switches = {}
 
-    progress = None
-    update_progress = set()
-
-    async def request_update(object_id):
-        """Request an update."""
-        nonlocal progress
-        update_progress.add(object_id)
-
-        if progress is not None:
-            return await progress
-
-        progress = asyncio.ensure_future(update_controller())
-        result = await progress
-        progress = None
-        update_progress.clear()
-        return result
-
-    async def update_controller():
-        """Update the values of the controller."""
-        tasks = [async_update_items(
-            controller, async_add_entities, request_update,
-            switches, update_progress
-        )]
-        await asyncio.wait(tasks)
-
-    await update_controller()
-
-
-async def async_update_items(controller, async_add_entities,
-                             request_controller_update, switches,
-                             progress_waiting):
-    """Update POE port state from the controller."""
-    import aiounifi
-
     @callback
-    def update_switch_state():
+    def update_controller():
-        """Tell switches to reload state."""
+        """Update the values of the controller."""
-        for client_id, client in switches.items():
+        update_items(controller, async_add_entities, switches)
-            if client_id not in progress_waiting:
-                client.async_schedule_update_ha_state()
 
-    try:
+    async_dispatcher_connect(hass, controller.event_update, update_controller)
-        with async_timeout.timeout(4):
-            await controller.api.clients.update()
-            await controller.api.devices.update()
 
-    except aiounifi.LoginRequired:
+    update_controller()
-        try:
-            with async_timeout.timeout(5):
-                await controller.api.login()
-        except (asyncio.TimeoutError, aiounifi.AiounifiException):
-            if controller.available:
-                controller.available = False
-                update_switch_state()
-            return
 
-    except (asyncio.TimeoutError, aiounifi.AiounifiException):
-        if controller.available:
-            LOGGER.error('Unable to reach controller %s', controller.host)
-            controller.available = False
-            update_switch_state()
-        return
-
-    if not controller.available:
-        LOGGER.info('Reconnected to controller %s', controller.host)
-        controller.available = True
 
+@callback
+def update_items(controller, async_add_entities, switches):
+    """Update POE port state from the controller."""
     new_switches = []
     devices = controller.api.devices
-    for client_id in controller.api.clients:
 
-        if client_id in progress_waiting:
+    for client_id in controller.api.clients:
-            continue
 
         if client_id in switches:
             LOGGER.debug("Updating UniFi switch %s (%s)",
@@ -137,8 +79,7 @@ async def async_update_items(controller, async_add_entities,
         if multi_clients_on_port:
             continue
 
-        switches[client_id] = UniFiSwitch(
+        switches[client_id] = UniFiSwitch(client, controller)
-            client, controller, request_controller_update)
         new_switches.append(switches[client_id])
         LOGGER.debug("New UniFi switch %s (%s)", client.hostname, client.mac)
 
@@ -149,18 +90,17 @@ async def async_update_items(controller, async_add_entities,
 class UniFiSwitch(SwitchDevice):
     """Representation of a client that uses POE."""
 
-    def __init__(self, client, controller, request_controller_update):
+    def __init__(self, client, controller):
         """Set up switch."""
         self.client = client
         self.controller = controller
         self.poe_mode = None
         if self.port.poe_mode != 'off':
             self.poe_mode = self.port.poe_mode
-        self.async_request_controller_update = request_controller_update
 
     async def async_update(self):
         """Synchronize state with controller."""
-        await self.async_request_controller_update(self.client.mac)
+        await self.controller.request_update()
 
     @property
     def name(self):

tests/components/unifi/test_controller.py

@@ -4,8 +4,7 @@ from unittest.mock import Mock, patch
 import pytest
 
 from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.components.unifi.const import (
+from homeassistant.components.unifi.const import CONF_CONTROLLER, CONF_SITE_ID
-    CONF_POE_CONTROL, CONF_CONTROLLER, CONF_SITE_ID)
 from homeassistant.const import (
     CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME, CONF_VERIFY_SSL)
 from homeassistant.components.unifi import controller, errors
@@ -22,8 +21,7 @@ CONTROLLER_DATA = {
 }
 
 ENTRY_CONFIG = {
-    CONF_CONTROLLER: CONTROLLER_DATA,
+    CONF_CONTROLLER: CONTROLLER_DATA
-    CONF_POE_CONTROL: True
 }
 
 
@@ -171,30 +169,6 @@ async def test_reset_unloads_entry_if_setup():
     assert len(hass.config_entries.async_forward_entry_unload.mock_calls) == 1
 
 
-async def test_reset_unloads_entry_without_poe_control():
-    """Calling reset while the entry has been setup."""
-    hass = Mock()
-    entry = Mock()
-    entry.data = dict(ENTRY_CONFIG)
-    entry.data[CONF_POE_CONTROL] = False
-    api = Mock()
-    api.initialize.return_value = mock_coro(True)
-
-    unifi_controller = controller.UniFiController(hass, entry)
-
-    with patch.object(controller, 'get_controller',
-                      return_value=mock_coro(api)):
-        assert await unifi_controller.async_setup() is True
-
-    assert not hass.config_entries.async_forward_entry_setup.mock_calls
-
-    hass.config_entries.async_forward_entry_unload.return_value = \
-        mock_coro(True)
-    assert await unifi_controller.async_reset()
-
-    assert not hass.config_entries.async_forward_entry_unload.mock_calls
-
-
 async def test_get_controller(hass):
     """Successful call."""
     with patch('aiounifi.Controller.login', return_value=mock_coro()):

tests/components/unifi/test_init.py

@@ -4,8 +4,7 @@ from unittest.mock import Mock, patch
 from homeassistant.components import unifi
 from homeassistant.components.unifi import config_flow
 from homeassistant.setup import async_setup_component
-from homeassistant.components.unifi.const import (
+from homeassistant.components.unifi.const import CONF_CONTROLLER, CONF_SITE_ID
-    CONF_POE_CONTROL, CONF_CONTROLLER, CONF_SITE_ID)
 from homeassistant.const import (
     CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME, CONF_VERIFY_SSL)
 
@@ -186,8 +185,7 @@ async def test_flow_works(hass, aioclient_mock):
             CONF_PORT: 1234,
             CONF_SITE_ID: 'default',
             CONF_VERIFY_SSL: True
-        },
+        }
-        CONF_POE_CONTROL: True
     }
 
 

tests/components/unifi/test_switch.py

@@ -4,22 +4,21 @@ from unittest.mock import Mock
 
 import pytest
 
+from tests.common import mock_coro
+
 import aiounifi
 from aiounifi.clients import Clients
 from aiounifi.devices import Devices
 
 from homeassistant import config_entries
 from homeassistant.components import unifi
-from homeassistant.components.unifi.const import (
+from homeassistant.components.unifi.const import CONF_CONTROLLER, CONF_SITE_ID
-    CONF_POE_CONTROL, CONF_CONTROLLER, CONF_SITE_ID)
 from homeassistant.setup import async_setup_component
 from homeassistant.const import (
     CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_USERNAME, CONF_VERIFY_SSL)
 
 import homeassistant.components.switch as switch
 
-from tests.common import mock_coro
-
 CLIENT_1 = {
     'hostname': 'client_1',
     'ip': '10.0.0.1',
@@ -180,8 +179,7 @@ CONTROLLER_DATA = {
 }
 
 ENTRY_CONFIG = {
-    CONF_CONTROLLER: CONTROLLER_DATA,
+    CONF_CONTROLLER: CONTROLLER_DATA
-    CONF_POE_CONTROL: True
 }
 
 CONTROLLER_ID = unifi.CONTROLLER_ID.format(host='mock-host', site='mock-site')
@@ -190,12 +188,9 @@ CONTROLLER_ID = unifi.CONTROLLER_ID.format(host='mock-host', site='mock-site')
 @pytest.fixture
 def mock_controller(hass):
     """Mock a UniFi Controller."""
-    controller = Mock(
+    controller = unifi.UniFiController(hass, None)
-        available=True,
+
-        api=Mock(),
+    controller.api = Mock()
-        spec=unifi.UniFiController
-    )
-    controller.mac = '10:00:00:00:00:01'
     controller.mock_requests = []
 
     controller.mock_client_responses = deque()
@@ -224,6 +219,9 @@ async def setup_controller(hass, mock_controller):
     config_entry = config_entries.ConfigEntry(
         1, unifi.DOMAIN, 'Mock Title', ENTRY_CONFIG, 'test',
         config_entries.CONN_CLASS_LOCAL_POLL)
+    mock_controller.config_entry = config_entry
+
+    await mock_controller.async_update()
     await hass.config_entries.async_forward_entry_setup(config_entry, 'switch')
     # To flush out the service call to update the group
     await hass.async_block_till_done()
@@ -242,6 +240,7 @@ async def test_platform_manually_configured(hass):
 async def test_no_clients(hass, mock_controller):
     """Test the update_clients function when no clients are found."""
     mock_controller.mock_client_responses.append({})
+    mock_controller.mock_device_responses.append({})
     await setup_controller(hass, mock_controller)
     assert len(mock_controller.mock_requests) == 2
     assert not hass.states.async_all()