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
pull/133660/head
shapournemati-iotty 2024-12-20 15:33:05 +01:00 committed by GitHub
parent f49111a4d9
commit 1c0135880d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 268 additions and 77 deletions

View File

@ -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 = [
IottyLightSwitch(
coordinator=coordinator, iotty_cloud=coordinator.iotty, iotty_device=d
lightswitch_entities = [
IottySwitch(
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 (
any(d.device_id == device.device_id for d in known_devices)
or device.device_type != LS_DEVICE_TYPE_UID
if any(d.device_id == device.device_id for d in known_devices) or (
device.device_type not in {LS_DEVICE_TYPE_UID, OU_DEVICE_TYPE_UID}
):
continue
iotty_entity = IottyLightSwitch(
coordinator=coordinator,
iotty_cloud=coordinator.iotty,
iotty_device=LightSwitch(
iotty_entity: SwitchEntity
iotty_device: LightSwitch | Outlet
if device.device_type == LS_DEVICE_TYPE_UID:
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):
"""Haas entity class for iotty LightSwitch."""
class IottySwitch(IottyEntity, SwitchEntity):
"""Haas entity class for iotty switch."""
_attr_device_class = SwitchDeviceClass.SWITCH
_iotty_device: LightSwitch
_attr_device_class: SwitchDeviceClass | None
_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(
self._iotty_device.device_id, self._iotty_device.cmd_turn_on()
)
await self._iotty_cloud.command(self._iotty_device.device_id, COMMAND_TURNON)
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(
self._iotty_device.device_id, self._iotty_device.cmd_turn_off()
)
await self._iotty_cloud.command(self._iotty_device.device_id, COMMAND_TURNOFF)
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.async_write_ha_state()

View File

@ -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."""

View File

@ -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([
])

View File

@ -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)
config_entry_oauth2_flow.async_register_implementation(
hass, DOMAIN, local_oauth_impl
await check_command_ok(
entity_id=entity_id,
initial_status=STATUS_OFF,
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))
assert state.state == STATUS_OFF
async def test_turn_on_outlet_ok(
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(
SWITCH_DOMAIN,
SERVICE_TURN_ON,
{ATTR_ENTITY_ID: entity_id},
blocking=True,
await check_command_ok(
entity_id=entity_id,
initial_status=STATUS_OFF,
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.async_block_till_done()
mock_command_fn.assert_called_once()
assert (state := hass.states.get(entity_id))
assert state.state == STATUS_ON
async def test_turn_off_ok(
async def test_turn_off_light_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)
config_entry_oauth2_flow.async_register_implementation(
hass, DOMAIN, local_oauth_impl
await check_command_ok(
entity_id=entity_id,
initial_status=STATUS_ON,
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))
assert state.state == STATUS_ON
async def test_turn_off_outlet_ok(
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(
SWITCH_DOMAIN,
SERVICE_TURN_OFF,
{ATTR_ENTITY_ID: entity_id},
blocking=True,
await check_command_ok(
entity_id=entity_id,
initial_status=STATUS_ON,
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.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,