Add outlet device class to iotty switch entity (#132912)

* upgrade iottycloud lib to 0.3.0

* Add outlet

* test outlet turn on and turn off

* test add outlet

* Refactor code to use only one SwitchEntity  with an EntityDescription to distinguish Outlet and Lightswitch

* Refactor switch entities to reduce duplicated code

* Refactor tests to reduce duplicated code

* Refactor code to improve abstraction layer using specific types instead of generics

* Remove print and redundant field

homeassistant/components/iotty/switch.py

@@ -3,13 +3,22 @@
 from __future__ import annotations
 
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
-from iottycloud.device import Device
 from iottycloud.lightswitch import LightSwitch
-from iottycloud.verbs import LS_DEVICE_TYPE_UID
+from iottycloud.outlet import Outlet
+from iottycloud.verbs import (
+    COMMAND_TURNOFF,
+    COMMAND_TURNON,
+    LS_DEVICE_TYPE_UID,
+    OU_DEVICE_TYPE_UID,
+)
 
-from homeassistant.components.switch import SwitchDeviceClass, SwitchEntity
+from homeassistant.components.switch import (
+    SwitchDeviceClass,
+    SwitchEntity,
+    SwitchEntityDescription,
+)
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
@@ -20,31 +29,62 @@ from .entity import IottyEntity
 
 _LOGGER = logging.getLogger(__name__)
 
+ENTITIES: dict[str, SwitchEntityDescription] = {
+    LS_DEVICE_TYPE_UID: SwitchEntityDescription(
+        key="light",
+        name=None,
+        device_class=SwitchDeviceClass.SWITCH,
+    ),
+    OU_DEVICE_TYPE_UID: SwitchEntityDescription(
+        key="outlet",
+        name=None,
+        device_class=SwitchDeviceClass.OUTLET,
+    ),
+}
+
 
 async def async_setup_entry(
     hass: HomeAssistant,
     config_entry: IottyConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
-    """Activate the iotty LightSwitch component."""
+    """Activate the iotty Switch component."""
     _LOGGER.debug("Setup SWITCH entry id is %s", config_entry.entry_id)
 
     coordinator = config_entry.runtime_data.coordinator
-    entities = [
+    lightswitch_entities = [
-        IottyLightSwitch(
+        IottySwitch(
-            coordinator=coordinator, iotty_cloud=coordinator.iotty, iotty_device=d
+            coordinator=coordinator,
+            iotty_cloud=coordinator.iotty,
+            iotty_device=d,
+            entity_description=ENTITIES[LS_DEVICE_TYPE_UID],
         )
         for d in coordinator.data.devices
         if d.device_type == LS_DEVICE_TYPE_UID
         if (isinstance(d, LightSwitch))
     ]
-    _LOGGER.debug("Found %d LightSwitches", len(entities))
+    _LOGGER.debug("Found %d LightSwitches", len(lightswitch_entities))
+
+    outlet_entities = [
+        IottySwitch(
+            coordinator=coordinator,
+            iotty_cloud=coordinator.iotty,
+            iotty_device=d,
+            entity_description=ENTITIES[OU_DEVICE_TYPE_UID],
+        )
+        for d in coordinator.data.devices
+        if d.device_type == OU_DEVICE_TYPE_UID
+        if (isinstance(d, Outlet))
+    ]
+    _LOGGER.debug("Found %d Outlets", len(outlet_entities))
+
+    entities = lightswitch_entities + outlet_entities
 
     async_add_entities(entities)
 
     known_devices: set = config_entry.runtime_data.known_devices
     for known_device in coordinator.data.devices:
-        if known_device.device_type == LS_DEVICE_TYPE_UID:
+        if known_device.device_type in {LS_DEVICE_TYPE_UID, OU_DEVICE_TYPE_UID}:
             known_devices.add(known_device)
 
     @callback
@@ -59,21 +99,37 @@ async def async_setup_entry(
 
         # Add entities for devices which we've not yet seen
         for device in devices:
-            if (
+            if any(d.device_id == device.device_id for d in known_devices) or (
-                any(d.device_id == device.device_id for d in known_devices)
+                device.device_type not in {LS_DEVICE_TYPE_UID, OU_DEVICE_TYPE_UID}
-                or device.device_type != LS_DEVICE_TYPE_UID
             ):
                 continue
 
-            iotty_entity = IottyLightSwitch(
+            iotty_entity: SwitchEntity
-                coordinator=coordinator,
+            iotty_device: LightSwitch | Outlet
-                iotty_cloud=coordinator.iotty,
+            if device.device_type == LS_DEVICE_TYPE_UID:
-                iotty_device=LightSwitch(
+                if TYPE_CHECKING:
+                    assert isinstance(device, LightSwitch)
+                iotty_device = LightSwitch(
                     device.device_id,
                     device.serial_number,
                     device.device_type,
                     device.device_name,
-                ),
+                )
+            else:
+                if TYPE_CHECKING:
+                    assert isinstance(device, Outlet)
+                iotty_device = Outlet(
+                    device.device_id,
+                    device.serial_number,
+                    device.device_type,
+                    device.device_name,
+                )
+
+            iotty_entity = IottySwitch(
+                coordinator=coordinator,
+                iotty_cloud=coordinator.iotty,
+                iotty_device=iotty_device,
+                entity_description=ENTITIES[device.device_type],
             )
 
             entities.extend([iotty_entity])
@@ -85,24 +141,27 @@ async def async_setup_entry(
     coordinator.async_add_listener(async_update_data)
 
 
-class IottyLightSwitch(IottyEntity, SwitchEntity):
+class IottySwitch(IottyEntity, SwitchEntity):
-    """Haas entity class for iotty LightSwitch."""
+    """Haas entity class for iotty switch."""
 
-    _attr_device_class = SwitchDeviceClass.SWITCH
+    _attr_device_class: SwitchDeviceClass | None
-    _iotty_device: LightSwitch
+    _iotty_device: LightSwitch | Outlet
 
     def __init__(
         self,
         coordinator: IottyDataUpdateCoordinator,
         iotty_cloud: IottyProxy,
-        iotty_device: LightSwitch,
+        iotty_device: LightSwitch | Outlet,
+        entity_description: SwitchEntityDescription,
     ) -> None:
-        """Initialize the LightSwitch device."""
+        """Initialize the Switch device."""
         super().__init__(coordinator, iotty_cloud, iotty_device)
+        self.entity_description = entity_description
+        self._attr_device_class = entity_description.device_class
 
     @property
     def is_on(self) -> bool:
-        """Return true if the LightSwitch is on."""
+        """Return true if the Switch is on."""
         _LOGGER.debug(
             "Retrieve device status for %s ? %s",
             self._iotty_device.device_id,
@@ -111,30 +170,25 @@ class IottyLightSwitch(IottyEntity, SwitchEntity):
         return self._iotty_device.is_on
 
     async def async_turn_on(self, **kwargs: Any) -> None:
-        """Turn the LightSwitch on."""
+        """Turn the Switch on."""
         _LOGGER.debug("[%s] Turning on", self._iotty_device.device_id)
-        await self._iotty_cloud.command(
+        await self._iotty_cloud.command(self._iotty_device.device_id, COMMAND_TURNON)
-            self._iotty_device.device_id, self._iotty_device.cmd_turn_on()
-        )
         await self.coordinator.async_request_refresh()
 
     async def async_turn_off(self, **kwargs: Any) -> None:
-        """Turn the LightSwitch off."""
+        """Turn the Switch off."""
         _LOGGER.debug("[%s] Turning off", self._iotty_device.device_id)
-        await self._iotty_cloud.command(
+        await self._iotty_cloud.command(self._iotty_device.device_id, COMMAND_TURNOFF)
-            self._iotty_device.device_id, self._iotty_device.cmd_turn_off()
-        )
         await self.coordinator.async_request_refresh()
 
     @callback
     def _handle_coordinator_update(self) -> None:
         """Handle updated data from the coordinator."""
 
-        device: Device = next(
+        device: LightSwitch | Outlet = next(  # type: ignore[assignment]
             device
             for device in self.coordinator.data.devices
             if device.device_id == self._iotty_device.device_id
         )
-        if isinstance(device, LightSwitch):
+        self._iotty_device.is_on = device.is_on
-            self._iotty_device.is_on = device.is_on
         self.async_write_ha_state()

tests/components/iotty/conftest.py

@@ -6,10 +6,12 @@ from unittest.mock import AsyncMock, MagicMock, patch
 from aiohttp import ClientSession
 from iottycloud.device import Device
 from iottycloud.lightswitch import LightSwitch
+from iottycloud.outlet import Outlet
 from iottycloud.shutter import Shutter
 from iottycloud.verbs import (
     LS_DEVICE_TYPE_UID,
     OPEN_PERCENTAGE,
+    OU_DEVICE_TYPE_UID,
     RESULT,
     SH_DEVICE_TYPE_UID,
     STATUS,
@@ -73,6 +75,22 @@ test_sh_one_added = [
     sh_2,
 ]
 
+ou_0 = Outlet("TestOU", "TEST_SERIAL_OU_0", OU_DEVICE_TYPE_UID, "[TEST] Outlet 0")
+
+ou_1 = Outlet("TestOU1", "TEST_SERIAL_OU_1", OU_DEVICE_TYPE_UID, "[TEST] Outlet 1")
+
+ou_2 = Outlet("TestOU2", "TEST_SERIAL_OU_2", OU_DEVICE_TYPE_UID, "[TEST] Outlet 2")
+
+test_ou = [ou_0, ou_1]
+
+test_ou_one_removed = [ou_0]
+
+test_ou_one_added = [
+    ou_0,
+    ou_1,
+    ou_2,
+]
+
 
 @pytest.fixture
 async def local_oauth_impl(hass: HomeAssistant):
@@ -175,6 +193,16 @@ def mock_get_devices_twolightswitches() -> Generator[AsyncMock]:
         yield mock_fn
 
 
+@pytest.fixture
+def mock_get_devices_two_outlets() -> Generator[AsyncMock]:
+    """Mock for get_devices, returning two outlets."""
+
+    with patch(
+        "iottycloud.cloudapi.CloudApi.get_devices", return_value=test_ou
+    ) as mock_fn:
+        yield mock_fn
+
+
 @pytest.fixture
 def mock_get_devices_twoshutters() -> Generator[AsyncMock]:
     """Mock for get_devices, returning two shutters."""

tests/components/iotty/snapshots/test_switch.ambr

@@ -120,6 +120,19 @@
     'switch.test_light_switch_2_test_serial_2',
   ])
 # ---
+# name: test_outlet_insertion_ok
+  list([
+    'switch.test_outlet_0_test_serial_ou_0',
+    'switch.test_outlet_1_test_serial_ou_1',
+  ])
+# ---
+# name: test_outlet_insertion_ok.1
+  list([
+    'switch.test_outlet_0_test_serial_ou_0',
+    'switch.test_outlet_1_test_serial_ou_1',
+    'switch.test_outlet_2_test_serial_ou_2',
+  ])
+# ---
 # name: test_setup_entry_ok_nodevices
   list([
   ])

tests/components/iotty/test_switch.py

@@ -20,12 +20,52 @@ from homeassistant.helpers import (
     entity_registry as er,
 )
 
-from .conftest import test_ls_one_added, test_ls_one_removed
+from .conftest import test_ls_one_added, test_ls_one_removed, test_ou_one_added
 
 from tests.common import MockConfigEntry, async_fire_time_changed
 
 
-async def test_turn_on_ok(
+async def check_command_ok(
+    entity_id: str,
+    initial_status: str,
+    final_status: str,
+    command: str,
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    local_oauth_impl: ClientSession,
+    mock_get_status,
+    mock_command_fn,
+) -> None:
+    """Issue a command."""
+
+    mock_config_entry.add_to_hass(hass)
+
+    config_entry_oauth2_flow.async_register_implementation(
+        hass, DOMAIN, local_oauth_impl
+    )
+
+    await hass.config_entries.async_setup(mock_config_entry.entry_id)
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == initial_status
+
+    mock_get_status.return_value = {RESULT: {STATUS: final_status}}
+
+    await hass.services.async_call(
+        SWITCH_DOMAIN,
+        command,
+        {ATTR_ENTITY_ID: entity_id},
+        blocking=True,
+    )
+
+    await hass.async_block_till_done()
+    mock_command_fn.assert_called_once()
+
+    assert (state := hass.states.get(entity_id))
+    assert state.state == final_status
+
+
+async def test_turn_on_light_ok(
     hass: HomeAssistant,
     mock_config_entry: MockConfigEntry,
     local_oauth_impl: ClientSession,
@@ -37,34 +77,45 @@ async def test_turn_on_ok(
 
     entity_id = "switch.test_light_switch_0_test_serial_0"
 
-    mock_config_entry.add_to_hass(hass)
+    await check_command_ok(
-
+        entity_id=entity_id,
-    config_entry_oauth2_flow.async_register_implementation(
+        initial_status=STATUS_OFF,
-        hass, DOMAIN, local_oauth_impl
+        final_status=STATUS_ON,
+        command=SERVICE_TURN_ON,
+        hass=hass,
+        mock_config_entry=mock_config_entry,
+        local_oauth_impl=local_oauth_impl,
+        mock_get_status=mock_get_status_filled_off,
+        mock_command_fn=mock_command_fn,
     )
 
-    await hass.config_entries.async_setup(mock_config_entry.entry_id)
 
-    assert (state := hass.states.get(entity_id))
+async def test_turn_on_outlet_ok(
-    assert state.state == STATUS_OFF
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    local_oauth_impl: ClientSession,
+    mock_get_devices_two_outlets,
+    mock_get_status_filled_off,
+    mock_command_fn,
+) -> None:
+    """Issue a turnon command."""
 
-    mock_get_status_filled_off.return_value = {RESULT: {STATUS: STATUS_ON}}
+    entity_id = "switch.test_outlet_0_test_serial_ou_0"
 
-    await hass.services.async_call(
+    await check_command_ok(
-        SWITCH_DOMAIN,
+        entity_id=entity_id,
-        SERVICE_TURN_ON,
+        initial_status=STATUS_OFF,
-        {ATTR_ENTITY_ID: entity_id},
+        final_status=STATUS_ON,
-        blocking=True,
+        command=SERVICE_TURN_ON,
+        hass=hass,
+        mock_config_entry=mock_config_entry,
+        local_oauth_impl=local_oauth_impl,
+        mock_get_status=mock_get_status_filled_off,
+        mock_command_fn=mock_command_fn,
     )
 
-    await hass.async_block_till_done()
-    mock_command_fn.assert_called_once()
 
-    assert (state := hass.states.get(entity_id))
+async def test_turn_off_light_ok(
-    assert state.state == STATUS_ON
-
-
-async def test_turn_off_ok(
     hass: HomeAssistant,
     mock_config_entry: MockConfigEntry,
     local_oauth_impl: ClientSession,
@@ -76,32 +127,43 @@ async def test_turn_off_ok(
 
     entity_id = "switch.test_light_switch_0_test_serial_0"
 
-    mock_config_entry.add_to_hass(hass)
+    await check_command_ok(
-
+        entity_id=entity_id,
-    config_entry_oauth2_flow.async_register_implementation(
+        initial_status=STATUS_ON,
-        hass, DOMAIN, local_oauth_impl
+        final_status=STATUS_OFF,
+        command=SERVICE_TURN_OFF,
+        hass=hass,
+        mock_config_entry=mock_config_entry,
+        local_oauth_impl=local_oauth_impl,
+        mock_get_status=mock_get_status_filled,
+        mock_command_fn=mock_command_fn,
     )
 
-    await hass.config_entries.async_setup(mock_config_entry.entry_id)
 
-    assert (state := hass.states.get(entity_id))
+async def test_turn_off_outlet_ok(
-    assert state.state == STATUS_ON
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    local_oauth_impl: ClientSession,
+    mock_get_devices_two_outlets,
+    mock_get_status_filled,
+    mock_command_fn,
+) -> None:
+    """Issue a turnoff command."""
 
-    mock_get_status_filled.return_value = {RESULT: {STATUS: STATUS_OFF}}
+    entity_id = "switch.test_outlet_0_test_serial_ou_0"
 
-    await hass.services.async_call(
+    await check_command_ok(
-        SWITCH_DOMAIN,
+        entity_id=entity_id,
-        SERVICE_TURN_OFF,
+        initial_status=STATUS_ON,
-        {ATTR_ENTITY_ID: entity_id},
+        final_status=STATUS_OFF,
-        blocking=True,
+        command=SERVICE_TURN_OFF,
+        hass=hass,
+        mock_config_entry=mock_config_entry,
+        local_oauth_impl=local_oauth_impl,
+        mock_get_status=mock_get_status_filled,
+        mock_command_fn=mock_command_fn,
     )
 
-    await hass.async_block_till_done()
-    mock_command_fn.assert_called_once()
-
-    assert (state := hass.states.get(entity_id))
-    assert state.state == STATUS_OFF
-
 
 async def test_setup_entry_ok_nodevices(
     hass: HomeAssistant,
@@ -229,6 +291,40 @@ async def test_devices_insertion_ok(
     assert hass.states.async_entity_ids() == snapshot
 
 
+async def test_outlet_insertion_ok(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    local_oauth_impl: ClientSession,
+    mock_get_devices_two_outlets,
+    mock_get_status_filled,
+    snapshot: SnapshotAssertion,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test iotty switch insertion."""
+
+    mock_config_entry.add_to_hass(hass)
+
+    config_entry_oauth2_flow.async_register_implementation(
+        hass, DOMAIN, local_oauth_impl
+    )
+
+    assert await hass.config_entries.async_setup(mock_config_entry.entry_id)
+
+    # Should have two devices
+    assert hass.states.async_entity_ids_count() == 2
+    assert hass.states.async_entity_ids() == snapshot
+
+    mock_get_devices_two_outlets.return_value = test_ou_one_added
+
+    freezer.tick(UPDATE_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+    # Should have three devices
+    assert hass.states.async_entity_ids_count() == 3
+    assert hass.states.async_entity_ids() == snapshot
+
+
 async def test_api_not_ok_entities_stay_the_same_as_before(
     hass: HomeAssistant,
     mock_config_entry: MockConfigEntry,