From af3bbdee822d8853beaf5d4207fea25c05a3f72e Mon Sep 17 00:00:00 2001
From: Tom <CoMPaTech@users.noreply.github.com>
Date: Mon, 22 May 2023 22:15:02 +0200
Subject: [PATCH] Plugwise add value_fn for switch (#93345)

* Plugwise add value_fn for switch

* Skip none as suggested
---
 homeassistant/components/plugwise/switch.py | 43 ++++++++++++++++-----
 1 file changed, 34 insertions(+), 9 deletions(-)

diff --git a/homeassistant/components/plugwise/switch.py b/homeassistant/components/plugwise/switch.py
index d6f8a420cc8..4204ab5a4d9 100644
--- a/homeassistant/components/plugwise/switch.py
+++ b/homeassistant/components/plugwise/switch.py
@@ -1,8 +1,12 @@
 """Plugwise Switch component for HomeAssistant."""
 from __future__ import annotations
 
+from collections.abc import Callable
+from dataclasses import dataclass
 from typing import Any
 
+from plugwise import SmileSwitches
+
 from homeassistant.components.switch import (
     SwitchDeviceClass,
     SwitchEntity,
@@ -18,29 +22,46 @@ from .coordinator import PlugwiseDataUpdateCoordinator
 from .entity import PlugwiseEntity
 from .util import plugwise_command
 
-SWITCHES: tuple[SwitchEntityDescription, ...] = (
-    SwitchEntityDescription(
+
+@dataclass
+class PlugwiseSwitchBaseMixin:
+    """Mixin for required Plugwise switch description keys."""
+
+    value_fn: Callable[[SmileSwitches], bool]
+
+
+@dataclass
+class PlugwiseSwitchEntityDescription(SwitchEntityDescription, PlugwiseSwitchBaseMixin):
+    """Describes Plugwise switch entity."""
+
+
+SWITCHES: tuple[PlugwiseSwitchEntityDescription, ...] = (
+    PlugwiseSwitchEntityDescription(
         key="dhw_cm_switch",
         translation_key="dhw_cm_switch",
         icon="mdi:water-plus",
         entity_category=EntityCategory.CONFIG,
+        value_fn=lambda data: data["dhw_cm_switch"],
     ),
-    SwitchEntityDescription(
+    PlugwiseSwitchEntityDescription(
         key="lock",
         translation_key="lock",
         icon="mdi:lock",
         entity_category=EntityCategory.CONFIG,
+        value_fn=lambda data: data["lock"],
     ),
-    SwitchEntityDescription(
+    PlugwiseSwitchEntityDescription(
         key="relay",
         translation_key="relay",
         device_class=SwitchDeviceClass.SWITCH,
+        value_fn=lambda data: data["relay"],
     ),
-    SwitchEntityDescription(
+    PlugwiseSwitchEntityDescription(
         key="cooling_ena_switch",
         name="Cooling",
         icon="mdi:snowflake-thermometer",
         entity_category=EntityCategory.CONFIG,
+        value_fn=lambda data: data["cooling_ena_switch"],
     ),
 )
 
@@ -54,8 +75,10 @@ async def async_setup_entry(
     coordinator = hass.data[DOMAIN][config_entry.entry_id]
     entities: list[PlugwiseSwitchEntity] = []
     for device_id, device in coordinator.data.devices.items():
+        if not (switches := device.get("switches")):
+            continue
         for description in SWITCHES:
-            if "switches" not in device or description.key not in device["switches"]:
+            if description.key not in switches:
                 continue
             entities.append(PlugwiseSwitchEntity(coordinator, device_id, description))
     async_add_entities(entities)
@@ -64,11 +87,13 @@ async def async_setup_entry(
 class PlugwiseSwitchEntity(PlugwiseEntity, SwitchEntity):
     """Representation of a Plugwise plug."""
 
+    entity_description: PlugwiseSwitchEntityDescription
+
     def __init__(
         self,
         coordinator: PlugwiseDataUpdateCoordinator,
         device_id: str,
-        description: SwitchEntityDescription,
+        description: PlugwiseSwitchEntityDescription,
     ) -> None:
         """Set up the Plugwise API."""
         super().__init__(coordinator, device_id)
@@ -76,9 +101,9 @@ class PlugwiseSwitchEntity(PlugwiseEntity, SwitchEntity):
         self._attr_unique_id = f"{device_id}-{description.key}"
 
     @property
-    def is_on(self) -> bool | None:
+    def is_on(self) -> bool:
         """Return True if entity is on."""
-        return self.device["switches"].get(self.entity_description.key)
+        return self.entity_description.value_fn(self.device["switches"])
 
     @plugwise_command
     async def async_turn_on(self, **kwargs: Any) -> None: