From 77ccf46c9b99d0d6fbbc9c2be74b2881dcbf1953 Mon Sep 17 00:00:00 2001
From: Dmitry Vasilyev <slydiman@gmail.com>
Date: Thu, 6 Jan 2022 15:01:16 +0400
Subject: [PATCH] Updated the support of Tuya Circuit Breaker 'dlq' (#63519)

Co-authored-by: Franck Nijhof <frenck@frenck.nl>
---
 homeassistant/components/tuya/base.py   | 13 ++++
 homeassistant/components/tuya/const.py  |  1 +
 homeassistant/components/tuya/sensor.py | 89 +++++++++++++++++++++++++
 homeassistant/components/tuya/switch.py |  2 +-
 4 files changed, 104 insertions(+), 1 deletion(-)

diff --git a/homeassistant/components/tuya/base.py b/homeassistant/components/tuya/base.py
index 38abac9a0fa..2a2a76cfe2f 100644
--- a/homeassistant/components/tuya/base.py
+++ b/homeassistant/components/tuya/base.py
@@ -1,8 +1,10 @@
 """Tuya Home Assistant Base Device Model."""
 from __future__ import annotations
 
+import base64
 from dataclasses import dataclass
 import json
+import struct
 from typing import Any
 
 from tuya_iot import TuyaDevice, TuyaDeviceManager
@@ -99,6 +101,17 @@ class ElectricityTypeData:
         """Load JSON string and return a ElectricityTypeData object."""
         return cls(**json.loads(data.lower()))
 
+    @classmethod
+    def from_raw(cls, data: str) -> ElectricityTypeData:
+        """Decode base64 string and return a ElectricityTypeData object."""
+        raw = base64.b64decode(data)
+        voltage = struct.unpack(">H", raw[0:2])[0] / 10.0
+        electriccurrent = struct.unpack(">L", b"\x00" + raw[2:5])[0] / 1000.0
+        power = struct.unpack(">L", b"\x00" + raw[5:8])[0] / 1000.0
+        return cls(
+            electriccurrent=str(electriccurrent), power=str(power), voltage=str(voltage)
+        )
+
 
 class TuyaEntity(Entity):
     """Tuya base device."""
diff --git a/homeassistant/components/tuya/const.py b/homeassistant/components/tuya/const.py
index c8fe6c07a95..c66cf5686ea 100644
--- a/homeassistant/components/tuya/const.py
+++ b/homeassistant/components/tuya/const.py
@@ -312,6 +312,7 @@ class DPCode(StrEnum):
     TOTAL_CLEAN_AREA = "total_clean_area"
     TOTAL_CLEAN_COUNT = "total_clean_count"
     TOTAL_CLEAN_TIME = "total_clean_time"
+    TOTAL_FORWARD_ENERGY = "total_forward_energy"
     UV = "uv"  # UV sterilization
     VA_BATTERY = "va_battery"
     VA_HUMIDITY = "va_humidity"
diff --git a/homeassistant/components/tuya/sensor.py b/homeassistant/components/tuya/sensor.py
index b738ae75cd0..3da1e75a80d 100644
--- a/homeassistant/components/tuya/sensor.py
+++ b/homeassistant/components/tuya/sensor.py
@@ -589,6 +589,88 @@ SENSORS: dict[str, tuple[TuyaSensorEntityDescription, ...]] = {
             subkey="voltage",
         ),
     ),
