UniFi - Store controller in config_entry.entry_id (#34553)

* Store controller in config_entry.entry_id

* Clean up imports
This commit is contained in:
Robert Svensson 2020-04-23 16:48:24 +02:00 committed by GitHub
parent a0fbf9ba47
commit c3689d7416
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
13 changed files with 190 additions and 182 deletions

View file

@ -7,7 +7,12 @@ import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
from .config_flow import get_controller_id_from_config_entry
from .const import ATTR_MANUFACTURER, DOMAIN, LOGGER, UNIFI_WIRELESS_CLIENTS
from .const import (
ATTR_MANUFACTURER,
DOMAIN as UNIFI_DOMAIN,
LOGGER,
UNIFI_WIRELESS_CLIENTS,
)
from .controller import UniFiController
SAVE_DELAY = 10
@ -15,7 +20,8 @@ STORAGE_KEY = "unifi_data"
STORAGE_VERSION = 1
CONFIG_SCHEMA = vol.Schema(
cv.deprecated(DOMAIN, invalidation_version="0.109"), {DOMAIN: cv.match_all}
cv.deprecated(UNIFI_DOMAIN, invalidation_version="0.109"),
{UNIFI_DOMAIN: cv.match_all},
)
@ -29,16 +35,13 @@ async def async_setup(hass, config):
async def async_setup_entry(hass, config_entry):
"""Set up the UniFi component."""
if DOMAIN not in hass.data:
hass.data[DOMAIN] = {}
hass.data.setdefault(UNIFI_DOMAIN, {})
controller = UniFiController(hass, config_entry)
if not await controller.async_setup():
return False
controller_id = get_controller_id_from_config_entry(config_entry)
hass.data[DOMAIN][controller_id] = controller
hass.data[UNIFI_DOMAIN][config_entry.entry_id] = controller
hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, controller.shutdown)
@ -62,8 +65,7 @@ async def async_setup_entry(hass, config_entry):
async def async_unload_entry(hass, config_entry):
"""Unload a config entry."""
controller_id = get_controller_id_from_config_entry(config_entry)
controller = hass.data[DOMAIN].pop(controller_id)
controller = hass.data[UNIFI_DOMAIN].pop(config_entry.entry_id)
return await controller.async_reset()
@ -90,15 +92,21 @@ class UnifiWirelessClients:
def get_data(self, config_entry):
"""Get data related to a specific controller."""
controller_id = get_controller_id_from_config_entry(config_entry)
data = self.data.get(controller_id, {"wireless_devices": []})
key = config_entry.entry_id
if controller_id in self.data:
key = controller_id
data = self.data.get(key, {"wireless_devices": []})
return set(data["wireless_devices"])
@callback
def update_data(self, data, config_entry):
"""Update data and schedule to save to file."""
controller_id = get_controller_id_from_config_entry(config_entry)
self.data[controller_id] = {"wireless_devices": list(data)}
if controller_id in self.data:
self.data.pop(controller_id)
self.data[config_entry.entry_id] = {"wireless_devices": list(data)}
self._store.async_delay_save(self._data_to_save, SAVE_DELAY)
@callback

View file

@ -28,7 +28,7 @@ from .const import (
CONF_TRACK_WIRED_CLIENTS,
CONTROLLER_ID,
DEFAULT_POE_CLIENTS,
DOMAIN,
DOMAIN as UNIFI_DOMAIN,
LOGGER,
)
from .controller import get_controller
@ -48,13 +48,7 @@ def get_controller_id_from_config_entry(config_entry):
)
@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):
class UnifiFlowHandler(config_entries.ConfigFlow, domain=UNIFI_DOMAIN):
"""Handle a UniFi config flow."""
VERSION = 1
@ -179,7 +173,7 @@ class UnifiOptionsFlowHandler(config_entries.OptionsFlow):
async def async_step_init(self, user_input=None):
"""Manage the UniFi options."""
self.controller = get_controller_from_config_entry(self.hass, self.config_entry)
self.controller = self.hass.data[UNIFI_DOMAIN][self.config_entry.entry_id]
self.options[CONF_BLOCK_CLIENT] = self.controller.option_block_clients
return await self.async_step_device_tracker()

View file

@ -17,7 +17,7 @@ from aiounifi.events import WIRELESS_CLIENT_CONNECTED, WIRELESS_GUEST_CONNECTED
from aiounifi.websocket import STATE_DISCONNECTED, STATE_RUNNING
import async_timeout
from homeassistant.components.device_tracker import DOMAIN as DT_DOMAIN
from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
from homeassistant.const import CONF_HOST
@ -46,14 +46,14 @@ from .const import (
DEFAULT_TRACK_CLIENTS,
DEFAULT_TRACK_DEVICES,
DEFAULT_TRACK_WIRED_CLIENTS,
DOMAIN,
DOMAIN as UNIFI_DOMAIN,
LOGGER,
UNIFI_WIRELESS_CLIENTS,
)
from .errors import AuthenticationRequired, CannotConnect
RETRY_TIMER = 15
SUPPORTED_PLATFORMS = [DT_DOMAIN, SENSOR_DOMAIN, SWITCH_DOMAIN]
SUPPORTED_PLATFORMS = [TRACKER_DOMAIN, SENSOR_DOMAIN, SWITCH_DOMAIN]
class UniFiController:
@ -283,14 +283,9 @@ class UniFiController:
return True
@staticmethod
async def async_config_entry_updated(hass, entry) -> None:
async def async_config_entry_updated(hass, config_entry) -> None:
"""Handle signals of config entry being updated."""
controller_id = CONTROLLER_ID.format(
host=entry.data[CONF_CONTROLLER][CONF_HOST],
site=entry.data[CONF_CONTROLLER][CONF_SITE_ID],
)
controller = hass.data[DOMAIN][controller_id]
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
async_dispatcher_send(hass, controller.signal_options_update)
@callback

