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

UniFi - Simplify getting controller from config entry (#26335)

Browse source 
* Simplify getting controller from config entry

* Lint ignore no longer needed

* Fix tests
This commit is contained in:
Robert Svensson 2019-09-01 17:57:25 +02:00 • committed by GitHub
parent 5b77a357e6
commit b5426761f4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 38 additions and 33 deletions
Download patch file Download diff file Expand all files Collapse all files

15
homeassistant/components/unifi/__init__.py
View file

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

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

@ -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."""

9
homeassistant/components/unifi/device_tracker.py
View file

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

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

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

3
tests/components/unifi/test_device_tracker.py
View file

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

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

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

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

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