From 41026b9227f624091406d30a6490dd21edff9c39 Mon Sep 17 00:00:00 2001
From: Bouwe Westerdijk <11290930+bouwew@users.noreply.github.com>
Date: Wed, 26 Jun 2024 11:04:00 +0200
Subject: [PATCH] Implement @plugwise_command for Plugwise Select platform
 (#120509)

---
 homeassistant/components/plugwise/const.py  |  1 +
 homeassistant/components/plugwise/select.py | 23 ++++++++-------------
 tests/components/plugwise/test_select.py    | 14 +++++++++----
 3 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/homeassistant/components/plugwise/const.py b/homeassistant/components/plugwise/const.py
index ed8cb2d2002..14599ce61fb 100644
--- a/homeassistant/components/plugwise/const.py
+++ b/homeassistant/components/plugwise/const.py
@@ -17,6 +17,7 @@ FLOW_SMILE: Final = "smile (Adam/Anna/P1)"
 FLOW_STRETCH: Final = "stretch (Stretch)"
 FLOW_TYPE: Final = "flow_type"
 GATEWAY: Final = "gateway"
+LOCATION: Final = "location"
 PW_TYPE: Final = "plugwise_type"
 SMILE: Final = "smile"
 STRETCH: Final = "stretch"
diff --git a/homeassistant/components/plugwise/select.py b/homeassistant/components/plugwise/select.py
index c8c9791c0da..99aecacb96b 100644
--- a/homeassistant/components/plugwise/select.py
+++ b/homeassistant/components/plugwise/select.py
@@ -2,27 +2,24 @@
 
 from __future__ import annotations
 
-from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
 
-from plugwise import Smile
-
 from homeassistant.components.select import SelectEntity, SelectEntityDescription
 from homeassistant.const import STATE_ON, EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import PlugwiseConfigEntry
-from .const import SelectOptionsType, SelectType
+from .const import LOCATION, SelectOptionsType, SelectType
 from .coordinator import PlugwiseDataUpdateCoordinator
 from .entity import PlugwiseEntity
+from .util import plugwise_command
 
 
 @dataclass(frozen=True, kw_only=True)
 class PlugwiseSelectEntityDescription(SelectEntityDescription):
     """Class describing Plugwise Select entities."""
 
-    command: Callable[[Smile, str, str], Awaitable[None]]
     key: SelectType
     options_key: SelectOptionsType
 
@@ -31,28 +28,24 @@ SELECT_TYPES = (
     PlugwiseSelectEntityDescription(
         key="select_schedule",
         translation_key="select_schedule",
-        command=lambda api, loc, opt: api.set_schedule_state(loc, STATE_ON, opt),
         options_key="available_schedules",
     ),
     PlugwiseSelectEntityDescription(
         key="select_regulation_mode",
         translation_key="regulation_mode",
         entity_category=EntityCategory.CONFIG,
-        command=lambda api, loc, opt: api.set_regulation_mode(opt),
         options_key="regulation_modes",
     ),
     PlugwiseSelectEntityDescription(
         key="select_dhw_mode",
         translation_key="dhw_mode",
         entity_category=EntityCategory.CONFIG,
-        command=lambda api, loc, opt: api.set_dhw_mode(opt),
         options_key="dhw_modes",
     ),
     PlugwiseSelectEntityDescription(
         key="select_gateway_mode",
         translation_key="gateway_mode",
         entity_category=EntityCategory.CONFIG,
-        command=lambda api, loc, opt: api.set_gateway_mode(opt),
         options_key="gateway_modes",
     ),
 )
@@ -109,10 +102,12 @@ class PlugwiseSelectEntity(PlugwiseEntity, SelectEntity):
         """Return the available select-options."""
         return self.device[self.entity_description.options_key]
 
+    @plugwise_command
     async def async_select_option(self, option: str) -> None:
-        """Change to the selected entity option."""
-        await self.entity_description.command(
-            self.coordinator.api, self.device["location"], option
-        )
+        """Change to the selected entity option.
 
-        await self.coordinator.async_request_refresh()
+        self.device[LOCATION] and STATE_ON are required for the thermostat-schedule select.
+        """
+        await self.coordinator.api.set_select(
+            self.entity_description.key, self.device[LOCATION], STATE_ON, option
+        )
diff --git a/tests/components/plugwise/test_select.py b/tests/components/plugwise/test_select.py
index 86b21af9e8b..a6245ff11e7 100644
--- a/tests/components/plugwise/test_select.py
+++ b/tests/components/plugwise/test_select.py
@@ -38,8 +38,9 @@ async def test_adam_change_select_entity(
         blocking=True,
     )
 
-    assert mock_smile_adam.set_schedule_state.call_count == 1
-    mock_smile_adam.set_schedule_state.assert_called_with(
+    assert mock_smile_adam.set_select.call_count == 1
+    mock_smile_adam.set_select.assert_called_with(
+        "select_schedule",
         "c50f167537524366a5af7aa3942feb1e",
         "on",
         "Badkamer Schema",
@@ -69,5 +70,10 @@ async def test_adam_select_regulation_mode(
         },
         blocking=True,
     )
-    assert mock_smile_adam_3.set_regulation_mode.call_count == 1
-    mock_smile_adam_3.set_regulation_mode.assert_called_with("heating")
+    assert mock_smile_adam_3.set_select.call_count == 1
+    mock_smile_adam_3.set_select.assert_called_with(
+        "select_regulation_mode",
+        "bc93488efab249e5bc54fd7e175a6f91",
+        "on",
+        "heating",
+    )