Skip to content
This repository has been archived by the owner on Dec 21, 2022. It is now read-only.

Commit

Permalink
UniFi - Simplify getting controller from config entry (home-assistant…
Browse files Browse the repository at this point in the history
…#26335)

* Simplify getting controller from config entry

* Lint ignore no longer needed

* Fix tests
  • Loading branch information
Kane610 authored Sep 1, 2019
1 parent 5b77a35 commit b542676
Show file tree
Hide file tree
Showing 7 changed files with 38 additions and 33 deletions.
15 changes: 5 additions & 10 deletions homeassistant/components/unifi/__init__.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -9,14 +12,12 @@
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,
)
Expand Down Expand Up @@ -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

Expand All @@ -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 changes: 19 additions & 3 deletions homeassistant/components/unifi/config_flow.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,14 @@
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,
Expand All @@ -33,6 +34,21 @@
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."""

Expand Down
9 changes: 2 additions & 7 deletions homeassistant/components/unifi/device_tracker.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
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
Expand All @@ -29,7 +29,6 @@
ATTR_MANUFACTURER,
CONF_CONTROLLER,
CONF_SITE_ID,
CONTROLLER_ID,
DOMAIN as UNIFI_DOMAIN,
)

Expand Down Expand Up @@ -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)
Expand Down
11 changes: 2 additions & 9 deletions homeassistant/components/unifi/switch.py
Original file line number Diff line number Diff line change
@@ -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__)


Expand All @@ -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
Expand Down
3 changes: 2 additions & 1 deletion tests/components/unifi/test_device_tracker.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
CONF_SSID_FILTER,
CONF_TRACK_DEVICES,
CONF_TRACK_WIRED_CLIENTS,
CONTROLLER_ID as CONF_CONTROLLER_ID,
UNIFI_CONFIG,
)
from homeassistant.const import (
Expand Down Expand Up @@ -101,7 +102,7 @@

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
Expand Down
8 changes: 6 additions & 2 deletions tests/components/unifi/test_init.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,11 @@
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,
Expand Down Expand Up @@ -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]


Expand Down
3 changes: 2 additions & 1 deletion tests/components/unifi/test_switch.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
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
Expand Down Expand Up @@ -213,7 +214,7 @@

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
Expand Down

0 comments on commit b542676

Please sign in to comment.