+    # Circuit Breaker
+    # https://developer.tuya.com/en/docs/iot/dlq?id=Kb0kidk9enyh8
+    "dlq": (
+        TuyaSensorEntityDescription(
+            key=DPCode.TOTAL_FORWARD_ENERGY,
+            name="Total Energy",
+            device_class=SensorDeviceClass.ENERGY,
+            state_class=SensorStateClass.TOTAL_INCREASING,
+        ),
+        TuyaSensorEntityDescription(
+            key=DPCode.PHASE_A,
+            name="Phase A Current",
+            device_class=SensorDeviceClass.CURRENT,
+            native_unit_of_measurement=ELECTRIC_CURRENT_AMPERE,
+            state_class=SensorStateClass.MEASUREMENT,
+            subkey="electriccurrent",
+        ),
+        TuyaSensorEntityDescription(
+            key=DPCode.PHASE_A,
+            name="Phase A Power",
+            device_class=SensorDeviceClass.POWER,
+            state_class=SensorStateClass.MEASUREMENT,
+            native_unit_of_measurement=POWER_KILO_WATT,
+            subkey="power",
+        ),
+        TuyaSensorEntityDescription(
+            key=DPCode.PHASE_A,
+            name="Phase A Voltage",
+            device_class=SensorDeviceClass.VOLTAGE,
+            state_class=SensorStateClass.MEASUREMENT,
+            native_unit_of_measurement=ELECTRIC_POTENTIAL_VOLT,
+            subkey="voltage",
+        ),
+        TuyaSensorEntityDescription(
+            key=DPCode.PHASE_B,
+            name="Phase B Current",
+            device_class=SensorDeviceClass.CURRENT,
+            native_unit_of_measurement=ELECTRIC_CURRENT_AMPERE,
+            state_class=SensorStateClass.MEASUREMENT,
+            subkey="electriccurrent",
+        ),
+        TuyaSensorEntityDescription(
+            key=DPCode.PHASE_B,
+            name="Phase B Power",
+            device_class=SensorDeviceClass.POWER,
+            state_class=SensorStateClass.MEASUREMENT,
+            native_unit_of_measurement=POWER_KILO_WATT,
+            subkey="power",
+        ),
+        TuyaSensorEntityDescription(
+            key=DPCode.PHASE_B,
+            name="Phase B Voltage",
+            device_class=SensorDeviceClass.VOLTAGE,
+            state_class=SensorStateClass.MEASUREMENT,
+            native_unit_of_measurement=ELECTRIC_POTENTIAL_VOLT,
+            subkey="voltage",
+        ),
+        TuyaSensorEntityDescription(
+            key=DPCode.PHASE_C,
+            name="Phase C Current",
+            device_class=SensorDeviceClass.CURRENT,
+            native_unit_of_measurement=ELECTRIC_CURRENT_AMPERE,
+            state_class=SensorStateClass.MEASUREMENT,
+            subkey="electriccurrent",
+        ),
+        TuyaSensorEntityDescription(
+            key=DPCode.PHASE_C,
+            name="Phase C Power",
+            device_class=SensorDeviceClass.POWER,
+            state_class=SensorStateClass.MEASUREMENT,
+            native_unit_of_measurement=POWER_KILO_WATT,
+            subkey="power",
+        ),
+        TuyaSensorEntityDescription(
+            key=DPCode.PHASE_C,
+            name="Phase C Voltage",
+            device_class=SensorDeviceClass.VOLTAGE,
+            state_class=SensorStateClass.MEASUREMENT,
+            native_unit_of_measurement=ELECTRIC_POTENTIAL_VOLT,
+            subkey="voltage",
+        ),
+    ),
     # Robot Vacuum
     # https://developer.tuya.com/en/docs/iot/fsd?id=K9gf487ck1tlo
     "sd": (
@@ -765,6 +847,7 @@ class TuyaSensorEntity(TuyaEntity, SensorEntity):
             "String",
             "Enum",
             "Json",
+            "Raw",
         ):
             return None
 
@@ -794,5 +877,11 @@ class TuyaSensorEntity(TuyaEntity, SensorEntity):
             values = ElectricityTypeData.from_json(value)
             return getattr(values, self.entity_description.subkey)
 
+        if self._status_range.type == "Raw":
+            if self.entity_description.subkey is None:
+                return None
+            values = ElectricityTypeData.from_raw(value)
+            return getattr(values, self.entity_description.subkey)
+
         # Valid string or enum value
         return value
diff --git a/homeassistant/components/tuya/switch.py b/homeassistant/components/tuya/switch.py
index 3085b7f439a..15d35601379 100644
--- a/homeassistant/components/tuya/switch.py
+++ b/homeassistant/components/tuya/switch.py
@@ -94,7 +94,7 @@ SWITCHES: dict[str, tuple[SwitchEntityDescription, ...]] = {
             entity_category=EntityCategory.CONFIG,
         ),
         SwitchEntityDescription(
-            key=DPCode.SWITCH_1,
+            key=DPCode.SWITCH,
             name="Switch",
         ),
     ),