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

Add device trackers to tplink_omada (#115601)

Browse source 
* Add device trackers to tplink_omada

* tplink_omada - Remove trackers and options flow

* Addressed code review feedback

* Run linter

* Use entity registry fixture
This commit is contained in:
MarkGodwin 2024-06-17 06:36:35 +01:00 • committed by GitHub
parent bd37ce6e9a
commit f09063d706
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 638 additions and 70 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
homeassistant/components/tplink_omada/__init__.py
View file

@ -19,7 +19,12 @@ from .config_flow import CONF_SITE, create_omada_client
from .const import DOMAIN
from .controller import OmadaSiteController
PLATFORMS: list[Platform] = [Platform.BINARY_SENSOR, Platform.SWITCH, Platform.UPDATE]
PLATFORMS: list[Platform] = [
Platform.BINARY_SENSOR,
Platform.DEVICE_TRACKER,
Platform.SWITCH,
Platform.UPDATE,
]
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
@ -50,10 +55,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
gateway_coordinator = await controller.get_gateway_coordinator()
if gateway_coordinator:
await gateway_coordinator.async_config_entry_first_refresh()
await controller.get_clients_coordinator().async_config_entry_first_refresh()
hass.data[DOMAIN][entry.entry_id] = controller
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
return True

65
homeassistant/components/tplink_omada/controller.py
View file

@ -1,58 +1,15 @@
"""Controller for sharing Omada API coordinators between platforms."""
from tplink_omada_client import OmadaSiteClient
from tplink_omada_client.devices import (
OmadaGateway,
OmadaSwitch,
OmadaSwitchPortDetails,
)
from tplink_omada_client.devices import OmadaSwitch
from homeassistant.core import HomeAssistant
from .coordinator import OmadaCoordinator
POLL_SWITCH_PORT = 300
POLL_GATEWAY = 300
class OmadaSwitchPortCoordinator(OmadaCoordinator[OmadaSwitchPortDetails]): # pylint: disable=hass-enforce-coordinator-module
"""Coordinator for getting details about ports on a switch."""
def __init__(
self,
hass: HomeAssistant,
omada_client: OmadaSiteClient,
network_switch: OmadaSwitch,
) -> None:
"""Initialize my coordinator."""
super().__init__(
hass, omada_client, f"{network_switch.name} Ports", POLL_SWITCH_PORT
)
self._network_switch = network_switch
async def poll_update(self) -> dict[str, OmadaSwitchPortDetails]:
"""Poll a switch's current state."""
ports = await self.omada_client.get_switch_ports(self._network_switch)
return {p.port_id: p for p in ports}
class OmadaGatewayCoordinator(OmadaCoordinator[OmadaGateway]): # pylint: disable=hass-enforce-coordinator-module
"""Coordinator for getting details about the site's gateway."""
def __init__(
self,
hass: HomeAssistant,
omada_client: OmadaSiteClient,
mac: str,
) -> None:
"""Initialize my coordinator."""
super().__init__(hass, omada_client, "Gateway", POLL_GATEWAY)
self.mac = mac
async def poll_update(self) -> dict[str, OmadaGateway]:
"""Poll a the gateway's current state."""
gateway = await self.omada_client.get_gateway(self.mac)
return {self.mac: gateway}
from .coordinator import (
OmadaClientsCoordinator,
OmadaGatewayCoordinator,
OmadaSwitchPortCoordinator,
)
class OmadaSiteController:
@ -60,6 +17,7 @@ class OmadaSiteController:
_gateway_coordinator: OmadaGatewayCoordinator | None = None
_initialized_gateway_coordinator = False
_clients_coordinator: OmadaClientsCoordinator | None = None
def __init__(self, hass: HomeAssistant, omada_client: OmadaSiteClient) -> None:
"""Create the controller."""
@ -98,3 +56,12 @@ class OmadaSiteController:
)
return self._gateway_coordinator
def get_clients_coordinator(self) -> OmadaClientsCoordinator:
"""Get coordinator for site's clients."""
if not self._clients_coordinator:
self._clients_coordinator = OmadaClientsCoordinator(
self._hass, self._omada_client
)
return self._clients_coordinator

