evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

UniFi simplify update (#24304)

Browse source
This commit is contained in:
Robert Svensson 2019-06-15 17:38:22 +02:00 • committed by GitHub
parent aa8ddeca34
commit e9b0f54a43
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 90 additions and 134 deletions
Download patch file Download diff file Expand all files Collapse all files

6
homeassistant/components/unifi/config_flow.py
View file

@ -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(

1
homeassistant/components/unifi/const.py
View file

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

72
homeassistant/components/unifi/controller.py
View file

@ -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:

86
homeassistant/components/unifi/switch.py
View file

@ -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):

30
tests/components/unifi/test_controller.py
View file

@ -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()):

6
tests/components/unifi/test_init.py
View file

@ -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
}
}

23
tests/components/unifi/test_switch.py
View file

@ -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()