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

@@ -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(
+    controller_id = get_controller_id_from_config_entry(config_entry)
-        host=config_entry.data[CONF_CONTROLLER][CONF_HOST],
-        site=config_entry.data[CONF_CONTROLLER][CONF_SITE_ID],
-    )
 
     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(
+    controller_id = get_controller_id_from_config_entry(config_entry)
-        host=config_entry.data[CONF_CONTROLLER][CONF_HOST],
-        site=config_entry.data[CONF_CONTROLLER][CONF_SITE_ID],
-    )
     controller = hass.data[DOMAIN].pop(controller_id)
     return await controller.async_reset()

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,
+    CONTROLLER_ID,
-    CONF_SITE_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."""
 

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(
+    controller = get_controller_from_config_entry(hass, config_entry)
-        host=config_entry.data[CONF_CONTROLLER][CONF_HOST],
-        site=config_entry.data[CONF_CONTROLLER][CONF_SITE_ID],
-    )
-    controller = hass.data[unifi.DOMAIN][controller_id]
     tracked = {}
 
     registry = await entity_registry.async_get_registry(hass)

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(
+    controller = get_controller_from_config_entry(hass, config_entry)
-        host=config_entry.data[CONF_CONTROLLER][CONF_HOST],
-        site=config_entry.data[CONF_CONTROLLER][CONF_SITE_ID],
-    )
-    controller = hass.data[unifi.DOMAIN][controller_id]
 
     if controller.site_role != "admin":
         return

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

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]
 
 

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