Refactor UniFi outlet switches (#80738)

* Rewrite UniFi outlet switches

* Bump aiounifi to v41

* Remove devices from items_added input
This commit is contained in:
Robert Svensson 2022-10-23 20:28:45 +02:00 committed by GitHub
parent 16d3cc905f
commit 0444dd71a6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 117 additions and 96 deletions

View File

@ -3,7 +3,7 @@
"name": "UniFi Network",
"config_flow": true,
"documentation": "https://www.home-assistant.io/integrations/unifi",
"requirements": ["aiounifi==40"],
"requirements": ["aiounifi==41"],
"codeowners": ["@Kane610"],
"quality_scale": "platinum",
"ssdp": [

View File

@ -20,7 +20,6 @@ from aiounifi.models.event import EventKey
from homeassistant.components.switch import DOMAIN, SwitchDeviceClass, SwitchEntity
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import ATTR_NAME
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers import entity_registry as er
from homeassistant.helpers.device_registry import (
@ -89,12 +88,9 @@ async def async_setup_entry(
@callback
def items_added(
clients: set = controller.api.clients,
devices: set = controller.api.devices,
dpi_groups: set = controller.api.dpi_groups,
) -> None:
"""Update the values of the controller."""
add_outlet_entities(controller, async_add_entities, devices)
if controller.option_block_clients:
add_block_entities(controller, async_add_entities, clients)
@ -112,6 +108,18 @@ async def async_setup_entry(
items_added()
known_poe_clients.clear()
@callback
def async_add_outlet_switch(_: ItemEvent, obj_id: str) -> None:
"""Add power outlet switch from UniFi controller."""
if not controller.api.outlets[obj_id].has_relay:
return
async_add_entities([UnifiOutletSwitch(obj_id, controller)])
controller.api.ports.subscribe(async_add_outlet_switch, ItemEvent.ADDED)
for index in controller.api.outlets:
async_add_outlet_switch(ItemEvent.ADDED, index)
@callback
def async_add_poe_switch(_: ItemEvent, obj_id: str) -> None:
"""Add port PoE switch from UniFi controller."""
@ -207,25 +215,6 @@ def add_dpi_entities(controller, async_add_entities, dpi_groups):
async_add_entities(switches)
@callback
def add_outlet_entities(controller, async_add_entities, devices):
"""Add new switch entities from the controller."""
switches = []
for mac in devices:
if (
mac in controller.entities[DOMAIN][OUTLET_SWITCH]
or not (device := controller.api.devices[mac]).outlet_table
):
continue
for outlet in device.outlets.values():
if outlet.has_relay:
switches.append(UniFiOutletSwitch(device, controller, outlet.index))
async_add_entities(switches)
class UniFiPOEClientSwitch(UniFiClient, SwitchEntity, RestoreEntity):
"""Representation of a client that uses POE."""
@ -506,64 +495,77 @@ class UniFiDPIRestrictionSwitch(UniFiBase, SwitchEntity):
)
class UniFiOutletSwitch(UniFiBase, SwitchEntity):
class UnifiOutletSwitch(SwitchEntity):
"""Representation of a outlet relay."""
DOMAIN = DOMAIN
TYPE = OUTLET_SWITCH
_attr_device_class = SwitchDeviceClass.OUTLET
_attr_has_entity_name = True
_attr_should_poll = False
def __init__(self, device, controller, index):
"""Set up outlet switch."""
super().__init__(device, controller)
def __init__(self, obj_id: str, controller) -> None:
"""Set up UniFi Network entity base."""
self._device_mac, index = obj_id.split("_", 1)
self._index = int(index)
self._obj_id = obj_id
self.controller = controller
self._outlet_index = index
outlet = self.controller.api.outlets[self._obj_id]
self._attr_name = outlet.name
self._attr_is_on = outlet.relay_state
self._attr_unique_id = f"{self._device_mac}-outlet-{index}"
self._attr_name = f"{device.name or device.model} {device.outlets[index].name}"
self._attr_unique_id = f"{device.mac}-outlet-{index}"
device = self.controller.api.devices[self._device_mac]
self._attr_available = controller.available and not device.disabled
self._attr_device_info = DeviceInfo(
connections={(CONNECTION_NETWORK_MAC, device.mac)},
manufacturer=ATTR_MANUFACTURER,
model=device.model,
name=device.name or None,
sw_version=device.version,
hw_version=device.board_revision,
)
@property
def is_on(self):
"""Return true if outlet is active."""
return self._item.outlets[self._outlet_index].relay_state
async def async_added_to_hass(self) -> None:
"""Entity created."""
self.async_on_remove(
self.controller.api.outlets.subscribe(self.async_signalling_callback)
)
self.async_on_remove(
async_dispatcher_connect(
self.hass,
self.controller.signal_reachable,
self.async_signal_reachable_callback,
)
)
@property
def available(self) -> bool:
"""Return if switch is available."""
return not self._item.disabled and self.controller.available
@callback
def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None:
"""Object has new event."""
device = self.controller.api.devices[self._device_mac]
outlet = self.controller.api.outlets[self._obj_id]
self._attr_available = self.controller.available and not device.disabled
self._attr_is_on = outlet.relay_state
self.async_write_ha_state()
@callback
def async_signal_reachable_callback(self) -> None:
"""Call when controller connection state change."""
self.async_signalling_callback(ItemEvent.ADDED, self._obj_id)
async def async_turn_on(self, **kwargs: Any) -> None:
"""Enable outlet relay."""
device = self.controller.api.devices[self._device_mac]
await self.controller.api.request(
DeviceSetOutletRelayRequest.create(self._item, self._outlet_index, True)
DeviceSetOutletRelayRequest.create(device, self._index, True)
)
async def async_turn_off(self, **kwargs: Any) -> None:
"""Disable outlet relay."""
device = self.controller.api.devices[self._device_mac]
await self.controller.api.request(
DeviceSetOutletRelayRequest.create(self._item, self._outlet_index, False)
DeviceSetOutletRelayRequest.create(device, self._index, False)
)
@property
def device_info(self) -> DeviceInfo:
"""Return a device description for device registry."""
info = DeviceInfo(
connections={(CONNECTION_NETWORK_MAC, self._item.mac)},
manufacturer=ATTR_MANUFACTURER,
model=self._item.model,
sw_version=self._item.version,
hw_version=self._item.board_revision,
)
if self._item.name:
info[ATTR_NAME] = self._item.name
return info
async def options_updated(self) -> None:
"""Config entry options are updated, no options to act on."""
class UnifiPoePortSwitch(SwitchEntity):
"""Representation of a Power-over-Ethernet source port on an UniFi device."""
@ -594,6 +596,7 @@ class UnifiPoePortSwitch(SwitchEntity):
model=device.model,
name=device.name or None,
sw_version=device.version,
hw_version=device.board_revision,
)
async def async_added_to_hass(self) -> None:

View File

@ -276,7 +276,7 @@ aiosyncthing==0.5.1
aiotractive==0.5.4
# homeassistant.components.unifi
aiounifi==40
aiounifi==41
# homeassistant.components.vlc_telnet
aiovlc==0.1.0

View File

@ -251,7 +251,7 @@ aiosyncthing==0.5.1
aiotractive==0.5.4
# homeassistant.components.unifi
aiounifi==40
aiounifi==41
# homeassistant.components.vlc_telnet
aiovlc==0.1.0

View File

@ -128,6 +128,7 @@ POE_SWITCH_CLIENTS = [
]
DEVICE_1 = {
"board_rev": 2,
"device_id": "mock-id",
"ip": "10.0.1.1",
"mac": "00:00:00:00:01:01",
@ -413,9 +414,16 @@ OUTLET_UP1 = {
"index": 1,
"has_relay": True,
"has_metering": False,
"relay_state": True,
"name": "Outlet 1",
},
{
"index": 2,
"has_relay": False,
"has_metering": False,
"relay_state": False,
"name": "Outlet 1",
}
},
],
"element_ap_serial": "44:d9:e7:90:f4:24",
"connected_at": 1641678609,
@ -910,34 +918,27 @@ async def test_dpi_switches_add_second_app(hass, aioclient_mock, mock_unifi_webs
async def test_outlet_switches(hass, aioclient_mock, mock_unifi_websocket):
"""Test the update_items function with some clients."""
"""Test the outlet entities."""
config_entry = await setup_unifi_integration(
hass,
aioclient_mock,
options={CONF_TRACK_DEVICES: False},
devices_response=[OUTLET_UP1],
hass, aioclient_mock, devices_response=[OUTLET_UP1]
)
controller = hass.data[UNIFI_DOMAIN][config_entry.entry_id]
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
outlet = hass.states.get("switch.plug_outlet_1")
assert outlet is not None
assert outlet.state == STATE_OFF
# Validate state object
switch_1 = hass.states.get("switch.plug_outlet_1")
assert switch_1 is not None
assert switch_1.state == STATE_ON
assert switch_1.attributes.get(ATTR_DEVICE_CLASS) == SwitchDeviceClass.OUTLET
# State change
outlet_up1 = deepcopy(OUTLET_UP1)
outlet_up1["outlet_table"][0]["relay_state"] = True
mock_unifi_websocket(message=MessageKey.DEVICE, data=outlet_up1)
# Update state object
device_1 = deepcopy(OUTLET_UP1)
device_1["outlet_table"][0]["relay_state"] = False
mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
await hass.async_block_till_done()
assert hass.states.get("switch.plug_outlet_1").state == STATE_OFF
outlet = hass.states.get("switch.plug_outlet_1")
assert outlet.state == STATE_ON
# Turn on and off outlet
# Turn off outlet
aioclient_mock.clear_requests()
aioclient_mock.put(
f"https://{controller.host}:1234/api/s/{controller.site}/rest/device/600c8356942a6ade50707b56",
@ -945,33 +946,50 @@ async def test_outlet_switches(hass, aioclient_mock, mock_unifi_websocket):
await hass.services.async_call(
SWITCH_DOMAIN,
SERVICE_TURN_ON,
SERVICE_TURN_OFF,
{ATTR_ENTITY_ID: "switch.plug_outlet_1"},
blocking=True,
)
assert aioclient_mock.call_count == 1
assert aioclient_mock.mock_calls[0][2] == {
"outlet_overrides": [{"index": 1, "name": "Outlet 1", "relay_state": True}]
"outlet_overrides": [{"index": 1, "name": "Outlet 1", "relay_state": False}]
}
# Turn on outlet
await hass.services.async_call(
SWITCH_DOMAIN,
SERVICE_TURN_OFF,
SERVICE_TURN_ON,
{ATTR_ENTITY_ID: "switch.plug_outlet_1"},
blocking=True,
)
assert aioclient_mock.call_count == 2
assert aioclient_mock.mock_calls[1][2] == {
"outlet_overrides": [{"index": 1, "name": "Outlet 1", "relay_state": False}]
"outlet_overrides": [{"index": 1, "name": "Outlet 1", "relay_state": True}]
}
# Changes to config entry options shouldn't affect outlets
hass.config_entries.async_update_entry(
config_entry,
options={CONF_BLOCK_CLIENT: []},
)
# Availability signalling
# Controller disconnects
mock_unifi_websocket(state=WebsocketState.DISCONNECTED)
await hass.async_block_till_done()
assert len(hass.states.async_entity_ids(SWITCH_DOMAIN)) == 1
assert hass.states.get("switch.plug_outlet_1").state == STATE_UNAVAILABLE
# Controller reconnects
mock_unifi_websocket(state=WebsocketState.RUNNING)
await hass.async_block_till_done()
assert hass.states.get("switch.plug_outlet_1").state == STATE_OFF
# Device gets disabled
device_1["disabled"] = True
mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
await hass.async_block_till_done()
assert hass.states.get("switch.plug_outlet_1").state == STATE_UNAVAILABLE
# Device gets re-enabled
device_1["disabled"] = False
mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
await hass.async_block_till_done()
assert hass.states.get("switch.plug_outlet_1").state == STATE_OFF
# Unload config entry
await hass.config_entries.async_unload(config_entry.entry_id)