Refactor UniFi outlet switches (#80738)

* Rewrite UniFi outlet switches

* Bump aiounifi to v41

* Remove devices from items_added input

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": [

homeassistant/components/unifi/switch.py

@@ -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):
+    def __init__(self, obj_id: str, controller) -> None:
-        """Set up outlet switch."""
+        """Set up UniFi Network entity base."""
-        super().__init__(device, controller)
+        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}"
+        device = self.controller.api.devices[self._device_mac]
-        self._attr_unique_id = f"{device.mac}-outlet-{index}"
+        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
+    async def async_added_to_hass(self) -> None:
-    def is_on(self):
+        """Entity created."""
-        """Return true if outlet is active."""
+        self.async_on_remove(
-        return self._item.outlets[self._outlet_index].relay_state
+            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
+    @callback
-    def available(self) -> bool:
+    def async_signalling_callback(self, event: ItemEvent, obj_id: str) -> None:
-        """Return if switch is available."""
+        """Object has new event."""
-        return not self._item.disabled and self.controller.available
+        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:

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

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

tests/components/unifi/test_switch.py

@@ -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,
+        hass, aioclient_mock, devices_response=[OUTLET_UP1]
-        aioclient_mock,
-        options={CONF_TRACK_DEVICES: False},
-        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")
+    # Validate state object
-    assert outlet is not None
+    switch_1 = hass.states.get("switch.plug_outlet_1")
-    assert outlet.state == STATE_OFF
+    assert switch_1 is not None
+    assert switch_1.state == STATE_ON
+    assert switch_1.attributes.get(ATTR_DEVICE_CLASS) == SwitchDeviceClass.OUTLET
 
-    # State change
+    # Update state object
-
+    device_1 = deepcopy(OUTLET_UP1)
-    outlet_up1 = deepcopy(OUTLET_UP1)
+    device_1["outlet_table"][0]["relay_state"] = False
-    outlet_up1["outlet_table"][0]["relay_state"] = True
+    mock_unifi_websocket(message=MessageKey.DEVICE, data=device_1)
-
-    mock_unifi_websocket(message=MessageKey.DEVICE, data=outlet_up1)
     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")
+    # Turn off outlet
-    assert outlet.state == STATE_ON
-
-    # Turn on and 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
+    # Availability signalling
-    hass.config_entries.async_update_entry(
+
-        config_entry,
+    # Controller disconnects
-        options={CONF_BLOCK_CLIENT: []},
+    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)