From b5426761f46ea02d3eee30ba5c1b2ddc5343c2d6 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Sun, 1 Sep 2019 17:57:25 +0200 Subject: [PATCH] UniFi - Simplify getting controller from config entry (#26335) * Simplify getting controller from config entry * Lint ignore no longer needed * Fix tests --- homeassistant/components/unifi/__init__.py | 15 +++++-------- homeassistant/components/unifi/config_flow.py | 22 ++++++++++++++++--- .../components/unifi/device_tracker.py | 9 ++------ homeassistant/components/unifi/switch.py | 11 ++-------- tests/components/unifi/test_device_tracker.py | 3 ++- tests/components/unifi/test_init.py | 8 +++++-- tests/components/unifi/test_switch.py | 3 ++- 7 files changed, 38 insertions(+), 33 deletions(-) diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py index 5ad60ddd835..db635828529 100644 --- a/homeassistant/components/unifi/__init__.py +++ b/homeassistant/components/unifi/__init__.py @@ -1,6 +1,9 @@ """Support for devices connected to UniFi POE.""" import voluptuous as vol +from homeassistant.components.unifi.config_flow import ( + get_controller_id_from_config_entry, +) from homeassistant.const import CONF_HOST from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC @@ -9,14 +12,12 @@ import homeassistant.helpers.config_validation as cv from .const import ( ATTR_MANUFACTURER, CONF_BLOCK_CLIENT, - CONF_CONTROLLER, CONF_DETECTION_TIME, CONF_DONT_TRACK_CLIENTS, CONF_DONT_TRACK_DEVICES, CONF_DONT_TRACK_WIRED_CLIENTS, CONF_SITE_ID, CONF_SSID_FILTER, - CONTROLLER_ID, DOMAIN, UNIFI_CONFIG, ) @@ -70,10 +71,7 @@ async def async_setup_entry(hass, config_entry): controller = UniFiController(hass, config_entry) - controller_id = CONTROLLER_ID.format( - host=config_entry.data[CONF_CONTROLLER][CONF_HOST], - site=config_entry.data[CONF_CONTROLLER][CONF_SITE_ID], - ) + controller_id = get_controller_id_from_config_entry(config_entry) hass.data[DOMAIN][controller_id] = controller @@ -98,9 +96,6 @@ async def async_setup_entry(hass, config_entry): async def async_unload_entry(hass, config_entry): """Unload a config entry.""" - controller_id = CONTROLLER_ID.format( - host=config_entry.data[CONF_CONTROLLER][CONF_HOST], - site=config_entry.data[CONF_CONTROLLER][CONF_SITE_ID], - ) + controller_id = get_controller_id_from_config_entry(config_entry) controller = hass.data[DOMAIN].pop(controller_id) return await controller.async_reset() diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py index c885f30af23..00b003746a2 100644 --- a/homeassistant/components/unifi/config_flow.py +++ b/homeassistant/components/unifi/config_flow.py @@ -11,13 +11,14 @@ from homeassistant.const import ( CONF_VERIFY_SSL, ) -from .const import ( # pylint: disable=unused-import +from .const import ( CONF_CONTROLLER, + CONF_DETECTION_TIME, + CONF_SITE_ID, CONF_TRACK_CLIENTS, CONF_TRACK_DEVICES, CONF_TRACK_WIRED_CLIENTS, - CONF_DETECTION_TIME, - CONF_SITE_ID, + CONTROLLER_ID, DEFAULT_TRACK_CLIENTS, DEFAULT_TRACK_DEVICES, DEFAULT_TRACK_WIRED_CLIENTS, @@ -33,6 +34,21 @@ DEFAULT_SITE_ID = "default" DEFAULT_VERIFY_SSL = False +@callback +def get_controller_id_from_config_entry(config_entry): + """Return controller with a matching bridge id.""" + return CONTROLLER_ID.format( + host=config_entry.data[CONF_CONTROLLER][CONF_HOST], + site=config_entry.data[CONF_CONTROLLER][CONF_SITE_ID], + ) + + +@callback +def get_controller_from_config_entry(hass, config_entry): + """Return controller with a matching bridge id.""" + return hass.data[DOMAIN][get_controller_id_from_config_entry(config_entry)] + + class UnifiFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): """Handle a UniFi config flow.""" diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index c4451546776..1f4aef754ad 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -5,7 +5,7 @@ import logging import voluptuous as vol from homeassistant import config_entries -from homeassistant.components import unifi +from homeassistant.components.unifi.config_flow import get_controller_from_config_entry from homeassistant.components.device_tracker import DOMAIN, PLATFORM_SCHEMA from homeassistant.components.device_tracker.config_entry import ScannerEntity from homeassistant.components.device_tracker.const import SOURCE_TYPE_ROUTER @@ -29,7 +29,6 @@ from .const import ( ATTR_MANUFACTURER, CONF_CONTROLLER, CONF_SITE_ID, - CONTROLLER_ID, DOMAIN as UNIFI_DOMAIN, ) @@ -106,11 +105,7 @@ async def async_setup_scanner(hass, config, sync_see, discovery_info): async def async_setup_entry(hass, config_entry, async_add_entities): """Set up device tracker for UniFi component.""" - controller_id = CONTROLLER_ID.format( - host=config_entry.data[CONF_CONTROLLER][CONF_HOST], - site=config_entry.data[CONF_CONTROLLER][CONF_SITE_ID], - ) - controller = hass.data[unifi.DOMAIN][controller_id] + controller = get_controller_from_config_entry(hass, config_entry) tracked = {} registry = await entity_registry.async_get_registry(hass) diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index ca4ae46f085..f46a55f671e 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -1,17 +1,14 @@ """Support for devices connected to UniFi POE.""" import logging -from homeassistant.components import unifi +from homeassistant.components.unifi.config_flow import get_controller_from_config_entry from homeassistant.components.switch import SwitchDevice -from homeassistant.const import CONF_HOST from homeassistant.core import callback from homeassistant.helpers import entity_registry from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC from homeassistant.helpers.dispatcher import async_dispatcher_connect from homeassistant.helpers.restore_state import RestoreEntity -from .const import CONF_CONTROLLER, CONF_SITE_ID, CONTROLLER_ID - LOGGER = logging.getLogger(__name__) @@ -25,11 +22,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): Switches are controlling network switch ports with Poe. """ - controller_id = CONTROLLER_ID.format( - host=config_entry.data[CONF_CONTROLLER][CONF_HOST], - site=config_entry.data[CONF_CONTROLLER][CONF_SITE_ID], - ) - controller = hass.data[unifi.DOMAIN][controller_id] + controller = get_controller_from_config_entry(hass, config_entry) if controller.site_role != "admin": return diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index e099286de7d..437fabf9689 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/tests/components/unifi/test_device_tracker.py @@ -17,6 +17,7 @@ from homeassistant.components.unifi.const import ( CONF_SSID_FILTER, CONF_TRACK_DEVICES, CONF_TRACK_WIRED_CLIENTS, + CONTROLLER_ID as CONF_CONTROLLER_ID, UNIFI_CONFIG, ) from homeassistant.const import ( @@ -101,7 +102,7 @@ CONTROLLER_DATA = { ENTRY_CONFIG = {CONF_CONTROLLER: CONTROLLER_DATA} -CONTROLLER_ID = unifi.CONTROLLER_ID.format(host="mock-host", site="mock-site") +CONTROLLER_ID = CONF_CONTROLLER_ID.format(host="mock-host", site="mock-site") @pytest.fixture diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py index b725e34f61d..ffd6d97e5b3 100644 --- a/tests/components/unifi/test_init.py +++ b/tests/components/unifi/test_init.py @@ -4,7 +4,11 @@ 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_CONTROLLER, CONF_SITE_ID +from homeassistant.components.unifi.const import ( + CONF_CONTROLLER, + CONF_SITE_ID, + CONTROLLER_ID as CONF_CONTROLLER_ID, +) from homeassistant.const import ( CONF_HOST, CONF_PASSWORD, @@ -113,7 +117,7 @@ async def test_controller_fail_setup(hass): mock_cntrlr.return_value.async_setup.return_value = mock_coro(False) assert await unifi.async_setup_entry(hass, entry) is False - controller_id = unifi.CONTROLLER_ID.format(host="0.0.0.0", site="default") + controller_id = CONF_CONTROLLER_ID.format(host="0.0.0.0", site="default") assert controller_id in hass.data[unifi.DOMAIN] diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index 3ac9ddb17dc..e660e57fc67 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -15,6 +15,7 @@ from homeassistant.components import unifi from homeassistant.components.unifi.const import ( CONF_CONTROLLER, CONF_SITE_ID, + CONTROLLER_ID as CONF_CONTROLLER_ID, UNIFI_CONFIG, ) from homeassistant.helpers import entity_registry @@ -213,7 +214,7 @@ CONTROLLER_DATA = { ENTRY_CONFIG = {CONF_CONTROLLER: CONTROLLER_DATA} -CONTROLLER_ID = unifi.CONTROLLER_ID.format(host="mock-host", site="mock-site") +CONTROLLER_ID = CONF_CONTROLLER_ID.format(host="mock-host", site="mock-site") @pytest.fixture