From 0444dd71a645f29e7285e9c0f1a064ea2aa826a7 Mon Sep 17 00:00:00 2001 From: Robert Svensson Date: Sun, 23 Oct 2022 20:28:45 +0200 Subject: [PATCH] Refactor UniFi outlet switches (#80738) * Rewrite UniFi outlet switches * Bump aiounifi to v41 * Remove devices from items_added input --- homeassistant/components/unifi/manifest.json | 2 +- homeassistant/components/unifi/switch.py | 127 ++++++++++--------- requirements_all.txt | 2 +- requirements_test_all.txt | 2 +- tests/components/unifi/test_switch.py | 80 +++++++----- 5 files changed, 117 insertions(+), 96 deletions(-) diff --git a/homeassistant/components/unifi/manifest.json b/homeassistant/components/unifi/manifest.json index ad5178c2d29c3b..5b96560f8c5bf3 100644 --- a/homeassistant/components/unifi/manifest.json +++ b/homeassistant/components/unifi/manifest.json @@ -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": [ diff --git a/homeassistant/components/unifi/switch.py b/homeassistant/components/unifi/switch.py index 68fe84b6b605bc..c88af43cc78149 100644 --- a/homeassistant/components/unifi/switch.py +++ b/homeassistant/components/unifi/switch.py @@ -20,7 +20,6 @@ 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 @@ def items_added( 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 @@ def device_info(self) -> DeviceInfo: ) -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 @@ def __init__(self, obj_id: str, controller) -> None: 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: diff --git a/requirements_all.txt b/requirements_all.txt index 0834663953bf38..e61e412f036834 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -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 diff --git a/requirements_test_all.txt b/requirements_test_all.txt index a451e4d10549ea..a0a95796a5cf84 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -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 diff --git a/tests/components/unifi/test_switch.py b/tests/components/unifi/test_switch.py index e0bac2c3eb3de6..12ef6f9b965aee 100644 --- a/tests/components/unifi/test_switch.py +++ b/tests/components/unifi/test_switch.py @@ -128,6 +128,7 @@ ] 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 @@ "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 - - # State change - - outlet_up1 = deepcopy(OUTLET_UP1) - outlet_up1["outlet_table"][0]["relay_state"] = True + # 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 - 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)