From 1087cb55b4b04a6ca1e0c2fc31e47df9404b10df Mon Sep 17 00:00:00 2001
From: Kai Gerken <Kai.Gerken@custom-coding.net>
Date: Wed, 8 Mar 2023 20:00:44 +0100
Subject: [PATCH] Added pzemdc reset energy action (#4481)

* remove unused sensors on pzemdc

* add reset energy action for pzemdc

* fix lint errors

* remove trailing space on pzemdc.h

* make method reset_energy of pzemdc public

* Apply suggestions from code review

---------

Co-authored-by: Jesse Hills <3060199+jesserockz@users.noreply.github.com>
---
 esphome/components/pzemdc/pzemdc.cpp |  8 ++++++++
 esphome/components/pzemdc/pzemdc.h   | 16 ++++++++++++----
 esphome/components/pzemdc/sensor.py  | 19 +++++++++++++++++++
 tests/test3.yaml                     |  6 ++++++
 4 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/esphome/components/pzemdc/pzemdc.cpp b/esphome/components/pzemdc/pzemdc.cpp
index c914edefc8..28e4210ff7 100644
--- a/esphome/components/pzemdc/pzemdc.cpp
+++ b/esphome/components/pzemdc/pzemdc.cpp
@@ -7,6 +7,7 @@ namespace pzemdc {
 static const char *const TAG = "pzemdc";
 
 static const uint8_t PZEM_CMD_READ_IN_REGISTERS = 0x04;
+static const uint8_t PZEM_CMD_RESET_ENERGY = 0x42;
 static const uint8_t PZEM_REGISTER_COUNT = 10;  // 10x 16-bit registers
 
 void PZEMDC::on_modbus_data(const std::vector<uint8_t> &data) {
@@ -61,5 +62,12 @@ void PZEMDC::dump_config() {
   LOG_SENSOR("", "Energy", this->energy_sensor_);
 }
 
+void PZEMDC::reset_energy() {
+  std::vector<uint8_t> cmd;
+  cmd.push_back(this->address_);
+  cmd.push_back(PZEM_CMD_RESET_ENERGY);
+  this->send_raw(cmd);
+}
+
 }  // namespace pzemdc
 }  // namespace esphome
diff --git a/esphome/components/pzemdc/pzemdc.h b/esphome/components/pzemdc/pzemdc.h
index 54df6c484c..21676e3422 100644
--- a/esphome/components/pzemdc/pzemdc.h
+++ b/esphome/components/pzemdc/pzemdc.h
@@ -15,8 +15,6 @@ class PZEMDC : public PollingComponent, public modbus::ModbusDevice {
   void set_current_sensor(sensor::Sensor *current_sensor) { current_sensor_ = current_sensor; }
   void set_power_sensor(sensor::Sensor *power_sensor) { power_sensor_ = power_sensor; }
   void set_energy_sensor(sensor::Sensor *energy_sensor) { energy_sensor_ = energy_sensor; }
-  void set_frequency_sensor(sensor::Sensor *frequency_sensor) { frequency_sensor_ = frequency_sensor; }
-  void set_powerfactor_sensor(sensor::Sensor *powerfactor_sensor) { power_factor_sensor_ = powerfactor_sensor; }
 
   void update() override;
 
@@ -24,14 +22,24 @@ class PZEMDC : public PollingComponent, public modbus::ModbusDevice {
 
   void dump_config() override;
 
+  void reset_energy();
+
  protected:
   sensor::Sensor *voltage_sensor_{nullptr};
   sensor::Sensor *current_sensor_{nullptr};
   sensor::Sensor *power_sensor_{nullptr};
-  sensor::Sensor *frequency_sensor_{nullptr};
-  sensor::Sensor *power_factor_sensor_{nullptr};
   sensor::Sensor *energy_sensor_{nullptr};
 };
 
+template<typename... Ts> class ResetEnergyAction : public Action<Ts...> {
+ public:
+  ResetEnergyAction(PZEMDC *pzemdc) : pzemdc_(pzemdc) {}
+
+  void play(Ts... x) override { this->pzemdc_->reset_energy(); }
+
+ protected:
+  PZEMDC *pzemdc_;
+};
+
 }  // namespace pzemdc
 }  // namespace esphome
diff --git a/esphome/components/pzemdc/sensor.py b/esphome/components/pzemdc/sensor.py
index 4234ce19ad..097b1c1cfd 100644
--- a/esphome/components/pzemdc/sensor.py
+++ b/esphome/components/pzemdc/sensor.py
@@ -1,5 +1,7 @@
 import esphome.codegen as cg
 import esphome.config_validation as cv
+from esphome import automation
+from esphome.automation import maybe_simple_id
 from esphome.components import sensor, modbus
 from esphome.const import (
     CONF_CURRENT,
@@ -24,6 +26,9 @@ AUTO_LOAD = ["modbus"]
 pzemdc_ns = cg.esphome_ns.namespace("pzemdc")
 PZEMDC = pzemdc_ns.class_("PZEMDC", cg.PollingComponent, modbus.ModbusDevice)
 
+# Actions
+ResetEnergyAction = pzemdc_ns.class_("ResetEnergyAction", automation.Action)
+
 CONFIG_SCHEMA = (
     cv.Schema(
         {
@@ -59,6 +64,20 @@ CONFIG_SCHEMA = (
 )
 
 
+@automation.register_action(
+    "pzemdc.reset_energy",
+    ResetEnergyAction,
+    maybe_simple_id(
+        {
+            cv.GenerateID(CONF_ID): cv.use_id(PZEMDC),
+        }
+    ),
+)
+async def reset_energy_to_code(config, action_id, template_arg, args):
+    paren = await cg.get_variable(config[CONF_ID])
+    return cg.new_Pvariable(action_id, template_arg, paren)
+
+
 async def to_code(config):
     var = cg.new_Pvariable(config[CONF_ID])
     await cg.register_component(var, config)
diff --git a/tests/test3.yaml b/tests/test3.yaml
index e62a171c95..f1b64b3389 100644
--- a/tests/test3.yaml
+++ b/tests/test3.yaml
@@ -579,6 +579,7 @@ sensor:
     power_factor:
       name: PZEMAC Power Factor
   - platform: pzemdc
+    id: pzemdc1
     voltage:
       name: PZEMDC Voltage
     current:
@@ -925,6 +926,11 @@ binary_sensor:
     on_press:
       then:
         - pzemac.reset_energy: pzemac1
+  - platform: template
+    id: pzemdc_reset_energy
+    on_press:
+      then:
+        - pzemdc.reset_energy: pzemdc1
 
   - platform: vbus
     model: deltasol_bs_plus