select api connected

Author: dmulcahey

tests/test_select.py

@@ -2,6 +2,7 @@
 
 from unittest.mock import call, patch
 
+import pytest
 from zhaquirks import (
     DEVICE_TYPE,
     ENDPOINTS,
@@ -27,14 +28,19 @@
     join_zigpy_device,
     send_attributes_report,
 )
+from tests.conftest import CombinedGateways
 from zha.application import Platform
-from zha.application.gateway import Gateway
-from zha.application.platforms import EntityCategory
+from zha.application.platforms import EntityCategory, PlatformEntity
 from zha.application.platforms.select import AqaraMotionSensitivities
 
 
-async def test_select(zha_gateway: Gateway) -> None:
+@pytest.mark.parametrize(
+    "gateway_type",
+    ["zha_gateway", "ws_gateway"],
+)
+async def test_select(zha_gateways: CombinedGateways, gateway_type: str) -> None:
     """Test zha select platform."""
+    zha_gateway = getattr(zha_gateways, gateway_type)
     zigpy_device = create_mock_zigpy_device(
         zha_gateway,
         {
@@ -63,7 +69,9 @@ async def test_select(zha_gateway: Gateway) -> None:
         "Fire Panic",
         "Emergency Panic",
     ]
-    assert entity._enum == security.IasWd.Warning.WarningMode
+
+    if isinstance(entity, PlatformEntity):
+        assert entity._enum == security.IasWd.Warning.WarningMode
 
     # change value from client
     await entity.async_select_option(security.IasWd.Warning.WarningMode.Burglar.name)
@@ -107,9 +115,16 @@ def __init__(self, *args, **kwargs):
     }
 
 
-async def test_on_off_select_attribute_report(zha_gateway: Gateway) -> None:
+@pytest.mark.parametrize(
+    "gateway_type",
+    ["zha_gateway", "ws_gateway"],
+)
+async def test_on_off_select_attribute_report(
+    zha_gateways: CombinedGateways, gateway_type: str
+) -> None:
     """Test ZHA attribute report parsing for select platform."""
 