View file

@ -4,16 +4,15 @@ import logging
from homeassistant.components.device_tracker import DOMAIN
from homeassistant.components.device_tracker.config_entry import ScannerEntity
from homeassistant.components.device_tracker.const import SOURCE_TYPE_ROUTER
from homeassistant.components.unifi.config_flow import get_controller_from_config_entry
from homeassistant.components.unifi.unifi_entity_base import UniFiBase
from homeassistant.core import callback
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.event import async_track_point_in_utc_time
import homeassistant.util.dt as dt_util
from .const import ATTR_MANUFACTURER
from .const import ATTR_MANUFACTURER, DOMAIN as UNIFI_DOMAIN
from .unifi_client import UniFiClient
from .unifi_entity_base import UniFiBase
LOGGER = logging.getLogger(__name__)
@ -45,7 +44,7 @@ DEVICE_TRACKER = "device"
async def async_setup_entry(hass, config_entry, async_add_entities):
"""Set up device tracker for UniFi component."""
controller = get_controller_from_config_entry(hass, config_entry)
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
controller.entities[DOMAIN] = {CLIENT_TRACKER: set(), DEVICE_TRACKER: set()}
# Restore clients that is not a part of active clients list.

View file

@ -2,11 +2,11 @@
import logging
from homeassistant.components.sensor import DOMAIN
from homeassistant.components.unifi.config_flow import get_controller_from_config_entry
from homeassistant.const import DATA_MEGABYTES
from homeassistant.core import callback
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from .const import DOMAIN as UNIFI_DOMAIN
from .unifi_client import UniFiClient
LOGGER = logging.getLogger(__name__)
@ -21,7 +21,7 @@ async def async_setup_platform(hass, config, async_add_entities, discovery_info=
async def async_setup_entry(hass, config_entry, async_add_entities):
"""Set up sensors for UniFi integration."""
controller = get_controller_from_config_entry(hass, config_entry)
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
controller.entities[DOMAIN] = {RX_SENSOR: set(), TX_SENSOR: set()}
@callback

View file

@ -2,11 +2,11 @@
import logging
from homeassistant.components.switch import DOMAIN, SwitchDevice
from homeassistant.components.unifi.config_flow import get_controller_from_config_entry
from homeassistant.core import callback
from homeassistant.helpers.dispatcher import async_dispatcher_connect
from homeassistant.helpers.restore_state import RestoreEntity
from .const import DOMAIN as UNIFI_DOMAIN
from .unifi_client import UniFiClient
LOGGER = logging.getLogger(__name__)
@ -24,7 +24,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
Switches are controlling network access and switch ports with POE.
"""
controller = get_controller_from_config_entry(hass, config_entry)
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
controller.entities[DOMAIN] = {BLOCK_SWITCH: set(), POE_SWITCH: set()}
if controller.site_role != "admin":

View file

@ -15,10 +15,11 @@ from aiounifi.events import (
WIRELESS_CLIENT_UNBLOCKED,
)
from homeassistant.components.unifi.unifi_entity_base import UniFiBase
from homeassistant.core import callback
from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC
from .unifi_entity_base import UniFiBase
LOGGER = logging.getLogger(__name__)
CLIENT_BLOCKED = (WIRED_CLIENT_BLOCKED, WIRELESS_CLIENT_BLOCKED)

View file

@ -3,8 +3,6 @@ import aiounifi
from asynctest import patch
from homeassistant import data_entry_flow
from homeassistant.components import unifi
from homeassistant.components.unifi import config_flow
from homeassistant.components.unifi.const import (
CONF_ALLOW_BANDWIDTH_SENSORS,
CONF_BLOCK_CLIENT,
@ -17,6 +15,7 @@ from homeassistant.components.unifi.const import (
CONF_TRACK_CLIENTS,
CONF_TRACK_DEVICES,
CONF_TRACK_WIRED_CLIENTS,
DOMAIN as UNIFI_DOMAIN,
)
from homeassistant.const import (
CONF_HOST,
@ -39,7 +38,7 @@ async def test_flow_works(hass, aioclient_mock, mock_discovery):
"""Test config flow."""
mock_discovery.return_value = "1"
result = await hass.config_entries.flow.async_init(
config_flow.DOMAIN, context={"source": "user"}
UNIFI_DOMAIN, context={"source": "user"}
)
assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
@ -97,7 +96,7 @@ async def test_flow_works(hass, aioclient_mock, mock_discovery):
async def test_flow_works_multiple_sites(hass, aioclient_mock):
"""Test config flow works when finding multiple sites."""
result = await hass.config_entries.flow.async_init(
config_flow.DOMAIN, context={"source": "user"}
UNIFI_DOMAIN, context={"source": "user"}
)
assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
@ -143,12 +142,12 @@ async def test_flow_works_multiple_sites(hass, aioclient_mock):
async def test_flow_fails_site_already_configured(hass, aioclient_mock):
"""Test config flow."""
entry = MockConfigEntry(
domain=unifi.DOMAIN, data={"controller": {"host": "1.2.3.4", "site": "site_id"}}
domain=UNIFI_DOMAIN, data={"controller": {"host": "1.2.3.4", "site": "site_id"}}
)
entry.add_to_hass(hass)
result = await hass.config_entries.flow.async_init(
config_flow.DOMAIN, context={"source": "user"}
UNIFI_DOMAIN, context={"source": "user"}
)
assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
@ -188,7 +187,7 @@ async def test_flow_fails_site_already_configured(hass, aioclient_mock):
async def test_flow_fails_user_credentials_faulty(hass, aioclient_mock):
"""Test config flow."""
result = await hass.config_entries.flow.async_init(
config_flow.DOMAIN, context={"source": "user"}
UNIFI_DOMAIN, context={"source": "user"}
)
assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
@ -215,7 +214,7 @@ async def test_flow_fails_user_credentials_faulty(hass, aioclient_mock):
async def test_flow_fails_controller_unavailable(hass, aioclient_mock):
"""Test config flow."""
result = await hass.config_entries.flow.async_init(
config_flow.DOMAIN, context={"source": "user"}
UNIFI_DOMAIN, context={"source": "user"}
)
assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
@ -242,7 +241,7 @@ async def test_flow_fails_controller_unavailable(hass, aioclient_mock):
async def test_flow_fails_unknown_problem(hass, aioclient_mock):
"""Test config flow."""
result = await hass.config_entries.flow.async_init(
config_flow.DOMAIN, context={"source": "user"}
UNIFI_DOMAIN, context={"source": "user"}
)
assert result["type"] == data_entry_flow.RESULT_TYPE_FORM

View file

@ -7,12 +7,25 @@ import aiounifi
from asynctest import patch
import pytest
from homeassistant.components import unifi
from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
from homeassistant.components.unifi.const import (
CONF_CONTROLLER,
CONF_SITE_ID,
DEFAULT_ALLOW_BANDWIDTH_SENSORS,
DEFAULT_DETECTION_TIME,
DEFAULT_TRACK_CLIENTS,
DEFAULT_TRACK_DEVICES,
DEFAULT_TRACK_WIRED_CLIENTS,
DOMAIN as UNIFI_DOMAIN,
UNIFI_WIRELESS_CLIENTS,
)
from homeassistant.components.unifi.controller import (
SUPPORTED_PLATFORMS,
get_controller,
)
from homeassistant.components.unifi.errors import AuthenticationRequired, CannotConnect
from homeassistant.const import (
CONF_HOST,
CONF_PASSWORD,
@ -68,10 +81,10 @@ async def setup_unifi_integration(
controllers=None,
):
"""Create the UniFi controller."""
assert await async_setup_component(hass, unifi.DOMAIN, {})
assert await async_setup_component(hass, UNIFI_DOMAIN, {})
config_entry = MockConfigEntry(
domain=unifi.DOMAIN,
domain=UNIFI_DOMAIN,
data=deepcopy(config),
options=deepcopy(options),
entry_id=1,
@ -124,10 +137,9 @@ async def setup_unifi_integration(
await hass.config_entries.async_setup(config_entry.entry_id)
await hass.async_block_till_done()
controller_id = unifi.get_controller_id_from_config_entry(config_entry)
if controller_id not in hass.data[unifi.DOMAIN]:
if config_entry.entry_id not in hass.data[UNIFI_DOMAIN]:
return None
controller = hass.data[unifi.DOMAIN][controller_id]
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
controller.mock_client_responses = mock_client_responses
controller.mock_device_responses = mock_device_responses
@ -147,31 +159,22 @@ async def test_controller_setup(hass):
controller = await setup_unifi_integration(hass)
entry = controller.config_entry
assert len(forward_entry_setup.mock_calls) == len(
unifi.controller.SUPPORTED_PLATFORMS
)
assert forward_entry_setup.mock_calls[0][1] == (entry, "device_tracker")
assert forward_entry_setup.mock_calls[1][1] == (entry, "sensor")
assert forward_entry_setup.mock_calls[2][1] == (entry, "switch")
assert len(forward_entry_setup.mock_calls) == len(SUPPORTED_PLATFORMS)
assert forward_entry_setup.mock_calls[0][1] == (entry, TRACKER_DOMAIN)
assert forward_entry_setup.mock_calls[1][1] == (entry, SENSOR_DOMAIN)
assert forward_entry_setup.mock_calls[2][1] == (entry, SWITCH_DOMAIN)
assert controller.host == CONTROLLER_DATA[CONF_HOST]
assert controller.site == CONTROLLER_DATA[CONF_SITE_ID]
assert controller.site_name in SITES
assert controller.site_role == SITES[controller.site_name]["role"]
assert (
controller.option_allow_bandwidth_sensors
== unifi.const.DEFAULT_ALLOW_BANDWIDTH_SENSORS
)
assert controller.option_allow_bandwidth_sensors == DEFAULT_ALLOW_BANDWIDTH_SENSORS
assert isinstance(controller.option_block_clients, list)
assert controller.option_track_clients == unifi.const.DEFAULT_TRACK_CLIENTS
assert controller.option_track_devices == unifi.const.DEFAULT_TRACK_DEVICES
assert (
controller.option_track_wired_clients == unifi.const.DEFAULT_TRACK_WIRED_CLIENTS
)
assert controller.option_detection_time == timedelta(
seconds=unifi.const.DEFAULT_DETECTION_TIME
)
assert controller.option_track_clients == DEFAULT_TRACK_CLIENTS
assert controller.option_track_devices == DEFAULT_TRACK_DEVICES
assert controller.option_track_wired_clients == DEFAULT_TRACK_WIRED_CLIENTS
assert controller.option_detection_time == timedelta(seconds=DEFAULT_DETECTION_TIME)
assert isinstance(controller.option_ssid_filter, list)
assert controller.mac is None
@ -184,23 +187,27 @@ async def test_controller_setup(hass):
async def test_controller_mac(hass):
"""Test that it is possible to identify controller mac."""
controller = await setup_unifi_integration(hass, clients_response=[CONTROLLER_HOST])
assert controller.mac == "10:00:00:00:00:01"
assert controller.mac == CONTROLLER_HOST["mac"]
async def test_controller_not_accessible(hass):
"""Retry to login gets scheduled when connection fails."""
with patch.object(
unifi.controller, "get_controller", side_effect=unifi.errors.CannotConnect
with patch(
"homeassistant.components.unifi.controller.get_controller",
side_effect=CannotConnect,
):
await setup_unifi_integration(hass)
assert hass.data[unifi.DOMAIN] == {}
assert hass.data[UNIFI_DOMAIN] == {}
async def test_controller_unknown_error(hass):
"""Unknown errors are handled."""
with patch.object(unifi.controller, "get_controller", side_effect=Exception):
with patch(
"homeassistant.components.unifi.controller.get_controller",
side_effect=Exception,
):
await setup_unifi_integration(hass)
assert hass.data[unifi.DOMAIN] == {}
assert hass.data[UNIFI_DOMAIN] == {}
async def test_reset_after_successful_setup(hass):
@ -245,7 +252,7 @@ async def test_get_controller(hass):
with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(
"aiounifi.Controller.login", return_value=True
):
assert await unifi.controller.get_controller(hass, **CONTROLLER_DATA)
assert await get_controller(hass, **CONTROLLER_DATA)
async def test_get_controller_verify_ssl_false(hass):
@ -255,28 +262,28 @@ async def test_get_controller_verify_ssl_false(hass):
with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(
"aiounifi.Controller.login", return_value=True
):
assert await unifi.controller.get_controller(hass, **controller_data)
assert await get_controller(hass, **controller_data)
async def test_get_controller_login_failed(hass):
"""Check that get_controller can handle a failed login."""
with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(
"aiounifi.Controller.login", side_effect=aiounifi.Unauthorized
), pytest.raises(unifi.errors.AuthenticationRequired):
await unifi.controller.get_controller(hass, **CONTROLLER_DATA)
), pytest.raises(AuthenticationRequired):
await get_controller(hass, **CONTROLLER_DATA)
async def test_get_controller_controller_unavailable(hass):
"""Check that get_controller can handle controller being unavailable."""
with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(
"aiounifi.Controller.login", side_effect=aiounifi.RequestError
), pytest.raises(unifi.errors.CannotConnect):
await unifi.controller.get_controller(hass, **CONTROLLER_DATA)
), pytest.raises(CannotConnect):
await get_controller(hass, **CONTROLLER_DATA)
async def test_get_controller_unknown_error(hass):
"""Check that get_controller can handle unknown errors."""
with patch("aiounifi.Controller.check_unifi_os", return_value=True), patch(
"aiounifi.Controller.login", side_effect=aiounifi.AiounifiException
), pytest.raises(unifi.errors.AuthenticationRequired):
await unifi.controller.get_controller(hass, **CONTROLLER_DATA)
), pytest.raises(AuthenticationRequired):
await get_controller(hass, **CONTROLLER_DATA)

View file

@ -7,8 +7,7 @@ from aiounifi.websocket import SIGNAL_DATA, STATE_DISCONNECTED, STATE_RUNNING
from asynctest import patch
from homeassistant import config_entries
from homeassistant.components import unifi
import homeassistant.components.device_tracker as device_tracker
from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
from homeassistant.components.unifi.const import (
CONF_BLOCK_CLIENT,
CONF_IGNORE_WIRED_BUG,
@ -16,6 +15,7 @@ from homeassistant.components.unifi.const import (
CONF_TRACK_CLIENTS,
CONF_TRACK_DEVICES,
CONF_TRACK_WIRED_CLIENTS,
DOMAIN as UNIFI_DOMAIN,
)
from homeassistant.const import STATE_UNAVAILABLE
from homeassistant.helpers import entity_registry
@ -99,18 +99,18 @@ async def test_platform_manually_configured(hass):
"""Test that nothing happens when configuring unifi through device tracker platform."""
assert (
await async_setup_component(
hass, device_tracker.DOMAIN, {device_tracker.DOMAIN: {"platform": "unifi"}}
hass, TRACKER_DOMAIN, {TRACKER_DOMAIN: {"platform": UNIFI_DOMAIN}}
)
is False
)
assert unifi.DOMAIN not in hass.data
assert UNIFI_DOMAIN not in hass.data
async def test_no_clients(hass):
"""Test the update_clients function when no clients are found."""
await setup_unifi_integration(hass)
assert len(hass.states.async_entity_ids("device_tracker")) == 0
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 0
async def test_tracked_devices(hass):
@ -125,7 +125,7 @@ async def test_tracked_devices(hass):
devices_response=[DEVICE_1, DEVICE_2],
known_wireless_clients=(CLIENT_4["mac"],),
)
assert len(hass.states.async_entity_ids("device_tracker")) == 5
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 5
client_1 = hass.states.get("device_tracker.client_1")
assert client_1 is not None
@ -186,7 +186,7 @@ async def test_remove_clients(hass):
controller = await setup_unifi_integration(
hass, clients_response=[CLIENT_1, CLIENT_2]
)
assert len(hass.states.async_entity_ids("device_tracker")) == 2
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
client_1 = hass.states.get("device_tracker.client_1")
assert client_1 is not None
@ -201,7 +201,7 @@ async def test_remove_clients(hass):
controller.api.session_handler(SIGNAL_DATA)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids("device_tracker")) == 1
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
client_1 = hass.states.get("device_tracker.client_1")
assert client_1 is None
@ -215,7 +215,7 @@ async def test_controller_state_change(hass):
controller = await setup_unifi_integration(
hass, clients_response=[CLIENT_1], devices_response=[DEVICE_1],
)
assert len(hass.states.async_entity_ids("device_tracker")) == 2
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
# Controller unavailable
controller.async_unifi_signalling_callback(
@ -245,7 +245,7 @@ async def test_option_track_clients(hass):
controller = await setup_unifi_integration(
hass, clients_response=[CLIENT_1, CLIENT_2], devices_response=[DEVICE_1],
)
assert len(hass.states.async_entity_ids("device_tracker")) == 3
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 3
client_1 = hass.states.get("device_tracker.client_1")
assert client_1 is not None
@ -290,7 +290,7 @@ async def test_option_track_wired_clients(hass):
controller = await setup_unifi_integration(
hass, clients_response=[CLIENT_1, CLIENT_2], devices_response=[DEVICE_1],
)
assert len(hass.states.async_entity_ids("device_tracker")) == 3
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 3
client_1 = hass.states.get("device_tracker.client_1")
assert client_1 is not None
@ -335,7 +335,7 @@ async def test_option_track_devices(hass):
controller = await setup_unifi_integration(
hass, clients_response=[CLIENT_1, CLIENT_2], devices_response=[DEVICE_1],
)
assert len(hass.states.async_entity_ids("device_tracker")) == 3
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 3
client_1 = hass.states.get("device_tracker.client_1")
assert client_1 is not None
@ -378,7 +378,7 @@ async def test_option_track_devices(hass):
async def test_option_ssid_filter(hass):
"""Test the SSID filter works."""
controller = await setup_unifi_integration(hass, clients_response=[CLIENT_3])
assert len(hass.states.async_entity_ids("device_tracker")) == 1
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
client_3 = hass.states.get("device_tracker.client_3")
assert client_3
@ -423,7 +423,7 @@ async def test_wireless_client_go_wired_issue(hass):
client_1_client["last_seen"] = dt_util.as_timestamp(dt_util.utcnow())
controller = await setup_unifi_integration(hass, clients_response=[client_1_client])
assert len(hass.states.async_entity_ids("device_tracker")) == 1
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
client_1 = hass.states.get("device_tracker.client_1")
assert client_1 is not None
@ -441,9 +441,7 @@ async def test_wireless_client_go_wired_issue(hass):
assert client_1.attributes["is_wired"] is False
with patch.object(
unifi.device_tracker.dt_util,
"utcnow",
return_value=(dt_util.utcnow() + timedelta(minutes=5)),
dt_util, "utcnow", return_value=(dt_util.utcnow() + timedelta(minutes=5)),
):
event = {"meta": {"message": "sta:sync"}, "data": [client_1_client]}
controller.api.message_handler(event)
@ -472,7 +470,7 @@ async def test_option_ignore_wired_bug(hass):
controller = await setup_unifi_integration(
hass, options={CONF_IGNORE_WIRED_BUG: True}, clients_response=[client_1_client]
)
assert len(hass.states.async_entity_ids("device_tracker")) == 1
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
client_1 = hass.states.get("device_tracker.client_1")
assert client_1 is not None
@ -504,7 +502,7 @@ async def test_restoring_client(hass):
"""Test the update_items function with some clients."""
config_entry = config_entries.ConfigEntry(
version=1,
domain=unifi.DOMAIN,
domain=UNIFI_DOMAIN,
title="Mock Title",
data=ENTRY_CONFIG,
source="test",
@ -516,16 +514,16 @@ async def test_restoring_client(hass):
registry = await entity_registry.async_get_registry(hass)
registry.async_get_or_create(
device_tracker.DOMAIN,
unifi.DOMAIN,
"{}-site_id".format(CLIENT_1["mac"]),
TRACKER_DOMAIN,
UNIFI_DOMAIN,
f'{CLIENT_1["mac"]}-site_id',
suggested_object_id=CLIENT_1["hostname"],
config_entry=config_entry,
)
registry.async_get_or_create(
device_tracker.DOMAIN,
unifi.DOMAIN,
"{}-site_id".format(CLIENT_2["mac"]),
TRACKER_DOMAIN,
UNIFI_DOMAIN,
f'{CLIENT_2["mac"]}-site_id',
suggested_object_id=CLIENT_2["hostname"],
config_entry=config_entry,
)
@ -536,7 +534,7 @@ async def test_restoring_client(hass):
clients_response=[CLIENT_2],
clients_all_response=[CLIENT_1],
)
assert len(hass.states.async_entity_ids("device_tracker")) == 2
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
device_1 = hass.states.get("device_tracker.client_1")
assert device_1 is not None
@ -546,11 +544,11 @@ async def test_dont_track_clients(hass):
"""Test don't track clients config works."""
await setup_unifi_integration(
hass,
options={unifi.controller.CONF_TRACK_CLIENTS: False},
options={CONF_TRACK_CLIENTS: False},
clients_response=[CLIENT_1],
devices_response=[DEVICE_1],
)
assert len(hass.states.async_entity_ids("device_tracker")) == 1
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
client_1 = hass.states.get("device_tracker.client_1")
assert client_1 is None
@ -564,11 +562,11 @@ async def test_dont_track_devices(hass):
"""Test don't track devices config works."""
await setup_unifi_integration(
hass,
options={unifi.controller.CONF_TRACK_DEVICES: False},
options={CONF_TRACK_DEVICES: False},
clients_response=[CLIENT_1],
devices_response=[DEVICE_1],
)
assert len(hass.states.async_entity_ids("device_tracker")) == 1
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
client_1 = hass.states.get("device_tracker.client_1")
assert client_1 is not None
@ -582,10 +580,10 @@ async def test_dont_track_wired_clients(hass):
"""Test don't track wired clients config works."""
await setup_unifi_integration(
hass,
options={unifi.controller.CONF_TRACK_WIRED_CLIENTS: False},
options={CONF_TRACK_WIRED_CLIENTS: False},
clients_response=[CLIENT_1, CLIENT_2],
)
assert len(hass.states.async_entity_ids("device_tracker")) == 1
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
client_1 = hass.states.get("device_tracker.client_1")
assert client_1 is not None

View file

@ -2,6 +2,7 @@
from unittest.mock import Mock, patch
from homeassistant.components import unifi
from homeassistant.components.unifi.const import DOMAIN as UNIFI_DOMAIN
from homeassistant.setup import async_setup_component
from .test_controller import setup_unifi_integration
@ -11,29 +12,29 @@ from tests.common import MockConfigEntry, mock_coro
async def test_setup_with_no_config(hass):
"""Test that we do not discover anything or try to set up a bridge."""
assert await async_setup_component(hass, unifi.DOMAIN, {}) is True
assert unifi.DOMAIN not in hass.data
assert await async_setup_component(hass, UNIFI_DOMAIN, {}) is True
assert UNIFI_DOMAIN not in hass.data
async def test_successful_config_entry(hass):
"""Test that configured options for a host are loaded via config entry."""
await setup_unifi_integration(hass)
assert hass.data[unifi.DOMAIN]
assert hass.data[UNIFI_DOMAIN]
async def test_controller_fail_setup(hass):
"""Test that a failed setup still stores controller."""
with patch.object(unifi, "UniFiController") as mock_cntrlr:
mock_cntrlr.return_value.async_setup.return_value = mock_coro(False)
with patch("homeassistant.components.unifi.UniFiController") as mock_controller:
mock_controller.return_value.async_setup.return_value = mock_coro(False)
await setup_unifi_integration(hass)
assert hass.data[unifi.DOMAIN] == {}
assert hass.data[UNIFI_DOMAIN] == {}
async def test_controller_no_mac(hass):
"""Test that configured options for a host are loaded via config entry."""
entry = MockConfigEntry(
domain=unifi.DOMAIN,
domain=UNIFI_DOMAIN,
data={
"controller": {
"host": "0.0.0.0",
@ -48,7 +49,9 @@ async def test_controller_no_mac(hass):
)
entry.add_to_hass(hass)
mock_registry = Mock()
with patch.object(unifi, "UniFiController") as mock_controller, patch(
with patch(
"homeassistant.components.unifi.UniFiController"
) as mock_controller, patch(
"homeassistant.helpers.device_registry.async_get_registry",
return_value=mock_coro(mock_registry),
):
@ -64,7 +67,7 @@ async def test_controller_no_mac(hass):
async def test_unload_entry(hass):
"""Test being able to unload an entry."""
controller = await setup_unifi_integration(hass)
assert hass.data[unifi.DOMAIN]
assert hass.data[UNIFI_DOMAIN]
assert await unifi.async_unload_entry(hass, controller.config_entry)
assert not hass.data[unifi.DOMAIN]
assert not hass.data[UNIFI_DOMAIN]

View file

@ -4,8 +4,14 @@ from copy import deepcopy
from aiounifi.controller import MESSAGE_CLIENT_REMOVED
from aiounifi.websocket import SIGNAL_DATA
from homeassistant.components import unifi
import homeassistant.components.sensor as sensor
from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
from homeassistant.components.sensor import DOMAIN as SENSOR_DOMAIN
from homeassistant.components.unifi.const import (
CONF_ALLOW_BANDWIDTH_SENSORS,
CONF_TRACK_CLIENTS,
CONF_TRACK_DEVICES,
DOMAIN as UNIFI_DOMAIN,
)
from homeassistant.setup import async_setup_component
from .test_controller import setup_unifi_integration
@ -44,21 +50,21 @@ async def test_platform_manually_configured(hass):
"""Test that we do not discover anything or try to set up a controller."""
assert (
await async_setup_component(
hass, sensor.DOMAIN, {sensor.DOMAIN: {"platform": "unifi"}}
hass, SENSOR_DOMAIN, {SENSOR_DOMAIN: {"platform": UNIFI_DOMAIN}}
)
is True
)
assert unifi.DOMAIN not in hass.data
assert UNIFI_DOMAIN not in hass.data
async def test_no_clients(hass):
"""Test the update_clients function when no clients are found."""
controller = await setup_unifi_integration(
hass, options={unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True},
hass, options={CONF_ALLOW_BANDWIDTH_SENSORS: True},
)
assert len(controller.mock_requests) == 4
assert len(hass.states.async_entity_ids("sensor")) == 0
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 0
async def test_sensors(hass):
@ -66,15 +72,15 @@ async def test_sensors(hass):
controller = await setup_unifi_integration(
hass,
options={
unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True,
unifi.const.CONF_TRACK_CLIENTS: False,
unifi.const.CONF_TRACK_DEVICES: False,
CONF_ALLOW_BANDWIDTH_SENSORS: True,
CONF_TRACK_CLIENTS: False,
CONF_TRACK_DEVICES: False,
},
clients_response=CLIENTS,
)
assert len(controller.mock_requests) == 4
assert len(hass.states.async_entity_ids("sensor")) == 4
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 4
wired_client_rx = hass.states.get("sensor.wired_client_name_rx")
assert wired_client_rx.state == "1234.0"
@ -104,8 +110,7 @@ async def test_sensors(hass):
assert wireless_client_tx.state == "6789.0"
hass.config_entries.async_update_entry(
controller.config_entry,
options={unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: False},
controller.config_entry, options={CONF_ALLOW_BANDWIDTH_SENSORS: False},
)
await hass.async_block_till_done()
@ -116,8 +121,7 @@ async def test_sensors(hass):
assert wireless_client_tx is None
hass.config_entries.async_update_entry(
controller.config_entry,
options={unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True},
controller.config_entry, options={CONF_ALLOW_BANDWIDTH_SENSORS: True},
)
await hass.async_block_till_done()
@ -131,12 +135,10 @@ async def test_sensors(hass):
async def test_remove_sensors(hass):
"""Test the remove_items function with some clients."""
controller = await setup_unifi_integration(
hass,
options={unifi.const.CONF_ALLOW_BANDWIDTH_SENSORS: True},
clients_response=CLIENTS,
hass, options={CONF_ALLOW_BANDWIDTH_SENSORS: True}, clients_response=CLIENTS,
)
assert len(hass.states.async_entity_ids("sensor")) == 4
assert len(hass.states.async_entity_ids("device_tracker")) == 2
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 4
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 2
wired_client_rx = hass.states.get("sensor.wired_client_name_rx")
assert wired_client_rx is not None
@ -155,8 +157,8 @@ async def test_remove_sensors(hass):
controller.api.session_handler(SIGNAL_DATA)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids("sensor")) == 2
assert len(hass.states.async_entity_ids("device_tracker")) == 1
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 2
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 1
wired_client_rx = hass.states.get("sensor.wired_client_name_rx")
assert wired_client_rx is None

View file

@ -5,13 +5,15 @@ from aiounifi.controller import MESSAGE_CLIENT_REMOVED
from aiounifi.websocket import SIGNAL_DATA
from homeassistant import config_entries
from homeassistant.components import unifi
import homeassistant.components.switch as switch
from homeassistant.components.device_tracker import DOMAIN as TRACKER_DOMAIN
from homeassistant.components.switch import DOMAIN as SWITCH_DOMAIN
from homeassistant.components.unifi.const import (
CONF_BLOCK_CLIENT,
CONF_TRACK_CLIENTS,
CONF_TRACK_DEVICES,
DOMAIN as UNIFI_DOMAIN,
)
from homeassistant.components.unifi.switch import POE_SWITCH
from homeassistant.helpers import entity_registry
from homeassistant.setup import async_setup_component
@ -200,11 +202,11 @@ async def test_platform_manually_configured(hass):
"""Test that we do not discover anything or try to set up a controller."""
assert (
await async_setup_component(
hass, switch.DOMAIN, {switch.DOMAIN: {"platform": "unifi"}}
hass, SWITCH_DOMAIN, {SWITCH_DOMAIN: {"platform": UNIFI_DOMAIN}}
)
is True
)
assert unifi.DOMAIN not in hass.data
assert UNIFI_DOMAIN not in hass.data
async def test_no_clients(hass):
@ -214,7 +216,7 @@ async def test_no_clients(hass):
)
assert len(controller.mock_requests) == 4
assert len(hass.states.async_entity_ids("switch")) == 0
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
async def test_controller_not_client(hass):
@ -227,7 +229,7 @@ async def test_controller_not_client(hass):
)
assert len(controller.mock_requests) == 4
assert len(hass.states.async_entity_ids("switch")) == 0
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
cloudkey = hass.states.get("switch.cloud_key")
assert cloudkey is None
@ -245,7 +247,7 @@ async def test_not_admin(hass):
)
assert len(controller.mock_requests) == 4
assert len(hass.states.async_entity_ids("switch")) == 0
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
async def test_switches(hass):
@ -263,13 +265,13 @@ async def test_switches(hass):
)
assert len(controller.mock_requests) == 4
assert len(hass.states.async_entity_ids("switch")) == 3
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 3
switch_1 = hass.states.get("switch.poe_client_1")
assert switch_1 is not None
assert switch_1.state == "on"
assert switch_1.attributes["power"] == "2.56"
assert switch_1.attributes["switch"] == "00:00:00:00:01:01"
assert switch_1.attributes[SWITCH_DOMAIN] == "00:00:00:00:01:01"
assert switch_1.attributes["port"] == 1
assert switch_1.attributes["poe_mode"] == "auto"
@ -285,7 +287,7 @@ async def test_switches(hass):
assert unblocked.state == "on"
await hass.services.async_call(
"switch", "turn_off", {"entity_id": "switch.block_client_1"}, blocking=True
SWITCH_DOMAIN, "turn_off", {"entity_id": "switch.block_client_1"}, blocking=True
)
assert len(controller.mock_requests) == 5
assert controller.mock_requests[4] == {
@ -295,7 +297,7 @@ async def test_switches(hass):
}
await hass.services.async_call(
"switch", "turn_on", {"entity_id": "switch.block_client_1"}, blocking=True
SWITCH_DOMAIN, "turn_on", {"entity_id": "switch.block_client_1"}, blocking=True
)
assert len(controller.mock_requests) == 6
assert controller.mock_requests[5] == {
@ -313,7 +315,7 @@ async def test_remove_switches(hass):
clients_response=[CLIENT_1, UNBLOCKED],
devices_response=[DEVICE_1],
)
assert len(hass.states.async_entity_ids("switch")) == 2
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
poe_switch = hass.states.get("switch.poe_client_1")
assert poe_switch is not None
@ -328,7 +330,7 @@ async def test_remove_switches(hass):
controller.api.session_handler(SIGNAL_DATA)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids("switch")) == 0
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
poe_switch = hass.states.get("switch.poe_client_1")
assert poe_switch is None
@ -349,7 +351,7 @@ async def test_new_client_discovered_on_block_control(hass):
)
assert len(controller.mock_requests) == 4
assert len(hass.states.async_entity_ids("switch")) == 0
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
blocked = hass.states.get("switch.block_client_1")
assert blocked is None
@ -361,7 +363,7 @@ async def test_new_client_discovered_on_block_control(hass):
controller.api.session_handler("data")
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids("switch")) == 1
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
blocked = hass.states.get("switch.block_client_1")
assert blocked is not None
@ -373,7 +375,7 @@ async def test_option_block_clients(hass):
options={CONF_BLOCK_CLIENT: [BLOCKED["mac"]]},
clients_all_response=[BLOCKED, UNBLOCKED],
)
assert len(hass.states.async_entity_ids("switch")) == 1
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
# Add a second switch
hass.config_entries.async_update_entry(
@ -381,28 +383,28 @@ async def test_option_block_clients(hass):
options={CONF_BLOCK_CLIENT: [BLOCKED["mac"], UNBLOCKED["mac"]]},
)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids("switch")) == 2
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
# Remove the second switch again
hass.config_entries.async_update_entry(
controller.config_entry, options={CONF_BLOCK_CLIENT: [BLOCKED["mac"]]},
)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids("switch")) == 1
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
# Enable one and remove another one
hass.config_entries.async_update_entry(
controller.config_entry, options={CONF_BLOCK_CLIENT: [UNBLOCKED["mac"]]},
)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids("switch")) == 1
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
# Remove one
hass.config_entries.async_update_entry(
controller.config_entry, options={CONF_BLOCK_CLIENT: []},
)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids("switch")) == 0
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
async def test_new_client_discovered_on_poe_control(hass):
@ -415,7 +417,7 @@ async def test_new_client_discovered_on_poe_control(hass):
)
assert len(controller.mock_requests) == 4
assert len(hass.states.async_entity_ids("switch")) == 1
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
controller.api.websocket._data = {
"meta": {"message": "sta:sync"},
@ -425,10 +427,10 @@ async def test_new_client_discovered_on_poe_control(hass):
# Calling a service will trigger the updates to run
await hass.services.async_call(
"switch", "turn_off", {"entity_id": "switch.poe_client_1"}, blocking=True
SWITCH_DOMAIN, "turn_off", {"entity_id": "switch.poe_client_1"}, blocking=True
)
assert len(controller.mock_requests) == 5
assert len(hass.states.async_entity_ids("switch")) == 2
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
assert controller.mock_requests[4] == {
"json": {
"port_overrides": [{"port_idx": 1, "portconf_id": "1a1", "poe_mode": "off"}]
@ -438,7 +440,7 @@ async def test_new_client_discovered_on_poe_control(hass):
}
await hass.services.async_call(
"switch", "turn_on", {"entity_id": "switch.poe_client_1"}, blocking=True
SWITCH_DOMAIN, "turn_on", {"entity_id": "switch.poe_client_1"}, blocking=True
)
assert len(controller.mock_requests) == 6
assert controller.mock_requests[4] == {
@ -467,7 +469,7 @@ async def test_ignore_multiple_poe_clients_on_same_port(hass):
)
assert len(controller.mock_requests) == 4
assert len(hass.states.async_entity_ids("device_tracker")) == 3
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 3
switch_1 = hass.states.get("switch.poe_client_1")
switch_2 = hass.states.get("switch.poe_client_2")
@ -479,7 +481,7 @@ async def test_restoring_client(hass):
"""Test the update_items function with some clients."""
config_entry = config_entries.ConfigEntry(
version=1,
domain=unifi.DOMAIN,
domain=UNIFI_DOMAIN,
title="Mock Title",
data=ENTRY_CONFIG,
source="test",
@ -491,16 +493,16 @@ async def test_restoring_client(hass):
registry = await entity_registry.async_get_registry(hass)
registry.async_get_or_create(
switch.DOMAIN,
unifi.DOMAIN,
"poe-{}".format(CLIENT_1["mac"]),
SWITCH_DOMAIN,
UNIFI_DOMAIN,
f'{POE_SWITCH}-{CLIENT_1["mac"]}',
suggested_object_id=CLIENT_1["hostname"],
config_entry=config_entry,
)
registry.async_get_or_create(
switch.DOMAIN,
unifi.DOMAIN,
"poe-{}".format(CLIENT_2["mac"]),
SWITCH_DOMAIN,
UNIFI_DOMAIN,
f'{POE_SWITCH}-{CLIENT_2["mac"]}',
suggested_object_id=CLIENT_2["hostname"],
config_entry=config_entry,
)
@ -518,7 +520,7 @@ async def test_restoring_client(hass):
)
assert len(controller.mock_requests) == 4
assert len(hass.states.async_entity_ids("switch")) == 2
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
device_1 = hass.states.get("switch.client_1")
assert device_1 is not None