From f4f7122c6687b5061f895c3d05e144f8180c7ba3 Mon Sep 17 00:00:00 2001
From: TheJulianJES <TheJulianJES@users.noreply.github.com>
Date: Sat, 22 Oct 2022 02:56:49 +0200
Subject: [PATCH] Add "power on state" config entity for Tuya plugs to ZHA
 (#80486)

* Add "power on state" config entity for Tuya plugs to ZHA

* Remove TS004F from models, as it doesn't support power-on-state
---
 .../components/zha/core/channels/general.py   | 13 ++++++++++++
 homeassistant/components/zha/select.py        | 20 +++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/homeassistant/components/zha/core/channels/general.py b/homeassistant/components/zha/core/channels/general.py
index d310157327b..f5dbab9e38e 100644
--- a/homeassistant/components/zha/core/channels/general.py
+++ b/homeassistant/components/zha/core/channels/general.py
@@ -331,6 +331,19 @@ class OnOffChannel(ZigbeeChannel):
         super().__init__(cluster, ch_pool)
         self._off_listener = None
 
+        if self.cluster.endpoint.model in (
+            "TS011F",
+            "TS0121",
+            "TS0001",
+            "TS0002",
+            "TS0003",
+            "TS0004",
+        ):
+            self.ZCL_INIT_ATTRS = (  # pylint: disable=invalid-name
+                self.ZCL_INIT_ATTRS.copy()
+            )
+            self.ZCL_INIT_ATTRS["power_on_state"] = True
+
     @property
     def on_off(self) -> bool | None:
         """Return cached value of on/off attribute."""
diff --git a/homeassistant/components/zha/select.py b/homeassistant/components/zha/select.py
index c2f315cd217..35746951b47 100644
--- a/homeassistant/components/zha/select.py
+++ b/homeassistant/components/zha/select.py
@@ -227,6 +227,26 @@ class ZHAStartupOnOffSelectEntity(
     _attr_name = "Start-up behavior"
 
 
+class TuyaPowerOnState(types.enum8):
+    """Tuya power on state enum."""
+
+    Off = 0x00
+    On = 0x01
+    LastState = 0x02
+
+
+@CONFIG_DIAGNOSTIC_MATCH(
+    channel_names=CHANNEL_ON_OFF,
+    models={"TS011F", "TS0121", "TS0001", "TS0002", "TS0003", "TS0004"},
+)
+class TuyaPowerOnStateSelectEntity(ZCLEnumSelectEntity, id_suffix="power_on_state"):
+    """Representation of a ZHA power on state select entity."""
+
+    _select_attr = "power_on_state"
+    _enum = TuyaPowerOnState
+    _attr_name = "Power on state"
+
+
 class AqaraMotionSensitivities(types.enum8):
     """Aqara motion sensitivities."""