diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py index 95af8376773..1ffb8d94210 100644 --- a/homeassistant/components/unifi/config_flow.py +++ b/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, - DOMAIN, LOGGER) +from .const import CONF_CONTROLLER, CONF_SITE_ID, 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_POE_CONTROL: True + CONF_CONTROLLER: self.config } return self.async_create_entry( diff --git a/homeassistant/components/unifi/const.py b/homeassistant/components/unifi/const.py index 4d65a0d223a..7353a9d302b 100644 --- a/homeassistant/components/unifi/const.py +++ b/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' diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 5105e33f1d6..d0600315c01 100644 --- a/homeassistant/components/unifi/controller.py +++ b/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.config_entries.async_forward_entry_setup( - self.config_entry, 'switch')) + hass.async_create_task( + hass.config_entries.async_forward_entry_setup( + 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( - self.config_entry, 'switch') - return True + return await self.hass.config_entries.async_forward_entry_unload( + self.config_entry, 'switch') 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: diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 5f33a9c08d3..dd6fc1ff1a2 100644 --- a/homeassistant/components/unifi/switch.py +++ b/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(): - """Tell switches to reload state.""" - for client_id, client in switches.items(): - if client_id not in progress_waiting: - client.async_schedule_update_ha_state() + def update_controller(): + """Update the values of the controller.""" + update_items(controller, async_add_entities, switches) - try: - with async_timeout.timeout(4): - await controller.api.clients.update() - await controller.api.devices.update() + async_dispatcher_connect(hass, controller.event_update, update_controller) - except aiounifi.LoginRequired: - 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 + update_controller() - 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: - continue + for client_id in controller.api.clients: 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( - client, controller, request_controller_update) + switches[client_id] = UniFiSwitch(client, controller) 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): diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py index d1db25a23cd..b708a69bb67 100644 --- a/tests/components/unifi/test_controller.py +++ b/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 ( - CONF_POE_CONTROL, CONF_CONTROLLER, CONF_SITE_ID) +from homeassistant.components.unifi.const import 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_POE_CONTROL: True + CONF_CONTROLLER: CONTROLLER_DATA } @@ -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()): diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py index ec5ab5a577b..fffdcb5fb98 100644 --- a/tests/components/unifi/test_init.py +++ b/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 ( - CONF_POE_CONTROL, CONF_CONTROLLER, CONF_SITE_ID) +from homeassistant.components.unifi.const import 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 + } } diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index 5a04b415f5d..4eba3aca61e 100644 --- a/tests/components/unifi/test_switch.py +++ b/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 ( - CONF_POE_CONTROL, CONF_CONTROLLER, CONF_SITE_ID) +from homeassistant.components.unifi.const import 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_POE_CONTROL: True + CONF_CONTROLLER: CONTROLLER_DATA } 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( - available=True, - api=Mock(), - spec=unifi.UniFiController - ) - controller.mac = '10:00:00:00:00:01' + controller = unifi.UniFiController(hass, None) + + controller.api = Mock() 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()