From 57b56010da6d790f56ca5fca3e76bd5f710e272b Mon Sep 17 00:00:00 2001
From: Frank Bakker <FrankBakkerNl@users.noreply.github.com>
Date: Sat, 13 Feb 2021 09:57:06 +0100
Subject: [PATCH] Added energy sensor to hlw8012 (#1198)

---
 esphome/components/hlw8012/hlw8012.cpp | 6 ++++++
 esphome/components/hlw8012/hlw8012.h   | 3 +++
 esphome/components/hlw8012/sensor.py   | 8 ++++++--
 tests/test1.yaml                       | 3 +++
 4 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/esphome/components/hlw8012/hlw8012.cpp b/esphome/components/hlw8012/hlw8012.cpp
index 2b3b4ec2d9..379dbe578e 100644
--- a/esphome/components/hlw8012/hlw8012.cpp
+++ b/esphome/components/hlw8012/hlw8012.cpp
@@ -79,6 +79,12 @@ void HLW8012Component::update() {
     this->power_sensor_->publish_state(power);
   }
 
+  if (this->energy_sensor_ != nullptr) {
+    cf_total_pulses_ += raw_cf;
+    float energy = cf_total_pulses_ * power_multiplier_micros / 3600 / 1000000.0f;
+    this->energy_sensor_->publish_state(energy);
+  }
+
   if (this->change_mode_at_++ == this->change_mode_every_) {
     this->current_mode_ = !this->current_mode_;
     ESP_LOGV(TAG, "Changing mode to %s mode", this->current_mode_ ? "CURRENT" : "VOLTAGE");
diff --git a/esphome/components/hlw8012/hlw8012.h b/esphome/components/hlw8012/hlw8012.h
index 4e5dc0f67f..af1f2e9a8c 100644
--- a/esphome/components/hlw8012/hlw8012.h
+++ b/esphome/components/hlw8012/hlw8012.h
@@ -29,6 +29,7 @@ class HLW8012Component : public PollingComponent {
   void set_voltage_sensor(sensor::Sensor *voltage_sensor) { voltage_sensor_ = voltage_sensor; }
   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; }
 
  protected:
   uint32_t nth_value_{0};
@@ -37,6 +38,7 @@ class HLW8012Component : public PollingComponent {
   uint32_t change_mode_every_{8};
   float current_resistor_{0.001};
   float voltage_divider_{2351};
+  uint64_t cf_total_pulses_{0};
   GPIOPin *sel_pin_;
   GPIOPin *cf_pin_;
   pulse_counter::PulseCounterStorage cf_store_;
@@ -45,6 +47,7 @@ class HLW8012Component : public PollingComponent {
   sensor::Sensor *voltage_sensor_{nullptr};
   sensor::Sensor *current_sensor_{nullptr};
   sensor::Sensor *power_sensor_{nullptr};
+  sensor::Sensor *energy_sensor_{nullptr};
 };
 
 }  // namespace hlw8012
diff --git a/esphome/components/hlw8012/sensor.py b/esphome/components/hlw8012/sensor.py
index e1f02b8fd2..e1a0ec6f18 100644
--- a/esphome/components/hlw8012/sensor.py
+++ b/esphome/components/hlw8012/sensor.py
@@ -3,8 +3,8 @@ import esphome.config_validation as cv
 from esphome import pins
 from esphome.components import sensor
 from esphome.const import CONF_CHANGE_MODE_EVERY, CONF_INITIAL_MODE, CONF_CURRENT, \
-    CONF_CURRENT_RESISTOR, CONF_ID, CONF_POWER, CONF_SEL_PIN, CONF_VOLTAGE, CONF_VOLTAGE_DIVIDER, \
-    ICON_FLASH, UNIT_VOLT, UNIT_AMPERE, UNIT_WATT
+    CONF_CURRENT_RESISTOR, CONF_ID, CONF_POWER, CONF_ENERGY, CONF_SEL_PIN, CONF_VOLTAGE, \
+    CONF_VOLTAGE_DIVIDER, ICON_FLASH, UNIT_VOLT, UNIT_AMPERE, UNIT_WATT, UNIT_WATT_HOURS
 
 AUTO_LOAD = ['pulse_counter']
 
@@ -29,6 +29,7 @@ CONFIG_SCHEMA = cv.Schema({
     cv.Optional(CONF_VOLTAGE): sensor.sensor_schema(UNIT_VOLT, ICON_FLASH, 1),
     cv.Optional(CONF_CURRENT): sensor.sensor_schema(UNIT_AMPERE, ICON_FLASH, 2),
     cv.Optional(CONF_POWER): sensor.sensor_schema(UNIT_WATT, ICON_FLASH, 1),
+    cv.Optional(CONF_ENERGY): sensor.sensor_schema(UNIT_WATT_HOURS, ICON_FLASH, 1),
 
     cv.Optional(CONF_CURRENT_RESISTOR, default=0.001): cv.resistance,
     cv.Optional(CONF_VOLTAGE_DIVIDER, default=2351): cv.positive_float,
@@ -57,6 +58,9 @@ def to_code(config):
     if CONF_POWER in config:
         sens = yield sensor.new_sensor(config[CONF_POWER])
         cg.add(var.set_power_sensor(sens))
+    if CONF_ENERGY in config:
+        sens = yield sensor.new_sensor(config[CONF_ENERGY])
+        cg.add(var.set_energy_sensor(sens))
     cg.add(var.set_current_resistor(config[CONF_CURRENT_RESISTOR]))
     cg.add(var.set_voltage_divider(config[CONF_VOLTAGE_DIVIDER]))
     cg.add(var.set_change_mode_every(config[CONF_CHANGE_MODE_EVERY]))
diff --git a/tests/test1.yaml b/tests/test1.yaml
index 5aed1dac44..fceff00e24 100644
--- a/tests/test1.yaml
+++ b/tests/test1.yaml
@@ -434,6 +434,9 @@ sensor:
     power:
       name: 'HLW8012 Power'
       id: hlw8012_power
+    energy:
+      name: "HLW8012 Energy"
+      id: hlw8012_energy
     update_interval: 15s
     current_resistor: 0.001 ohm
     voltage_divider: 2351