+    zha_gateway = getattr(zha_gateways, gateway_type)
     zigpy_device = create_mock_zigpy_device(
         zha_gateway,
         {
@@ -126,7 +141,7 @@ async def test_on_off_select_attribute_report(zha_gateway: Gateway) -> None:
 
     zigpy_device = get_device(zigpy_device)
     aqara_sensor = await join_zigpy_device(zha_gateway, zigpy_device)
-    cluster = aqara_sensor.device.endpoints.get(1).opple_cluster
+    cluster = zigpy_device.endpoints.get(1).opple_cluster
 
     entity = get_entity(aqara_sensor, platform=Platform.SELECT)
     assert entity.state["state"] == AqaraMotionSensitivities.Medium.name
@@ -160,11 +175,16 @@ async def test_on_off_select_attribute_report(zha_gateway: Gateway) -> None:
 )
 
 
+@pytest.mark.parametrize(
+    "gateway_type",
+    ["zha_gateway", "ws_gateway"],
+)
 async def test_on_off_select_attribute_report_v2(
-    zha_gateway: Gateway,
+    zha_gateways: CombinedGateways, gateway_type: str
 ) -> None:
     """Test ZHA attribute report parsing for select platform."""
 
+    zha_gateway = getattr(zha_gateways, gateway_type)
     zigpy_device = create_mock_zigpy_device(
         zha_gateway,
         {
@@ -184,7 +204,7 @@ async def test_on_off_select_attribute_report_v2(
 
     zha_device = await join_zigpy_device(zha_gateway, zigpy_device)
     cluster = zigpy_device.endpoints[1].opple_cluster
-    assert isinstance(zha_device.device, CustomDeviceV2)
+    assert isinstance(zigpy_device, CustomDeviceV2)
 
     entity = get_entity(zha_device, platform=Platform.SELECT)
 
@@ -228,8 +248,15 @@ async def test_on_off_select_attribute_report_v2(
         )
 
 
-async def test_non_zcl_select_state_restoration(zha_gateway: Gateway) -> None:
+@pytest.mark.parametrize(
+    "gateway_type",
+    ["zha_gateway", "ws_gateway"],
+)
+async def test_non_zcl_select_state_restoration(
+    zha_gateways: CombinedGateways, gateway_type: str
+) -> None:
     """Test the non-ZCL select state restoration."""
+    zha_gateway = getattr(zha_gateways, gateway_type)
     zigpy_device = create_mock_zigpy_device(
         zha_gateway,
         {
@@ -251,9 +278,11 @@ async def test_non_zcl_select_state_restoration(zha_gateway: Gateway) -> None:
     entity.restore_external_state_attributes(
         state=security.IasWd.Warning.WarningMode.Burglar.name
     )
+    await zha_gateway.async_block_till_done()  # needed for WS operations
     assert entity.state["state"] == security.IasWd.Warning.WarningMode.Burglar.name
 
     entity.restore_external_state_attributes(
         state=security.IasWd.Warning.WarningMode.Fire.name
     )
+    await zha_gateway.async_block_till_done()  # needed for WS operations
     assert entity.state["state"] == security.IasWd.Warning.WarningMode.Fire.name

zha/application/platforms/select/__init__.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+import asyncio
 from enum import Enum
 import functools
 import logging
@@ -116,21 +117,18 @@ def current_option(self) -> str | None:
             return None
         return option.name.replace("_", " ")
 
-    async def async_select_option(self, option: str) -> None:
+    async def async_select_option(self, option: str, **kwargs) -> None:
         """Change the selected option."""
         self._cluster_handler.data_cache[self._attribute_name] = self._enum[
             option.replace(" ", "_")
         ]
         self.maybe_emit_state_changed_event()
 
-    def restore_external_state_attributes(
-        self,
-        *,
-        state: str,
-    ) -> None:
+    def restore_external_state_attributes(self, *, state: str, **kwargs) -> None:
         """Restore extra state attributes that are stored outside of the ZCL cache."""
         value = state.replace(" ", "_")
         self._cluster_handler.data_cache[self._attribute_name] = self._enum[value]
+        self.maybe_emit_state_changed_event()
 
 
 class NonZCLSelectEntity(EnumSelectEntity):
@@ -262,7 +260,7 @@ def current_option(self) -> str | None:
         option = self._enum(option)
         return option.name.replace("_", " ")
 
-    async def async_select_option(self, option: str) -> None:
+    async def async_select_option(self, option: str, **kwargs) -> None:
         """Change the selected option."""
         await self._cluster_handler.write_attributes_safe(
             {self._attribute_name: self._enum[option.replace(" ", "_")]}
@@ -916,17 +914,26 @@ def __init__(
     ) -> None:
         """Initialize the ZHA select entity."""
         super().__init__(entity_info, device)
+        self._tasks: list[asyncio.Task] = []
 
     @property
     def current_option(self) -> str | None:
         """Return the selected entity option to represent the entity state."""
 
     async def async_select_option(self, option: str) -> None:
         """Change the selected option."""
+        await self._device.gateway.selects.select_option(self.info_object, option)
 
     def restore_external_state_attributes(
         self,
         *,
         state: str,
     ) -> None:
         """Restore extra state attributes."""
+        task = asyncio.create_task(
+            self._device.gateway.selects.restore_external_state_attributes(
+                self.info_object, state
+            )
+        )
+        self._tasks.append(task)
+        task.add_done_callback(self._tasks.remove)

zha/application/platforms/select/websocket_api.py

@@ -38,6 +38,28 @@ async def select_option(
     )
 
 
+class SelectRestoreExternalStateAttributesCommand(PlatformEntityCommand):
+    """Select restore external state command."""
+
+    command: Literal[APICommands.SELECT_RESTORE_EXTERNAL_STATE_ATTRIBUTES] = (
+        APICommands.SELECT_RESTORE_EXTERNAL_STATE_ATTRIBUTES
+    )
+    platform: str = Platform.SELECT
+    state: str
+
+
+@decorators.websocket_command(SelectRestoreExternalStateAttributesCommand)
+@decorators.async_response
+async def restore_lock_external_state_attributes(
+    server: Server, client: Client, command: SelectRestoreExternalStateAttributesCommand
+) -> None:
+    """Restore externally preserved state for selects."""
+    await execute_platform_entity_command(
+        server, client, command, "restore_external_state_attributes"
+    )
+
+
 def load_api(server: Server) -> None:
     """Load the api command handlers."""
     register_api_command(server, select_option)
+    register_api_command(server, restore_lock_external_state_attributes)