64
homeassistant/components/tplink_omada/coordinator.py
View file

@ -4,7 +4,9 @@ import asyncio
from datetime import timedelta
import logging
from tplink_omada_client import OmadaSiteClient
from tplink_omada_client import OmadaSiteClient, OmadaSwitchPortDetails
from tplink_omada_client.clients import OmadaWirelessClient
from tplink_omada_client.devices import OmadaGateway, OmadaSwitch
from tplink_omada_client.exceptions import OmadaClientException
from homeassistant.core import HomeAssistant
@ -12,6 +14,10 @@ from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, Upda
_LOGGER = logging.getLogger(__name__)
POLL_SWITCH_PORT = 300
POLL_GATEWAY = 300
POLL_CLIENTS = 300
class OmadaCoordinator[_T](DataUpdateCoordinator[dict[str, _T]]):
"""Coordinator for synchronizing bulk Omada data."""
@ -43,3 +49,59 @@ class OmadaCoordinator[_T](DataUpdateCoordinator[dict[str, _T]]):
async def poll_update(self) -> dict[str, _T]:
"""Poll the current data from the controller."""
raise NotImplementedError("Update method not implemented")
class OmadaSwitchPortCoordinator(OmadaCoordinator[OmadaSwitchPortDetails]):
"""Coordinator for getting details about ports on a switch."""
def __init__(
self,
hass: HomeAssistant,
omada_client: OmadaSiteClient,
network_switch: OmadaSwitch,
) -> None:
"""Initialize my coordinator."""
super().__init__(
hass, omada_client, f"{network_switch.name} Ports", POLL_SWITCH_PORT
)
self._network_switch = network_switch
async def poll_update(self) -> dict[str, OmadaSwitchPortDetails]:
"""Poll a switch's current state."""
ports = await self.omada_client.get_switch_ports(self._network_switch)
return {p.port_id: p for p in ports}
class OmadaGatewayCoordinator(OmadaCoordinator[OmadaGateway]):
"""Coordinator for getting details about the site's gateway."""
def __init__(
self,
hass: HomeAssistant,
omada_client: OmadaSiteClient,
mac: str,
) -> None:
"""Initialize my coordinator."""
super().__init__(hass, omada_client, "Gateway", POLL_GATEWAY)
self.mac = mac
async def poll_update(self) -> dict[str, OmadaGateway]:
"""Poll a the gateway's current state."""
gateway = await self.omada_client.get_gateway(self.mac)
return {self.mac: gateway}
class OmadaClientsCoordinator(OmadaCoordinator[OmadaWirelessClient]):
"""Coordinator for getting details about the site's connected clients."""
def __init__(self, hass: HomeAssistant, omada_client: OmadaSiteClient) -> None:
"""Initialize my coordinator."""
super().__init__(hass, omada_client, "ClientsList", POLL_CLIENTS)
async def poll_update(self) -> dict[str, OmadaWirelessClient]:
"""Poll the site's current active wi-fi clients."""
return {
c.mac: c
async for c in self.omada_client.get_connected_clients()
if isinstance(c, OmadaWirelessClient)
}

107
homeassistant/components/tplink_omada/device_tracker.py Normal file
View file

