From 0a7d3c367b9e494b24baf66d5373370718ec9fca Mon Sep 17 00:00:00 2001
From: Clyde Stubbs <2366188+clydebarrow@users.noreply.github.com>
Date: Mon, 27 Nov 2023 08:36:52 +1100
Subject: [PATCH] Add 'enable_at_startup' feature to power_supply (#5826)

---
 esphome/components/power_supply/__init__.py   |  4 +++
 .../components/power_supply/power_supply.cpp  | 29 ++++++++-----------
 .../components/power_supply/power_supply.h    |  3 +-
 tests/test1.1.yaml                            | 13 +++++----
 4 files changed, 25 insertions(+), 24 deletions(-)

diff --git a/esphome/components/power_supply/__init__.py b/esphome/components/power_supply/__init__.py
index f7dd8bca84..6735eddff3 100644
--- a/esphome/components/power_supply/__init__.py
+++ b/esphome/components/power_supply/__init__.py
@@ -8,6 +8,8 @@ power_supply_ns = cg.esphome_ns.namespace("power_supply")
 PowerSupply = power_supply_ns.class_("PowerSupply", cg.Component)
 MULTI_CONF = True
 
+CONF_ENABLE_ON_BOOT = "enable_on_boot"
+
 CONFIG_SCHEMA = cv.Schema(
     {
         cv.Required(CONF_ID): cv.declare_id(PowerSupply),
@@ -18,6 +20,7 @@ CONFIG_SCHEMA = cv.Schema(
         cv.Optional(
             CONF_KEEP_ON_TIME, default="10s"
         ): cv.positive_time_period_milliseconds,
+        cv.Optional(CONF_ENABLE_ON_BOOT, default=False): cv.boolean,
     }
 ).extend(cv.COMPONENT_SCHEMA)
 
@@ -30,5 +33,6 @@ async def to_code(config):
     cg.add(var.set_pin(pin))
     cg.add(var.set_enable_time(config[CONF_ENABLE_TIME]))
     cg.add(var.set_keep_on_time(config[CONF_KEEP_ON_TIME]))
+    cg.add(var.set_enable_on_boot(config[CONF_ENABLE_ON_BOOT]))
 
     cg.add_define("USE_POWER_SUPPLY")
diff --git a/esphome/components/power_supply/power_supply.cpp b/esphome/components/power_supply/power_supply.cpp
index c4d157615a..7474075302 100644
--- a/esphome/components/power_supply/power_supply.cpp
+++ b/esphome/components/power_supply/power_supply.cpp
@@ -11,47 +11,42 @@ void PowerSupply::setup() {
 
   this->pin_->setup();
   this->pin_->digital_write(false);
-  this->enabled_ = false;
+  if (this->enable_on_boot_)
+    this->request_high_power();
 }
 void PowerSupply::dump_config() {
   ESP_LOGCONFIG(TAG, "Power Supply:");
   LOG_PIN("  Pin: ", this->pin_);
   ESP_LOGCONFIG(TAG, "  Time to enable: %" PRIu32 " ms", this->enable_time_);
   ESP_LOGCONFIG(TAG, "  Keep on time: %.1f s", this->keep_on_time_ / 1000.0f);
+  if (this->enable_on_boot_)
+    ESP_LOGCONFIG(TAG, "  Enabled at startup: True");
 }
 
 float PowerSupply::get_setup_priority() const { return setup_priority::IO; }
 
-bool PowerSupply::is_enabled() const { return this->enabled_; }
+bool PowerSupply::is_enabled() const { return this->active_requests_ != 0; }
 
 void PowerSupply::request_high_power() {
-  this->cancel_timeout("power-supply-off");
-  this->pin_->digital_write(true);
-
   if (this->active_requests_ == 0) {
-    // we need to enable the power supply.
-    // cancel old timeout if it exists because we now definitely have a high power mode.
+    this->cancel_timeout("power-supply-off");
     ESP_LOGD(TAG, "Enabling power supply.");
+    this->pin_->digital_write(true);
     delay(this->enable_time_);
   }
-  this->enabled_ = true;
-  // increase active requests
   this->active_requests_++;
 }
 
 void PowerSupply::unrequest_high_power() {
-  this->active_requests_--;
-  if (this->active_requests_ < 0) {
-    // we're just going to use 0 as our new counter.
-    this->active_requests_ = 0;
-  }
-
   if (this->active_requests_ == 0) {
-    // set timeout for power supply off
+    ESP_LOGW(TAG, "Invalid call to unrequest_high_power");
+    return;
+  }
+  this->active_requests_--;
+  if (this->active_requests_ == 0) {
     this->set_timeout("power-supply-off", this->keep_on_time_, [this]() {
       ESP_LOGD(TAG, "Disabling power supply.");
       this->pin_->digital_write(false);
-      this->enabled_ = false;
     });
   }
 }
diff --git a/esphome/components/power_supply/power_supply.h b/esphome/components/power_supply/power_supply.h
index 49d905ba3a..0b06105ae9 100644
--- a/esphome/components/power_supply/power_supply.h
+++ b/esphome/components/power_supply/power_supply.h
@@ -13,6 +13,7 @@ class PowerSupply : public Component {
   void set_pin(GPIOPin *pin) { pin_ = pin; }
   void set_enable_time(uint32_t enable_time) { enable_time_ = enable_time; }
   void set_keep_on_time(uint32_t keep_on_time) { keep_on_time_ = keep_on_time; }
+  void set_enable_on_boot(bool enable_on_boot) { enable_on_boot_ = enable_on_boot; }
 
   /// Is this power supply currently on?
   bool is_enabled() const;
@@ -35,7 +36,7 @@ class PowerSupply : public Component {
 
  protected:
   GPIOPin *pin_;
-  bool enabled_{false};
+  bool enable_on_boot_{false};
   uint32_t enable_time_;
   uint32_t keep_on_time_;
   int16_t active_requests_{0};  // use signed integer to make catching negative requests easier.
diff --git a/tests/test1.1.yaml b/tests/test1.1.yaml
index f4ad89897b..e2e7bd5d63 100644
--- a/tests/test1.1.yaml
+++ b/tests/test1.1.yaml
@@ -44,12 +44,13 @@ network:
 e131:
 
 power_supply:
-  id: atx_power_supply
-  enable_time: 20ms
-  keep_on_time: 10s
-  pin:
-    number: 13
-    inverted: true
+  - id: atx_power_supply
+    enable_time: 20ms
+    keep_on_time: 10s
+    enable_on_boot: true
+    pin:
+      number: 13
+      inverted: true
 
 i2c:
   sda: 21