Fixture cleanup in UniFi tests (#119227)

* Make sure all mock_device_registry are used with usefixtuers

* Make sure to use name with fixtures and rename functions to start with fixture_

* Streamline config_entry_setup

* Type all *_payload

* Mark @pytest.mark.usefixtures("mock_default_requests")

* Clean up unnecessary newlines after docstring
This commit is contained in:
Robert Svensson 2024-06-09 22:07:36 +02:00 committed by GitHub
parent 7dfaa05793
commit 325352e197
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 95 additions and 137 deletions

View File

@ -54,8 +54,8 @@ CONTROLLER_HOST = {
} }
@pytest.fixture(autouse=True) @pytest.fixture(autouse=True, name="mock_discovery")
def mock_discovery(): def fixture_discovery():
"""No real network traffic allowed.""" """No real network traffic allowed."""
with patch( with patch(
"homeassistant.components.unifi.config_flow._async_discover_unifi", "homeassistant.components.unifi.config_flow._async_discover_unifi",
@ -64,8 +64,8 @@ def mock_discovery():
yield mock yield mock
@pytest.fixture @pytest.fixture(name="mock_device_registry")
def mock_device_registry(hass, device_registry: dr.DeviceRegistry): def fixture_device_registry(hass, device_registry: dr.DeviceRegistry):
"""Mock device registry.""" """Mock device registry."""
config_entry = MockConfigEntry(domain="something_else") config_entry = MockConfigEntry(domain="something_else")
config_entry.add_to_hass(hass) config_entry.add_to_hass(hass)
@ -93,7 +93,7 @@ def mock_device_registry(hass, device_registry: dr.DeviceRegistry):
@pytest.fixture(name="config_entry") @pytest.fixture(name="config_entry")
def config_entry_fixture( def fixture_config_entry(
hass: HomeAssistant, hass: HomeAssistant,
config_entry_data: MappingProxyType[str, Any], config_entry_data: MappingProxyType[str, Any],
config_entry_options: MappingProxyType[str, Any], config_entry_options: MappingProxyType[str, Any],
@ -111,7 +111,7 @@ def config_entry_fixture(
@pytest.fixture(name="config_entry_data") @pytest.fixture(name="config_entry_data")
def config_entry_data_fixture() -> MappingProxyType[str, Any]: def fixture_config_entry_data() -> MappingProxyType[str, Any]:
"""Define a config entry data fixture.""" """Define a config entry data fixture."""
return { return {
CONF_HOST: DEFAULT_HOST, CONF_HOST: DEFAULT_HOST,
@ -124,7 +124,7 @@ def config_entry_data_fixture() -> MappingProxyType[str, Any]:
@pytest.fixture(name="config_entry_options") @pytest.fixture(name="config_entry_options")
def config_entry_options_fixture() -> MappingProxyType[str, Any]: def fixture_config_entry_options() -> MappingProxyType[str, Any]:
"""Define a config entry options fixture.""" """Define a config entry options fixture."""
return {} return {}
@ -133,13 +133,13 @@ def config_entry_options_fixture() -> MappingProxyType[str, Any]:
@pytest.fixture(name="known_wireless_clients") @pytest.fixture(name="known_wireless_clients")
def known_wireless_clients_fixture() -> list[str]: def fixture_known_wireless_clients() -> list[str]:
"""Known previously observed wireless clients.""" """Known previously observed wireless clients."""
return [] return []
@pytest.fixture(autouse=True) @pytest.fixture(autouse=True, name="mock_wireless_client_storage")
def mock_wireless_client_storage(hass_storage, known_wireless_clients: list[str]): def fixture_wireless_client_storage(hass_storage, known_wireless_clients: list[str]):
"""Mock the known wireless storage.""" """Mock the known wireless storage."""
data: dict[str, list[str]] = ( data: dict[str, list[str]] = (
{"wireless_clients": known_wireless_clients} if known_wireless_clients else {} {"wireless_clients": known_wireless_clients} if known_wireless_clients else {}
@ -151,7 +151,7 @@ def mock_wireless_client_storage(hass_storage, known_wireless_clients: list[str]
@pytest.fixture(name="mock_requests") @pytest.fixture(name="mock_requests")
def request_fixture( def fixture_request(
aioclient_mock: AiohttpClientMocker, aioclient_mock: AiohttpClientMocker,
client_payload: list[dict[str, Any]], client_payload: list[dict[str, Any]],
clients_all_payload: list[dict[str, Any]], clients_all_payload: list[dict[str, Any]],
@ -198,49 +198,49 @@ def request_fixture(
@pytest.fixture(name="client_payload") @pytest.fixture(name="client_payload")
def client_data_fixture() -> list[dict[str, Any]]: def fixture_client_data() -> list[dict[str, Any]]:
"""Client data.""" """Client data."""
return [] return []
@pytest.fixture(name="clients_all_payload") @pytest.fixture(name="clients_all_payload")
def clients_all_data_fixture() -> list[dict[str, Any]]: def fixture_clients_all_data() -> list[dict[str, Any]]:
"""Clients all data.""" """Clients all data."""
return [] return []
@pytest.fixture(name="device_payload") @pytest.fixture(name="device_payload")
def device_data_fixture() -> list[dict[str, Any]]: def fixture_device_data() -> list[dict[str, Any]]:
"""Device data.""" """Device data."""
return [] return []
@pytest.fixture(name="dpi_app_payload") @pytest.fixture(name="dpi_app_payload")
def dpi_app_data_fixture() -> list[dict[str, Any]]: def fixture_dpi_app_data() -> list[dict[str, Any]]:
"""DPI app data.""" """DPI app data."""
return [] return []
@pytest.fixture(name="dpi_group_payload") @pytest.fixture(name="dpi_group_payload")
def dpi_group_data_fixture() -> list[dict[str, Any]]: def fixture_dpi_group_data() -> list[dict[str, Any]]:
"""DPI group data.""" """DPI group data."""
return [] return []
@pytest.fixture(name="port_forward_payload") @pytest.fixture(name="port_forward_payload")
def port_forward_data_fixture() -> list[dict[str, Any]]: def fixture_port_forward_data() -> list[dict[str, Any]]:
"""Port forward data.""" """Port forward data."""
return [] return []
@pytest.fixture(name="site_payload") @pytest.fixture(name="site_payload")
def site_data_fixture() -> list[dict[str, Any]]: def fixture_site_data() -> list[dict[str, Any]]:
"""Site data.""" """Site data."""
return [{"desc": "Site name", "name": "site_id", "role": "admin", "_id": "1"}] return [{"desc": "Site name", "name": "site_id", "role": "admin", "_id": "1"}]
@pytest.fixture(name="system_information_payload") @pytest.fixture(name="system_information_payload")
def system_information_data_fixture() -> list[dict[str, Any]]: def fixture_system_information_data() -> list[dict[str, Any]]:
"""System information data.""" """System information data."""
return [ return [
{ {
@ -262,13 +262,13 @@ def system_information_data_fixture() -> list[dict[str, Any]]:
@pytest.fixture(name="wlan_payload") @pytest.fixture(name="wlan_payload")
def wlan_data_fixture() -> list[dict[str, Any]]: def fixture_wlan_data() -> list[dict[str, Any]]:
"""WLAN data.""" """WLAN data."""
return [] return []
@pytest.fixture(name="mock_default_requests") @pytest.fixture(name="mock_default_requests")
def default_requests_fixture( def fixture_default_requests(
mock_requests: Callable[[str, str], None], mock_requests: Callable[[str, str], None],
) -> None: ) -> None:
"""Mock UniFi requests responses with default host and site.""" """Mock UniFi requests responses with default host and site."""
@ -276,7 +276,7 @@ def default_requests_fixture(
@pytest.fixture(name="config_entry_factory") @pytest.fixture(name="config_entry_factory")
async def config_entry_factory_fixture( async def fixture_config_entry_factory(
hass: HomeAssistant, hass: HomeAssistant,
config_entry: ConfigEntry, config_entry: ConfigEntry,
mock_requests: Callable[[str, str], None], mock_requests: Callable[[str, str], None],
@ -293,7 +293,7 @@ async def config_entry_factory_fixture(
@pytest.fixture(name="config_entry_setup") @pytest.fixture(name="config_entry_setup")
async def config_entry_setup_fixture( async def fixture_config_entry_setup(
hass: HomeAssistant, config_entry_factory: Callable[[], ConfigEntry] hass: HomeAssistant, config_entry_factory: Callable[[], ConfigEntry]
) -> ConfigEntry: ) -> ConfigEntry:
"""Fixture providing a set up instance of UniFi network integration.""" """Fixture providing a set up instance of UniFi network integration."""

View File

@ -95,9 +95,8 @@ DPI_GROUPS = [
] ]
async def test_flow_works( @pytest.mark.usefixtures("mock_default_requests")
hass: HomeAssistant, mock_discovery, mock_default_requests: None async def test_flow_works(hass: HomeAssistant, mock_discovery) -> None:
) -> None:
"""Test config flow.""" """Test config flow."""
mock_discovery.return_value = "1" mock_discovery.return_value = "1"
result = await hass.config_entries.flow.async_init( result = await hass.config_entries.flow.async_init(
@ -165,9 +164,8 @@ async def test_flow_works_negative_discovery(
] ]
], ],
) )
async def test_flow_multiple_sites( @pytest.mark.usefixtures("mock_default_requests")
hass: HomeAssistant, mock_default_requests: None async def test_flow_multiple_sites(hass: HomeAssistant) -> None:
) -> None:
"""Test config flow works when finding multiple sites.""" """Test config flow works when finding multiple sites."""
result = await hass.config_entries.flow.async_init( result = await hass.config_entries.flow.async_init(
UNIFI_DOMAIN, context={"source": config_entries.SOURCE_USER} UNIFI_DOMAIN, context={"source": config_entries.SOURCE_USER}
@ -193,9 +191,8 @@ async def test_flow_multiple_sites(
assert result["data_schema"]({"site": "2"}) assert result["data_schema"]({"site": "2"})
async def test_flow_raise_already_configured( @pytest.mark.usefixtures("config_entry_setup")
hass: HomeAssistant, config_entry_setup: ConfigEntry async def test_flow_raise_already_configured(hass: HomeAssistant) -> None:
) -> None:
"""Test config flow aborts since a connected config entry already exists.""" """Test config flow aborts since a connected config entry already exists."""
result = await hass.config_entries.flow.async_init( result = await hass.config_entries.flow.async_init(
UNIFI_DOMAIN, context={"source": config_entries.SOURCE_USER} UNIFI_DOMAIN, context={"source": config_entries.SOURCE_USER}
@ -219,9 +216,8 @@ async def test_flow_raise_already_configured(
assert result["reason"] == "already_configured" assert result["reason"] == "already_configured"
async def test_flow_aborts_configuration_updated( @pytest.mark.usefixtures("config_entry_setup")
hass: HomeAssistant, config_entry_setup: ConfigEntry async def test_flow_aborts_configuration_updated(hass: HomeAssistant) -> None:
) -> None:
"""Test config flow aborts since a connected config entry already exists.""" """Test config flow aborts since a connected config entry already exists."""
result = await hass.config_entries.flow.async_init( result = await hass.config_entries.flow.async_init(
UNIFI_DOMAIN, context={"source": config_entries.SOURCE_USER} UNIFI_DOMAIN, context={"source": config_entries.SOURCE_USER}
@ -249,9 +245,8 @@ async def test_flow_aborts_configuration_updated(
assert result["reason"] == "configuration_updated" assert result["reason"] == "configuration_updated"
async def test_flow_fails_user_credentials_faulty( @pytest.mark.usefixtures("mock_default_requests")
hass: HomeAssistant, mock_default_requests: None async def test_flow_fails_user_credentials_faulty(hass: HomeAssistant) -> None:
) -> None:
"""Test config flow.""" """Test config flow."""
result = await hass.config_entries.flow.async_init( result = await hass.config_entries.flow.async_init(
UNIFI_DOMAIN, context={"source": config_entries.SOURCE_USER} UNIFI_DOMAIN, context={"source": config_entries.SOURCE_USER}
@ -276,9 +271,8 @@ async def test_flow_fails_user_credentials_faulty(
assert result["errors"] == {"base": "faulty_credentials"} assert result["errors"] == {"base": "faulty_credentials"}
async def test_flow_fails_hub_unavailable( @pytest.mark.usefixtures("mock_default_requests")
hass: HomeAssistant, mock_default_requests: None async def test_flow_fails_hub_unavailable(hass: HomeAssistant) -> None:
) -> None:
"""Test config flow.""" """Test config flow."""
result = await hass.config_entries.flow.async_init( result = await hass.config_entries.flow.async_init(
UNIFI_DOMAIN, context={"source": config_entries.SOURCE_USER} UNIFI_DOMAIN, context={"source": config_entries.SOURCE_USER}
@ -465,7 +459,6 @@ async def test_simple_option_flow(
async def test_form_ssdp(hass: HomeAssistant) -> None: async def test_form_ssdp(hass: HomeAssistant) -> None:
"""Test we get the form with ssdp source.""" """Test we get the form with ssdp source."""
result = await hass.config_entries.flow.async_init( result = await hass.config_entries.flow.async_init(
UNIFI_DOMAIN, UNIFI_DOMAIN,
context={"source": config_entries.SOURCE_SSDP}, context={"source": config_entries.SOURCE_SSDP},
@ -507,7 +500,6 @@ async def test_form_ssdp_aborts_if_host_already_exists(
hass: HomeAssistant, config_entry: ConfigEntry hass: HomeAssistant, config_entry: ConfigEntry
) -> None: ) -> None:
"""Test we abort if the host is already configured.""" """Test we abort if the host is already configured."""
result = await hass.config_entries.flow.async_init( result = await hass.config_entries.flow.async_init(
UNIFI_DOMAIN, UNIFI_DOMAIN,
context={"source": config_entries.SOURCE_SSDP}, context={"source": config_entries.SOURCE_SSDP},
@ -551,7 +543,6 @@ async def test_form_ssdp_aborts_if_serial_already_exists(
async def test_form_ssdp_gets_form_with_ignored_entry(hass: HomeAssistant) -> None: async def test_form_ssdp_gets_form_with_ignored_entry(hass: HomeAssistant) -> None:
"""Test we can still setup if there is an ignored never configured entry.""" """Test we can still setup if there is an ignored never configured entry."""
entry = MockConfigEntry( entry = MockConfigEntry(
domain=UNIFI_DOMAIN, domain=UNIFI_DOMAIN,
data={"not_controller_key": None}, data={"not_controller_key": None},

View File

@ -545,7 +545,6 @@ async def test_option_track_clients(
hass: HomeAssistant, config_entry_setup: ConfigEntry hass: HomeAssistant, config_entry_setup: ConfigEntry
) -> None: ) -> None:
"""Test the tracking of clients can be turned off.""" """Test the tracking of clients can be turned off."""
assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 3 assert len(hass.states.async_entity_ids(TRACKER_DOMAIN)) == 3
assert hass.states.get("device_tracker.wireless_client") assert hass.states.get("device_tracker.wireless_client")
assert hass.states.get("device_tracker.wired_client") assert hass.states.get("device_tracker.wired_client")

View File

@ -77,9 +77,9 @@ async def test_reset_fails(
assert config_entry_setup.state is ConfigEntryState.LOADED assert config_entry_setup.state is ConfigEntryState.LOADED
@pytest.mark.usefixtures("mock_device_registry")
async def test_connection_state_signalling( async def test_connection_state_signalling(
hass: HomeAssistant, hass: HomeAssistant,
mock_device_registry,
mock_websocket_state, mock_websocket_state,
config_entry_factory: Callable[[], ConfigEntry], config_entry_factory: Callable[[], ConfigEntry],
client_payload: list[dict[str, Any]], client_payload: list[dict[str, Any]],

View File

@ -1045,7 +1045,6 @@ async def test_device_system_stats(
device_payload: list[dict[str, Any]], device_payload: list[dict[str, Any]],
) -> None: ) -> None:
"""Verify that device stats sensors are working as expected.""" """Verify that device stats sensors are working as expected."""
assert len(hass.states.async_all()) == 8 assert len(hass.states.async_all()) == 8
assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 4 assert len(hass.states.async_entity_ids(SENSOR_DOMAIN)) == 4
@ -1134,14 +1133,13 @@ async def test_device_system_stats(
] ]
], ],
) )
@pytest.mark.usefixtures("config_entry_setup")
async def test_bandwidth_port_sensors( async def test_bandwidth_port_sensors(
hass: HomeAssistant, hass: HomeAssistant,
entity_registry: er.EntityRegistry, entity_registry: er.EntityRegistry,
mock_websocket_message, mock_websocket_message,
config_entry_setup: ConfigEntry, config_entry_setup: ConfigEntry,
config_entry_options: MappingProxyType[str, Any], config_entry_options: MappingProxyType[str, Any],
device_payload, device_payload: list[dict[str, Any]],
) -> None: ) -> None:
"""Verify that port bandwidth sensors are working as expected.""" """Verify that port bandwidth sensors are working as expected."""
assert len(hass.states.async_all()) == 5 assert len(hass.states.async_all()) == 5

View File

@ -29,16 +29,14 @@ async def test_reconnect_client(
client_payload: list[dict[str, Any]], client_payload: list[dict[str, Any]],
) -> None: ) -> None:
"""Verify call to reconnect client is performed as expected.""" """Verify call to reconnect client is performed as expected."""
config_entry = config_entry_setup
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.post( aioclient_mock.post(
f"https://{config_entry.data[CONF_HOST]}:1234" f"https://{config_entry_setup.data[CONF_HOST]}:1234"
f"/api/s/{config_entry.data[CONF_SITE_ID]}/cmd/stamgr", f"/api/s/{config_entry_setup.data[CONF_SITE_ID]}/cmd/stamgr",
) )
device_entry = device_registry.async_get_or_create( device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id, config_entry_id=config_entry_setup.entry_id,
connections={(dr.CONNECTION_NETWORK_MAC, client_payload[0]["mac"])}, connections={(dr.CONNECTION_NETWORK_MAC, client_payload[0]["mac"])},
) )
@ -74,12 +72,10 @@ async def test_reconnect_device_without_mac(
config_entry_setup: ConfigEntry, config_entry_setup: ConfigEntry,
) -> None: ) -> None:
"""Verify no call is made if device does not have a known mac.""" """Verify no call is made if device does not have a known mac."""
config_entry = config_entry_setup
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
device_entry = device_registry.async_get_or_create( device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id, config_entry_id=config_entry_setup.entry_id,
connections={("other connection", "not mac")}, connections={("other connection", "not mac")},
) )
@ -103,16 +99,14 @@ async def test_reconnect_client_hub_unavailable(
client_payload: list[dict[str, Any]], client_payload: list[dict[str, Any]],
) -> None: ) -> None:
"""Verify no call is made if hub is unavailable.""" """Verify no call is made if hub is unavailable."""
config_entry = config_entry_setup
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.post( aioclient_mock.post(
f"https://{config_entry.data[CONF_HOST]}:1234" f"https://{config_entry_setup.data[CONF_HOST]}:1234"
f"/api/s/{config_entry.data[CONF_SITE_ID]}/cmd/stamgr", f"/api/s/{config_entry_setup.data[CONF_SITE_ID]}/cmd/stamgr",
) )
device_entry = device_registry.async_get_or_create( device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id, config_entry_id=config_entry_setup.entry_id,
connections={(dr.CONNECTION_NETWORK_MAC, client_payload[0]["mac"])}, connections={(dr.CONNECTION_NETWORK_MAC, client_payload[0]["mac"])},
) )
@ -136,12 +130,9 @@ async def test_reconnect_client_unknown_mac(
config_entry_setup: ConfigEntry, config_entry_setup: ConfigEntry,
) -> None: ) -> None:
"""Verify no call is made if trying to reconnect a mac unknown to hub.""" """Verify no call is made if trying to reconnect a mac unknown to hub."""
config_entry = config_entry_setup
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
device_entry = device_registry.async_get_or_create( device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id, config_entry_id=config_entry_setup.entry_id,
connections={(dr.CONNECTION_NETWORK_MAC, "mac unknown to hub")}, connections={(dr.CONNECTION_NETWORK_MAC, "mac unknown to hub")},
) )
@ -165,12 +156,9 @@ async def test_reconnect_wired_client(
client_payload: list[dict[str, Any]], client_payload: list[dict[str, Any]],
) -> None: ) -> None:
"""Verify no call is made if client is wired.""" """Verify no call is made if client is wired."""
config_entry = config_entry_setup
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
device_entry = device_registry.async_get_or_create( device_entry = device_registry.async_get_or_create(
config_entry_id=config_entry.entry_id, config_entry_id=config_entry_setup.entry_id,
connections={(dr.CONNECTION_NETWORK_MAC, client_payload[0]["mac"])}, connections={(dr.CONNECTION_NETWORK_MAC, client_payload[0]["mac"])},
) )
@ -219,12 +207,10 @@ async def test_remove_clients(
config_entry_setup: ConfigEntry, config_entry_setup: ConfigEntry,
) -> None: ) -> None:
"""Verify removing different variations of clients work.""" """Verify removing different variations of clients work."""
config_entry = config_entry_setup
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.post( aioclient_mock.post(
f"https://{config_entry.data[CONF_HOST]}:1234" f"https://{config_entry_setup.data[CONF_HOST]}:1234"
f"/api/s/{config_entry.data[CONF_SITE_ID]}/cmd/stamgr", f"/api/s/{config_entry_setup.data[CONF_SITE_ID]}/cmd/stamgr",
) )
await hass.services.async_call(UNIFI_DOMAIN, SERVICE_REMOVE_CLIENTS, blocking=True) await hass.services.async_call(UNIFI_DOMAIN, SERVICE_REMOVE_CLIENTS, blocking=True)
@ -233,7 +219,7 @@ async def test_remove_clients(
"macs": ["00:00:00:00:00:00", "00:00:00:00:00:01"], "macs": ["00:00:00:00:00:00", "00:00:00:00:00:01"],
} }
assert await hass.config_entries.async_unload(config_entry.entry_id) assert await hass.config_entries.async_unload(config_entry_setup.entry_id)
@pytest.mark.parametrize( @pytest.mark.parametrize(
@ -254,7 +240,6 @@ async def test_remove_clients_hub_unavailable(
) -> None: ) -> None:
"""Verify no call is made if UniFi Network is unavailable.""" """Verify no call is made if UniFi Network is unavailable."""
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
with patch( with patch(
"homeassistant.components.unifi.UnifiHub.available", new_callable=PropertyMock "homeassistant.components.unifi.UnifiHub.available", new_callable=PropertyMock
) as ws_mock: ) as ws_mock:
@ -283,6 +268,5 @@ async def test_remove_clients_no_call_on_empty_list(
) -> None: ) -> None:
"""Verify no call is made if no fitting client has been added to the list.""" """Verify no call is made if no fitting client has been added to the list."""
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
await hass.services.async_call(UNIFI_DOMAIN, SERVICE_REMOVE_CLIENTS, blocking=True) await hass.services.async_call(UNIFI_DOMAIN, SERVICE_REMOVE_CLIENTS, blocking=True)
assert aioclient_mock.call_count == 0 assert aioclient_mock.call_count == 0

View File

@ -1,7 +1,9 @@
"""UniFi Network switch platform tests.""" """UniFi Network switch platform tests."""
from collections.abc import Callable
from copy import deepcopy from copy import deepcopy
from datetime import timedelta from datetime import timedelta
from typing import Any
from aiounifi.models.message import MessageKey from aiounifi.models.message import MessageKey
import pytest import pytest
@ -20,7 +22,7 @@ from homeassistant.components.unifi.const import (
CONF_TRACK_DEVICES, CONF_TRACK_DEVICES,
DOMAIN as UNIFI_DOMAIN, DOMAIN as UNIFI_DOMAIN,
) )
from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY from homeassistant.config_entries import RELOAD_AFTER_UPDATE_DELAY, ConfigEntry
from homeassistant.const import ( from homeassistant.const import (
ATTR_DEVICE_CLASS, ATTR_DEVICE_CLASS,
ATTR_ENTITY_ID, ATTR_ENTITY_ID,
@ -800,11 +802,9 @@ async def test_switches(
hass: HomeAssistant, hass: HomeAssistant,
entity_registry: er.EntityRegistry, entity_registry: er.EntityRegistry,
aioclient_mock: AiohttpClientMocker, aioclient_mock: AiohttpClientMocker,
config_entry_setup, config_entry_setup: ConfigEntry,
) -> None: ) -> None:
"""Test the update_items function with some clients.""" """Test the update_items function with some clients."""
config_entry = config_entry_setup
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 3 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 3
switch_4 = hass.states.get("switch.poe_client_4") switch_4 = hass.states.get("switch.poe_client_4")
@ -831,8 +831,8 @@ async def test_switches(
# Block and unblock client # Block and unblock client
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.post( aioclient_mock.post(
f"https://{config_entry.data[CONF_HOST]}:1234" f"https://{config_entry_setup.data[CONF_HOST]}:1234"
f"/api/s/{config_entry.data[CONF_SITE_ID]}/cmd/stamgr", f"/api/s/{config_entry_setup.data[CONF_SITE_ID]}/cmd/stamgr",
) )
await hass.services.async_call( await hass.services.async_call(
@ -856,8 +856,8 @@ async def test_switches(
# Enable and disable DPI # Enable and disable DPI
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.put( aioclient_mock.put(
f"https://{config_entry.data[CONF_HOST]}:1234" f"https://{config_entry_setup.data[CONF_HOST]}:1234"
f"/api/s/{config_entry.data[CONF_SITE_ID]}/rest/dpiapp/{DPI_APPS[0]['_id']}", f"/api/s/{config_entry_setup.data[CONF_SITE_ID]}/rest/dpiapp/{DPI_APPS[0]['_id']}",
) )
await hass.services.async_call( await hass.services.async_call(
@ -924,11 +924,9 @@ async def test_block_switches(
hass: HomeAssistant, hass: HomeAssistant,
aioclient_mock: AiohttpClientMocker, aioclient_mock: AiohttpClientMocker,
mock_websocket_message, mock_websocket_message,
config_entry_setup, config_entry_setup: ConfigEntry,
) -> None: ) -> None:
"""Test the update_items function with some clients.""" """Test the update_items function with some clients."""
config_entry = config_entry_setup
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 2
blocked = hass.states.get("switch.block_client_1") blocked = hass.states.get("switch.block_client_1")
@ -959,8 +957,8 @@ async def test_block_switches(
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.post( aioclient_mock.post(
f"https://{config_entry.data[CONF_HOST]}:1234" f"https://{config_entry_setup.data[CONF_HOST]}:1234"
f"/api/s/{config_entry.data[CONF_SITE_ID]}/cmd/stamgr", f"/api/s/{config_entry_setup.data[CONF_SITE_ID]}/cmd/stamgr",
) )
await hass.services.async_call( await hass.services.async_call(
@ -1085,16 +1083,14 @@ async def test_outlet_switches(
hass: HomeAssistant, hass: HomeAssistant,
aioclient_mock: AiohttpClientMocker, aioclient_mock: AiohttpClientMocker,
mock_websocket_message, mock_websocket_message,
config_entry_setup, config_entry_setup: ConfigEntry,
device_payload, device_payload: list[dict[str, Any]],
mock_websocket_state, mock_websocket_state,
entity_id: str, entity_id: str,
outlet_index: int, outlet_index: int,
expected_switches: int, expected_switches: int,
) -> None: ) -> None:
"""Test the outlet entities.""" """Test the outlet entities."""
config_entry = config_entry_setup
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == expected_switches assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == expected_switches
# Validate state object # Validate state object
@ -1114,8 +1110,8 @@ async def test_outlet_switches(
device_id = device_payload[0]["device_id"] device_id = device_payload[0]["device_id"]
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.put( aioclient_mock.put(
f"https://{config_entry.data[CONF_HOST]}:1234" f"https://{config_entry_setup.data[CONF_HOST]}:1234"
f"/api/s/{config_entry.data[CONF_SITE_ID]}/rest/device/{device_id}", f"/api/s/{config_entry_setup.data[CONF_SITE_ID]}/rest/device/{device_id}",
) )
await hass.services.async_call( await hass.services.async_call(
@ -1171,11 +1167,11 @@ async def test_outlet_switches(
assert hass.states.get(f"switch.{entity_id}").state == STATE_OFF assert hass.states.get(f"switch.{entity_id}").state == STATE_OFF
# Unload config entry # Unload config entry
await hass.config_entries.async_unload(config_entry.entry_id) await hass.config_entries.async_unload(config_entry_setup.entry_id)
assert hass.states.get(f"switch.{entity_id}").state == STATE_UNAVAILABLE assert hass.states.get(f"switch.{entity_id}").state == STATE_UNAVAILABLE
# Remove config entry # Remove config entry
await hass.config_entries.async_remove(config_entry.entry_id) await hass.config_entries.async_remove(config_entry_setup.entry_id)
await hass.async_block_till_done() await hass.async_block_till_done()
assert hass.states.get(f"switch.{entity_id}") is None assert hass.states.get(f"switch.{entity_id}") is None
@ -1211,16 +1207,14 @@ async def test_new_client_discovered_on_block_control(
) )
@pytest.mark.parametrize("clients_all_payload", [[BLOCKED, UNBLOCKED]]) @pytest.mark.parametrize("clients_all_payload", [[BLOCKED, UNBLOCKED]])
async def test_option_block_clients( async def test_option_block_clients(
hass: HomeAssistant, config_entry_setup, clients_all_payload hass: HomeAssistant, config_entry_setup: ConfigEntry, clients_all_payload
) -> None: ) -> None:
"""Test the changes to option reflects accordingly.""" """Test the changes to option reflects accordingly."""
config_entry = config_entry_setup
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
# Add a second switch # Add a second switch
hass.config_entries.async_update_entry( hass.config_entries.async_update_entry(
config_entry, config_entry_setup,
options={ options={
CONF_BLOCK_CLIENT: [ CONF_BLOCK_CLIENT: [
clients_all_payload[0]["mac"], clients_all_payload[0]["mac"],
@ -1233,24 +1227,21 @@ async def test_option_block_clients(
# Remove the second switch again # Remove the second switch again
hass.config_entries.async_update_entry( hass.config_entries.async_update_entry(
config_entry, config_entry_setup, options={CONF_BLOCK_CLIENT: [clients_all_payload[0]["mac"]]}
options={CONF_BLOCK_CLIENT: [clients_all_payload[0]["mac"]]},
) )
await hass.async_block_till_done() await hass.async_block_till_done()
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
# Enable one and remove the other one # Enable one and remove the other one
hass.config_entries.async_update_entry( hass.config_entries.async_update_entry(
config_entry, config_entry_setup, options={CONF_BLOCK_CLIENT: [clients_all_payload[1]["mac"]]}
options={CONF_BLOCK_CLIENT: [clients_all_payload[1]["mac"]]},
) )
await hass.async_block_till_done() await hass.async_block_till_done()
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
# Remove one # Remove one
hass.config_entries.async_update_entry( hass.config_entries.async_update_entry(
config_entry, config_entry_setup, options={CONF_BLOCK_CLIENT: []}
options={CONF_BLOCK_CLIENT: []},
) )
await hass.async_block_till_done() await hass.async_block_till_done()
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
@ -1263,16 +1254,15 @@ async def test_option_block_clients(
@pytest.mark.parametrize("client_payload", [[CLIENT_1]]) @pytest.mark.parametrize("client_payload", [[CLIENT_1]])
@pytest.mark.parametrize("dpi_app_payload", [DPI_APPS]) @pytest.mark.parametrize("dpi_app_payload", [DPI_APPS])
@pytest.mark.parametrize("dpi_group_payload", [DPI_GROUPS]) @pytest.mark.parametrize("dpi_group_payload", [DPI_GROUPS])
async def test_option_remove_switches(hass: HomeAssistant, config_entry_setup) -> None: async def test_option_remove_switches(
hass: HomeAssistant, config_entry_setup: ConfigEntry
) -> None:
"""Test removal of DPI switch when options updated.""" """Test removal of DPI switch when options updated."""
config_entry = config_entry_setup
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
# Disable DPI Switches # Disable DPI Switches
hass.config_entries.async_update_entry( hass.config_entries.async_update_entry(
config_entry, config_entry_setup, options={CONF_DPI_RESTRICTIONS: False}
options={CONF_DPI_RESTRICTIONS: False},
) )
await hass.async_block_till_done() await hass.async_block_till_done()
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
@ -1285,12 +1275,10 @@ async def test_poe_port_switches(
aioclient_mock: AiohttpClientMocker, aioclient_mock: AiohttpClientMocker,
mock_websocket_message, mock_websocket_message,
mock_websocket_state, mock_websocket_state,
config_entry_setup, config_entry_setup: ConfigEntry,
device_payload, device_payload: list[dict[str, Any]],
) -> None: ) -> None:
"""Test PoE port entities work.""" """Test PoE port entities work."""
config_entry = config_entry_setup
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 0
ent_reg_entry = entity_registry.async_get("switch.mock_name_port_1_poe") ent_reg_entry = entity_registry.async_get("switch.mock_name_port_1_poe")
@ -1328,8 +1316,8 @@ async def test_poe_port_switches(
# Turn off PoE # Turn off PoE
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.put( aioclient_mock.put(
f"https://{config_entry.data[CONF_HOST]}:1234" f"https://{config_entry_setup.data[CONF_HOST]}:1234"
f"/api/s/{config_entry.data[CONF_SITE_ID]}/rest/device/mock-id", f"/api/s/{config_entry_setup.data[CONF_SITE_ID]}/rest/device/mock-id",
) )
await hass.services.async_call( await hass.services.async_call(
@ -1398,12 +1386,10 @@ async def test_wlan_switches(
aioclient_mock: AiohttpClientMocker, aioclient_mock: AiohttpClientMocker,
mock_websocket_message, mock_websocket_message,
mock_websocket_state, mock_websocket_state,
config_entry_setup, config_entry_setup: ConfigEntry,
wlan_payload, wlan_payload: list[dict[str, Any]],
) -> None: ) -> None:
"""Test control of UniFi WLAN availability.""" """Test control of UniFi WLAN availability."""
config_entry = config_entry_setup
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
ent_reg_entry = entity_registry.async_get("switch.ssid_1") ent_reg_entry = entity_registry.async_get("switch.ssid_1")
@ -1426,8 +1412,8 @@ async def test_wlan_switches(
# Disable WLAN # Disable WLAN
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.put( aioclient_mock.put(
f"https://{config_entry.data[CONF_HOST]}:1234" f"https://{config_entry_setup.data[CONF_HOST]}:1234"
f"/api/s/{config_entry.data[CONF_SITE_ID]}/rest/wlanconf/{wlan['_id']}", f"/api/s/{config_entry_setup.data[CONF_SITE_ID]}/rest/wlanconf/{wlan['_id']}",
) )
await hass.services.async_call( await hass.services.async_call(
@ -1485,11 +1471,10 @@ async def test_port_forwarding_switches(
aioclient_mock: AiohttpClientMocker, aioclient_mock: AiohttpClientMocker,
mock_websocket_message, mock_websocket_message,
mock_websocket_state, mock_websocket_state,
config_entry_setup, config_entry_setup: ConfigEntry,
port_forward_payload, port_forward_payload: list[dict[str, Any]],
) -> None: ) -> None:
"""Test control of UniFi port forwarding.""" """Test control of UniFi port forwarding."""
config_entry = config_entry_setup
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1 assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
ent_reg_entry = entity_registry.async_get("switch.unifi_network_plex") ent_reg_entry = entity_registry.async_get("switch.unifi_network_plex")
@ -1512,8 +1497,8 @@ async def test_port_forwarding_switches(
# Disable port forward # Disable port forward
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.put( aioclient_mock.put(
f"https://{config_entry.data[CONF_HOST]}:1234" f"https://{config_entry_setup.data[CONF_HOST]}:1234"
f"/api/s/{config_entry.data[CONF_SITE_ID]}/rest/portforward/{data['_id']}", f"/api/s/{config_entry_setup.data[CONF_SITE_ID]}/rest/portforward/{data['_id']}",
) )
await hass.services.async_call( await hass.services.async_call(
@ -1594,8 +1579,8 @@ async def test_port_forwarding_switches(
async def test_updating_unique_id( async def test_updating_unique_id(
hass: HomeAssistant, hass: HomeAssistant,
entity_registry: er.EntityRegistry, entity_registry: er.EntityRegistry,
config_entry_factory, config_entry_factory: Callable[[], ConfigEntry],
config_entry, config_entry: ConfigEntry,
device_payload, device_payload,
) -> None: ) -> None:
"""Verify outlet control and poe control unique ID update works.""" """Verify outlet control and poe control unique ID update works."""

View File

@ -16,6 +16,7 @@ from homeassistant.components.update import (
UpdateDeviceClass, UpdateDeviceClass,
UpdateEntityFeature, UpdateEntityFeature,
) )
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ( from homeassistant.const import (
ATTR_DEVICE_CLASS, ATTR_DEVICE_CLASS,
ATTR_ENTITY_ID, ATTR_ENTITY_ID,
@ -138,18 +139,18 @@ async def test_not_admin(hass: HomeAssistant) -> None:
@pytest.mark.parametrize("device_payload", [[DEVICE_1]]) @pytest.mark.parametrize("device_payload", [[DEVICE_1]])
async def test_install( async def test_install(
hass: HomeAssistant, aioclient_mock: AiohttpClientMocker, config_entry_setup hass: HomeAssistant,
aioclient_mock: AiohttpClientMocker,
config_entry_setup: ConfigEntry,
) -> None: ) -> None:
"""Test the device update install call.""" """Test the device update install call."""
config_entry = config_entry_setup
assert len(hass.states.async_entity_ids(UPDATE_DOMAIN)) == 1 assert len(hass.states.async_entity_ids(UPDATE_DOMAIN)) == 1
device_state = hass.states.get("update.device_1") device_state = hass.states.get("update.device_1")
assert device_state.state == STATE_ON assert device_state.state == STATE_ON
url = ( url = (
f"https://{config_entry.data[CONF_HOST]}:1234" f"https://{config_entry_setup.data[CONF_HOST]}:1234"
f"/api/s/{config_entry.data[CONF_SITE_ID]}/cmd/devmgr" f"/api/s/{config_entry_setup.data[CONF_SITE_ID]}/cmd/devmgr"
) )
aioclient_mock.clear_requests() aioclient_mock.clear_requests()
aioclient_mock.post(url) aioclient_mock.post(url)