From c3689d741622d669deb161d5c242d7c495813c66 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Thu, 23 Apr 2020 16:48:24 +0200 Subject: [PATCH] UniFi - Store controller in config_entry.entry_id (#34553) * Store controller in config_entry.entry_id * Clean up imports --- homeassistant/components/unifi/__init__.py | 30 ++++--- homeassistant/components/unifi/config_flow.py | 12 +-- homeassistant/components/unifi/controller.py | 15 ++-- .../components/unifi/device_tracker.py | 7 +- homeassistant/components/unifi/sensor.py | 4 +- homeassistant/components/unifi/switch.py | 4 +- .../components/unifi/unifi_client.py | 3 +- tests/components/unifi/test_config_flow.py | 17 ++-- tests/components/unifi/test_controller.py | 83 ++++++++++--------- tests/components/unifi/test_device_tracker.py | 62 +++++++------- tests/components/unifi/test_init.py | 23 ++--- tests/components/unifi/test_sensor.py | 44 +++++----- tests/components/unifi/test_switch.py | 68 +++++++-------- 13 files changed, 190 insertions(+), 182 deletions(-) diff --git a/homeassistant/components/unifi/__init__.py b/homeassistant/components/unifi/__init__.py index e3225a2d210..610d39b0fb9 100644 --- a/homeassistant/components/unifi/__init__.py +++ b/homeassistant/components/unifi/__init__.py @@ -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 diff --git a/homeassistant/components/unifi/config_flow.py b/homeassistant/components/unifi/config_flow.py index e91001e51be..f42acf54e9d 100644 --- a/homeassistant/components/unifi/config_flow.py +++ b/homeassistant/components/unifi/config_flow.py @@ -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() diff --git a/homeassistant/components/unifi/controller.py b/homeassistant/components/unifi/controller.py index 96f1194df24..807e727777c 100644 --- a/homeassistant/components/unifi/controller.py +++ b/homeassistant/components/unifi/controller.py @@ -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 diff --git a/homeassistant/components/unifi/device_tracker.py b/homeassistant/components/unifi/device_tracker.py index 2ac516fac55..4b72c520b31 100644 --- a/homeassistant/components/unifi/device_tracker.py +++ b/homeassistant/components/unifi/device_tracker.py @@ -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. diff --git a/homeassistant/components/unifi/sensor.py b/homeassistant/components/unifi/sensor.py index 964db5820b8..9077db49dac 100644 --- a/homeassistant/components/unifi/sensor.py +++ b/homeassistant/components/unifi/sensor.py @@ -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 diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 5fb6daf524c..73e1dd131bc 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -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": diff --git a/homeassistant/components/unifi/unifi_client.py b/homeassistant/components/unifi/unifi_client.py index 8f91d4e1de3..c9bd038dd77 100644 --- a/homeassistant/components/unifi/unifi_client.py +++ b/homeassistant/components/unifi/unifi_client.py @@ -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) diff --git a/tests/components/unifi/test_config_flow.py b/tests/components/unifi/test_config_flow.py index c6b4f27e7f4..bad191e1600 100644 --- a/tests/components/unifi/test_config_flow.py +++ b/tests/components/unifi/test_config_flow.py @@ -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 diff --git a/tests/components/unifi/test_controller.py b/tests/components/unifi/test_controller.py index 844eaa5d222..824ffbffc41 100644 --- a/tests/components/unifi/test_controller.py +++ b/tests/components/unifi/test_controller.py @@ -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) diff --git a/tests/components/unifi/test_device_tracker.py b/tests/components/unifi/test_device_tracker.py index c204c75a122..df7c49a1bf7 100644 --- a/tests/components/unifi/test_device_tracker.py +++ b/tests/components/unifi/test_device_tracker.py @@ -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 diff --git a/tests/components/unifi/test_init.py b/tests/components/unifi/test_init.py index 0ccc89cdb89..db6c1e30748 100644 --- a/tests/components/unifi/test_init.py +++ b/tests/components/unifi/test_init.py @@ -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] diff --git a/tests/components/unifi/test_sensor.py b/tests/components/unifi/test_sensor.py index 3c801474235..cc879a3d713 100644 --- a/tests/components/unifi/test_sensor.py +++ b/tests/components/unifi/test_sensor.py @@ -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 diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index 5ea76472739..4dfd3bc4c9d 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -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