From 88ccd60a085926fc874fbcd0c71edfd628252a4a Mon Sep 17 00:00:00 2001
From: Brandon Davidson <brad@oatmail.org>
Date: Thu, 6 Jun 2019 13:44:52 -0700
Subject: [PATCH] Allow setting the initial mode of HLW8012 sensors (#611)

* Allow setting the initial mode of hlw8012 sensors

* Changes as per code review
---
 esphome/components/hlw8012/hlw8012.h | 5 +++++
 esphome/components/hlw8012/sensor.py | 9 ++++++++-
 esphome/const.py                     | 1 +
 tests/test1.yaml                     | 1 +
 4 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/esphome/components/hlw8012/hlw8012.h b/esphome/components/hlw8012/hlw8012.h
index b9321b51c6..4e5dc0f67f 100644
--- a/esphome/components/hlw8012/hlw8012.h
+++ b/esphome/components/hlw8012/hlw8012.h
@@ -8,6 +8,8 @@
 namespace esphome {
 namespace hlw8012 {
 
+enum HLW8012InitialMode { HLW8012_INITIAL_MODE_CURRENT = 0, HLW8012_INITIAL_MODE_VOLTAGE };
+
 class HLW8012Component : public PollingComponent {
  public:
   void setup() override;
@@ -15,6 +17,9 @@ class HLW8012Component : public PollingComponent {
   float get_setup_priority() const override;
   void update() override;
 
+  void set_initial_mode(HLW8012InitialMode initial_mode) {
+    current_mode_ = initial_mode == HLW8012_INITIAL_MODE_CURRENT;
+  }
   void set_change_mode_every(uint32_t change_mode_every) { change_mode_every_ = change_mode_every; }
   void set_current_resistor(float current_resistor) { current_resistor_ = current_resistor; }
   void set_voltage_divider(float voltage_divider) { voltage_divider_ = voltage_divider; }
diff --git a/esphome/components/hlw8012/sensor.py b/esphome/components/hlw8012/sensor.py
index 697c34f9d2..e1f02b8fd2 100644
--- a/esphome/components/hlw8012/sensor.py
+++ b/esphome/components/hlw8012/sensor.py
@@ -2,7 +2,7 @@ import esphome.codegen as cg
 import esphome.config_validation as cv
 from esphome import pins
 from esphome.components import sensor
-from esphome.const import CONF_CHANGE_MODE_EVERY, CONF_CURRENT, \
+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
 
@@ -10,6 +10,11 @@ AUTO_LOAD = ['pulse_counter']
 
 hlw8012_ns = cg.esphome_ns.namespace('hlw8012')
 HLW8012Component = hlw8012_ns.class_('HLW8012Component', cg.PollingComponent)
+HLW8012InitialMode = hlw8012_ns.enum('HLW8012InitialMode')
+INITIAL_MODES = {
+    CONF_CURRENT: HLW8012InitialMode.HLW8012_INITIAL_MODE_CURRENT,
+    CONF_VOLTAGE: HLW8012InitialMode.HLW8012_INITIAL_MODE_VOLTAGE,
+}
 
 CONF_CF1_PIN = 'cf1_pin'
 CONF_CF_PIN = 'cf_pin'
@@ -28,6 +33,7 @@ CONFIG_SCHEMA = cv.Schema({
     cv.Optional(CONF_CURRENT_RESISTOR, default=0.001): cv.resistance,
     cv.Optional(CONF_VOLTAGE_DIVIDER, default=2351): cv.positive_float,
     cv.Optional(CONF_CHANGE_MODE_EVERY, default=8): cv.All(cv.uint32_t, cv.Range(min=1)),
+    cv.Optional(CONF_INITIAL_MODE, default=CONF_VOLTAGE): cv.one_of(*INITIAL_MODES, lower=True),
 }).extend(cv.polling_component_schema('60s'))
 
 
@@ -54,3 +60,4 @@ def to_code(config):
     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]))
+    cg.add(var.set_initial_mode(INITIAL_MODES[config[CONF_INITIAL_MODE]]))
diff --git a/esphome/const.py b/esphome/const.py
index b99ab2014b..a53dbfbf78 100644
--- a/esphome/const.py
+++ b/esphome/const.py
@@ -182,6 +182,7 @@ CONF_IIR_FILTER = 'iir_filter'
 CONF_ILLUMINANCE = 'illuminance'
 CONF_INCLUDES = 'includes'
 CONF_INDEX = 'index'
+CONF_INITIAL_MODE = 'initial_mode'
 CONF_INITIAL_VALUE = 'initial_value'
 CONF_INTEGRATION_TIME = 'integration_time'
 CONF_INTENSITY = 'intensity'
diff --git a/tests/test1.yaml b/tests/test1.yaml
index 47b6aaeab5..2e92f817f8 100644
--- a/tests/test1.yaml
+++ b/tests/test1.yaml
@@ -347,6 +347,7 @@ sensor:
     current_resistor: 0.001 ohm
     voltage_divider: 2351
     change_mode_every: 16
+    initial_mode: VOLTAGE
   - platform: total_daily_energy
     power_id: hlw8012_power
     name: "HLW8012 Total Daily Energy"