@ -0,0 +1,107 @@
"""Connected Wi-Fi device scanners for TP-Link Omada access points."""
import logging
from tplink_omada_client.clients import OmadaWirelessClient
from homeassistant.components.device_tracker import ScannerEntity, SourceType
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.update_coordinator import CoordinatorEntity
from .config_flow import CONF_SITE
from .const import DOMAIN
from .controller import OmadaClientsCoordinator, OmadaSiteController
_LOGGER = logging.getLogger(__name__)
async def async_setup_entry(
hass: HomeAssistant,
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up device trackers and scanners."""
controller: OmadaSiteController = hass.data[DOMAIN][config_entry.entry_id]
clients_coordinator = controller.get_clients_coordinator()
site_id = config_entry.data[CONF_SITE]
# Add all known WiFi devices as potentially tracked devices. They will only be
# tracked if the user enables the entity.
async_add_entities(
[
OmadaClientScannerEntity(
site_id, client.mac, client.name, clients_coordinator
)
async for client in controller.omada_client.get_known_clients()
if isinstance(client, OmadaWirelessClient)
]
)
class OmadaClientScannerEntity(
CoordinatorEntity[OmadaClientsCoordinator], ScannerEntity
):
"""Entity for a client connected to the Omada network."""
_client_details: OmadaWirelessClient | None = None
def __init__(
self,
site_id: str,
client_id: str,
display_name: str,
coordinator: OmadaClientsCoordinator,
) -> None:
"""Initialize the scanner."""
super().__init__(coordinator)
self._site_id = site_id
self._client_id = client_id
self._attr_name = display_name
@property
def source_type(self) -> SourceType:
"""Return the source type of the device."""
return SourceType.ROUTER
def _do_update(self) -> None:
self._client_details = self.coordinator.data.get(self._client_id)
async def async_added_to_hass(self) -> None:
"""When entity is added to hass."""
await super().async_added_to_hass()
self._do_update()
@callback
def _handle_coordinator_update(self) -> None:
"""Handle updated data from the coordinator."""
self._do_update()
self.async_write_ha_state()
@property
def ip_address(self) -> str | None:
"""Return the primary ip address of the device."""
return self._client_details.ip if self._client_details else None
@property
def mac_address(self) -> str | None:
"""Return the mac address of the device."""
return self._client_id
@property
def hostname(self) -> str | None:
"""Return hostname of the device."""
return self._client_details.host_name if self._client_details else None
@property
def is_connected(self) -> bool:
"""Return true if the device is connected to the network."""
return self._client_details.is_active if self._client_details else False
@property
def unique_id(self) -> str | None:
"""Return the unique id of the device."""
return f"scanner_{self._site_id}_{self._client_id}"

98
tests/components/tplink_omada/conftest.py
View file

@ -1,9 +1,16 @@
"""Test fixtures for TP-Link Omada integration."""
from collections.abc import AsyncIterable
import json
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from tplink_omada_client.clients import (
OmadaConnectedClient,
OmadaNetworkClient,
OmadaWiredClient,
OmadaWirelessClient,
)
from tplink_omada_client.devices import (
OmadaGateway,
OmadaListDevice,
@ -49,29 +56,82 @@ def mock_setup_entry() -> Generator[AsyncMock]:
@pytest.fixture
def mock_omada_site_client() -> Generator[AsyncMock]:
"""Mock Omada site client."""
site_client = AsyncMock()
site_client = MagicMock()
gateway_data = json.loads(load_fixture("gateway-TL-ER7212PC.json", DOMAIN))
gateway = OmadaGateway(gateway_data)
site_client.get_gateway.return_value = gateway
site_client.get_gateway = AsyncMock(return_value=gateway)
switch1_data = json.loads(load_fixture("switch-TL-SG3210XHP-M2.json", DOMAIN))
switch1 = OmadaSwitch(switch1_data)
site_client.get_switches.return_value = [switch1]
site_client.get_switches = AsyncMock(return_value=[switch1])
devices_data = json.loads(load_fixture("devices.json", DOMAIN))
devices = [OmadaListDevice(d) for d in devices_data]
site_client.get_devices.return_value = devices
site_client.get_devices = AsyncMock(return_value=devices)
switch1_ports_data = json.loads(
load_fixture("switch-ports-TL-SG3210XHP-M2.json", DOMAIN)
)
switch1_ports = [OmadaSwitchPortDetails(p) for p in switch1_ports_data]
site_client.get_switch_ports.return_value = switch1_ports
site_client.get_switch_ports = AsyncMock(return_value=switch1_ports)
async def async_empty() -> AsyncIterable:
for c in []:
yield c
site_client.get_known_clients.return_value = async_empty()
site_client.get_connected_clients.return_value = async_empty()
return site_client
@pytest.fixture
def mock_omada_clients_only_site_client() -> Generator[AsyncMock]:
"""Mock Omada site client containing only client connection data."""
site_client = MagicMock()
site_client.get_switches = AsyncMock(return_value=[])
site_client.get_devices = AsyncMock(return_value=[])
site_client.get_switch_ports = AsyncMock(return_value=[])
site_client.get_client = AsyncMock(side_effect=_get_mock_client)
site_client.get_known_clients.side_effect = _get_mock_known_clients
site_client.get_connected_clients.side_effect = _get_mock_connected_clients
return site_client
async def _get_mock_known_clients() -> AsyncIterable[OmadaNetworkClient]:
"""Mock known clients of the Omada network."""
known_clients_data = json.loads(load_fixture("known-clients.json", DOMAIN))
for c in known_clients_data:
if c["wireless"]:
yield OmadaWirelessClient(c)
else:
yield OmadaWiredClient(c)
async def _get_mock_connected_clients() -> AsyncIterable[OmadaConnectedClient]:
"""Mock connected clients of the Omada network."""
connected_clients_data = json.loads(load_fixture("connected-clients.json", DOMAIN))
for c in connected_clients_data:
if c["wireless"]:
yield OmadaWirelessClient(c)
else:
yield OmadaWiredClient(c)
def _get_mock_client(mac: str) -> OmadaNetworkClient:
"""Mock an Omada client."""
connected_clients_data = json.loads(load_fixture("connected-clients.json", DOMAIN))
for c in connected_clients_data:
if c["mac"] == mac:
if c["wireless"]:
return OmadaWirelessClient(c)
return OmadaWiredClient(c)
@pytest.fixture
def mock_omada_client(mock_omada_site_client: AsyncMock) -> Generator[MagicMock]:
"""Mock Omada client."""
@ -85,13 +145,39 @@ def mock_omada_client(mock_omada_site_client: AsyncMock) -> Generator[MagicMock]
yield client
@pytest.fixture
def mock_omada_clients_only_client(
mock_omada_clients_only_site_client: AsyncMock,
) -> Generator[MagicMock]:
"""Mock Omada client."""
with patch(
"homeassistant.components.tplink_omada.create_omada_client",
autospec=True,
) as client_mock:
client = client_mock.return_value
client.get_site_client.return_value = mock_omada_clients_only_site_client
yield client
@pytest.fixture
async def init_integration(
hass: HomeAssistant,
mock_config_entry: MockConfigEntry,
mock_omada_client: MagicMock,
) -> MockConfigEntry:
"""Set up the TP-Link Omada integration for testing."""
mock_config_entry = MockConfigEntry(
title="Test Omada Controller",
domain=DOMAIN,
data={
CONF_HOST: "127.0.0.1",
CONF_PASSWORD: "mocked-password",
CONF_USERNAME: "mocked-user",
CONF_VERIFY_SSL: False,
CONF_SITE: "Default",
},
unique_id="12345",
)
mock_config_entry.add_to_hass(hass)
await hass.config_entries.async_setup(mock_config_entry.entry_id)

120
tests/components/tplink_omada/fixtures/connected-clients.json Normal file
View file

@ -0,0 +1,120 @@
[
{
"mac": "16-32-50-ED-FB-15",
"name": "16-32-50-ED-FB-15",
"deviceType": "unknown",
"ip": "192.168.1.177",
"connectType": 1,
"connectDevType": "ap",
"connectedToWirelessRouter": false,
"wireless": true,
"ssid": "OFFICE_SSID",
"signalLevel": 62,
"healthScore": -1,
"signalRank": 4,
"wifiMode": 4,
"apName": "Office",
"apMac": "E8-48-B8-7E-C7-1A",
"radioId": 0,
"channel": 1,
"rxRate": 65000,
"txRate": 72000,
"powerSave": false,
"rssi": -65,
"snr": 30,
"stackableSwitch": false,
"vid": 0,
"activity": 96,
"trafficDown": 25412800785,
"trafficUp": 1636427981,
"uptime": 621441,
"lastSeen": 1713109713169,
"authStatus": 0,
"guest": false,
"active": true,
"manager": false,
"downPacket": 30179275,
"upPacket": 14288106,
"support5g2": false,
"multiLink": []
},
{
"mac": "2E-DC-E1-C4-37-D3",
"name": "Apple",
"deviceType": "unknown",
"ip": "192.168.1.192",
"connectType": 1,
"connectDevType": "ap",
"connectedToWirelessRouter": false,
"wireless": true,
"ssid": "ROAMING_SSID",
"signalLevel": 67,
"healthScore": -1,
"signalRank": 4,
"wifiMode": 5,
"apName": "Spare Room",
"apMac": "C0-C9-E3-4B-AF-0E",
"radioId": 1,
"channel": 44,
"rxRate": 7000,
"txRate": 390000,
"powerSave": false,
"rssi": -63,
"snr": 32,
"stackableSwitch": false,
"vid": 0,
"activity": 0,
"trafficDown": 3327229,
"trafficUp": 746841,
"uptime": 2091,
"lastSeen": 1713109728764,
"authStatus": 0,
"guest": false,
"active": true,
"manager": false,
"downPacket": 5128,
"upPacket": 3611,
"support5g2": false,
"multiLink": []
},
{
"mac": "2C-71-FF-ED-34-83",
"name": "Banana",
"hostName": "testhost",
"deviceType": "unknown",
"ip": "192.168.1.102",
"connectType": 1,
"connectDevType": "ap",
"connectedToWirelessRouter": false,
"wireless": true,
"ssid": "ROAMING_SSID",
"signalLevel": 57,
"healthScore": -1,
"signalRank": 3,
"wifiMode": 5,
"apName": "Living Room",
"apMac": "C0-C9-E3-4B-A7-FE",
"radioId": 1,
"channel": 36,
"rxRate": 6000,
"txRate": 390000,
"powerSave": false,
"rssi": -67,
"snr": 28,
"stackableSwitch": false,
"vid": 0,
"activity": 39,
"trafficDown": 407300090,
"trafficUp": 94910187,
"uptime": 621461,
"lastSeen": 1713109729576,
"authStatus": 0,
"guest": false,
"active": true,
"manager": false,
"downPacket": 477858,
"upPacket": 501956,
"support5g2": false,
"multiLink": []
}
]

67
tests/components/tplink_omada/fixtures/known-clients.json Normal file
View file

@ -0,0 +1,67 @@
[
{
"name": "16-32-50-ED-FB-15",
"mac": "16-32-50-ED-FB-15",
"wireless": true,
"guest": false,
"download": 259310931013,
"upload": 43957031162,
"duration": 6832173,
"lastSeen": 1712488285622,
"block": false,
"manager": false,
"lockToAp": false
},
{
"name": "Banana",
"mac": "2C-71-FF-ED-34-83",
"wireless": true,
"guest": false,
"download": 22093851790,
"upload": 6961197401,
"duration": 16192898,
"lastSeen": 1712488285767,
"block": false,
"manager": false,
"lockToAp": false
},
{
"name": "Pear",
"mac": "2C-D2-6B-BA-9C-94",
"wireless": true,
"guest": false,
"download": 0,
"upload": 0,
"duration": 23,
"lastSeen": 1713083620997,
"block": false,
"manager": false,
"lockToAp": false
},
{
"name": "Apple",
"mac": "2E-DC-E1-C4-37-D3",
"wireless": true,
"guest": false,
"download": 1366833567,
"upload": 30126947,
"duration": 60255,
"lastSeen": 1713107649827,
"block": false,
"manager": false,
"lockToAp": false
},
{
"name": "32-39-24-B1-67-23",
"mac": "32-39-24-B1-67-23",
"wireless": false,
"guest": false,
"download": 1621140542,
"upload": 433306522,
"duration": 60571,
"lastSeen": 1713107438528,
"block": false,
"manager": false,
"lockToAp": false
}
]

33
tests/components/tplink_omada/snapshots/test_device_tracker.ambr Normal file
View file

@ -0,0 +1,33 @@
# serializer version: 1
# name: test_device_scanner_created
StateSnapshot({
'attributes': ReadOnlyDict({
'friendly_name': 'Banana',
'host_name': 'testhost',
'ip': '192.168.1.102',
'mac': '2C-71-FF-ED-34-83',
'source_type': <SourceType.ROUTER: 'router'>,
}),
'context': <ANY>,
'entity_id': 'device_tracker.banana',
'last_changed': <ANY>,
'last_reported': <ANY>,
'last_updated': <ANY>,
'state': 'home',
})
# ---
# name: test_device_scanner_update_to_away_nulls_properties
StateSnapshot({
'attributes': ReadOnlyDict({
'friendly_name': 'Banana',
'mac': '2C-71-FF-ED-34-83',
'source_type': <SourceType.ROUTER: 'router'>,
}),
'context': <ANY>,
'entity_id': 'device_tracker.banana',
'last_changed': <ANY>,
'last_reported': <ANY>,
'last_updated': <ANY>,
'state': 'not_home',
})
# ---

117
tests/components/tplink_omada/test_device_tracker.py Normal file
View file

@ -0,0 +1,117 @@
"""Tests for TP-Link Omada device tracker entities."""
from collections.abc import AsyncIterable
from datetime import timedelta
from unittest.mock import MagicMock
import pytest
from syrupy.assertion import SnapshotAssertion
from tplink_omada_client.clients import OmadaConnectedClient
from homeassistant.components.tplink_omada.const import DOMAIN
from homeassistant.components.tplink_omada.coordinator import POLL_CLIENTS
from homeassistant.core import HomeAssistant
from homeassistant.helpers import entity_registry as er
from homeassistant.util.dt import utcnow
from tests.common import MockConfigEntry, async_fire_time_changed
UPDATE_INTERVAL = timedelta(seconds=10)
POLL_INTERVAL = timedelta(seconds=POLL_CLIENTS + 10)
MOCK_ENTRY_DATA = {
"host": "https://fake.omada.host",
"verify_ssl": True,
"site": "SiteId",
"username": "test-username",
"password": "test-password",
}
@pytest.fixture
async def init_integration(
hass: HomeAssistant,
mock_omada_clients_only_client: MagicMock,
) -> MockConfigEntry:
"""Set up the TP-Link Omada integration for testing."""
mock_config_entry = MockConfigEntry(
title="Test Omada Controller",
domain=DOMAIN,
data=dict(MOCK_ENTRY_DATA),
unique_id="12345",
)
mock_config_entry.add_to_hass(hass)
await hass.config_entries.async_setup(mock_config_entry.entry_id)
await hass.async_block_till_done()
return mock_config_entry
async def test_device_scanner_created(
hass: HomeAssistant,
init_integration: MockConfigEntry,
snapshot: SnapshotAssertion,
entity_registry: er.EntityRegistry,
) -> None:
"""Test gateway connected switches."""
entity_id = "device_tracker.banana"
updated_entity = entity_registry.async_update_entity(entity_id, disabled_by=None)
assert not updated_entity.disabled
async_fire_time_changed(hass, utcnow() + POLL_INTERVAL)
await hass.async_block_till_done()
entity = hass.states.get(entity_id)
assert entity is not None
assert entity == snapshot
async def test_device_scanner_update_to_away_nulls_properties(
hass: HomeAssistant,
mock_omada_clients_only_site_client: MagicMock,
init_integration: MockConfigEntry,
snapshot: SnapshotAssertion,
entity_registry: er.EntityRegistry,
) -> None:
"""Test gateway connected switches."""
entity_id = "device_tracker.banana"
updated_entity = entity_registry.async_update_entity(entity_id, disabled_by=None)
assert not updated_entity.disabled
async_fire_time_changed(hass, utcnow() + POLL_INTERVAL)
await hass.async_block_till_done()
entity = hass.states.get(entity_id)
await _setup_client_disconnect(
mock_omada_clients_only_site_client, "2C-71-FF-ED-34-83"
)
async_fire_time_changed(hass, utcnow() + (POLL_INTERVAL * 2))
await hass.async_block_till_done()
entity = hass.states.get(entity_id)
assert entity is not None
assert entity == snapshot
mock_omada_clients_only_site_client.get_connected_clients.assert_called_once()
async def _setup_client_disconnect(
mock_omada_site_client: MagicMock,
client_mac: str,
):
original_clients = [
c
async for c in mock_omada_site_client.get_connected_clients()
if c.mac != client_mac
]
async def get_filtered_clients() -> AsyncIterable[OmadaConnectedClient]:
for c in original_clients:
yield c
mock_omada_site_client.get_connected_clients.reset_mock()
mock_omada_site_client.get_connected_clients.side_effect = get_filtered_clients

28
tests/components/tplink_omada/test_switch.py
View file

@ -2,7 +2,7 @@
from datetime import timedelta
from typing import Any
from unittest.mock import MagicMock
from unittest.mock import AsyncMock, MagicMock
from syrupy.assertion import SnapshotAssertion
from tplink_omada_client import SwitchPortOverrides
@ -17,7 +17,7 @@ from tplink_omada_client.devices import (
from tplink_omada_client.exceptions import InvalidDevice
from homeassistant.components import switch
from homeassistant.components.tplink_omada.controller import POLL_GATEWAY
from homeassistant.components.tplink_omada.coordinator import POLL_GATEWAY
from homeassistant.const import ATTR_ENTITY_ID
from homeassistant.core import HomeAssistant
from homeassistant.helpers import entity_registry as er
@ -34,6 +34,7 @@ async def test_poe_switches(
mock_omada_site_client: MagicMock,
init_integration: MockConfigEntry,
snapshot: SnapshotAssertion,
entity_registry: er.EntityRegistry,
) -> None:
"""Test PoE switch."""
poe_switch_mac = "54-AF-97-00-00-01"
@ -44,6 +45,7 @@ async def test_poe_switches(
poe_switch_mac,
1,
snapshot,
entity_registry,
)
await _test_poe_switch(
@ -53,6 +55,7 @@ async def test_poe_switches(
poe_switch_mac,
2,
snapshot,
entity_registry,
)
@ -84,10 +87,11 @@ async def test_gateway_connect_ipv4_switch(
port_status = test_gateway.port_status[3]
assert port_status.port_number == 4
mock_omada_site_client.set_gateway_wan_port_connect_state.reset_mock()
mock_omada_site_client.set_gateway_wan_port_connect_state.return_value = (
_get_updated_gateway_port_status(
mock_omada_site_client, test_gateway, 3, "internetState", 0
mock_omada_site_client.set_gateway_wan_port_connect_state = AsyncMock(
return_value=(
_get_updated_gateway_port_status(
mock_omada_site_client, test_gateway, 3, "internetState", 0
)
)
)
await call_service(hass, "turn_off", entity_id)
@ -136,8 +140,8 @@ async def test_gateway_port_poe_switch(
port_config = test_gateway.port_configs[4]
assert port_config.port_number == 5
mock_omada_site_client.set_gateway_port_settings.return_value = (
OmadaGatewayPortConfig(port_config.raw_data, poe_enabled=False)
mock_omada_site_client.set_gateway_port_settings = AsyncMock(
return_value=(OmadaGatewayPortConfig(port_config.raw_data, poe_enabled=False))
)
await call_service(hass, "turn_off", entity_id)
_assert_gateway_poe_set(mock_omada_site_client, test_gateway, False)
@ -239,9 +243,8 @@ async def _test_poe_switch(
network_switch_mac: str,
port_num: int,
snapshot: SnapshotAssertion,
entity_registry: er.EntityRegistry,
) -> None:
entity_registry = er.async_get(hass)
def assert_update_switch_port(
device: OmadaSwitch,
switch_port_details: OmadaSwitchPortDetails,
@ -260,9 +263,8 @@ async def _test_poe_switch(
entry = entity_registry.async_get(entity_id)
assert entry == snapshot
mock_omada_site_client.update_switch_port.reset_mock()
mock_omada_site_client.update_switch_port.return_value = await _update_port_details(
mock_omada_site_client, port_num, False
mock_omada_site_client.update_switch_port = AsyncMock(
return_value=await _update_port_details(mock_omada_site_client, port_num, False)
)
await call_service(hass, "turn_off", entity_id)