[safe_mode] Add feature to explicitly mark a boot as successful (#14306)

Co-authored-by: Jonathan Swoboda <154711427+swoboda1337@users.noreply.github.com>

esphome/components/safe_mode/__init__.py

@@ -21,6 +21,7 @@ CONF_ON_SAFE_MODE = "on_safe_mode"
 safe_mode_ns = cg.esphome_ns.namespace("safe_mode")
 SafeModeComponent = safe_mode_ns.class_("SafeModeComponent", cg.Component)
 SafeModeTrigger = safe_mode_ns.class_("SafeModeTrigger", automation.Trigger.template())
+MarkSuccessfulAction = safe_mode_ns.class_("MarkSuccessfulAction", automation.Action)
 
 
 def _remove_id_if_disabled(value):
@@ -53,6 +54,22 @@ CONFIG_SCHEMA = cv.All(
 )
 
 
+@automation.register_action(
+    "safe_mode.mark_successful",
+    MarkSuccessfulAction,
+    cv.Schema(
+        {
+            cv.GenerateID(): cv.use_id(SafeModeComponent),
+        }
+    ),
+)
+async def safe_mode_mark_successful_to_code(config, action_id, template_arg, args):
+    parent = await cg.get_variable(config[CONF_ID])
+    var = cg.new_Pvariable(action_id, template_arg)
+    cg.add(var.set_parent(parent))
+    return var
+
+
 @coroutine_with_priority(CoroPriority.APPLICATION)
 async def to_code(config):
     if not config[CONF_DISABLED]:

esphome/components/safe_mode/automation.h

@@ -1,20 +1,22 @@
 #pragma once
 #include "esphome/core/defines.h"
-
-#ifdef USE_SAFE_MODE_CALLBACK
-#include "safe_mode.h"
-
 #include "esphome/core/automation.h"
+#include "safe_mode.h"
 
 namespace esphome::safe_mode {
 
+#ifdef USE_SAFE_MODE_CALLBACK
 class SafeModeTrigger final : public Trigger<> {
  public:
   explicit SafeModeTrigger(SafeModeComponent *parent) {
     parent->add_on_safe_mode_callback([this]() { trigger(); });
   }
 };
+#endif  // USE_SAFE_MODE_CALLBACK
+
+template<typename... Ts> class MarkSuccessfulAction : public Action<Ts...>, public Parented<SafeModeComponent> {
+ public:
+  void play(const Ts &...x) override { this->parent_->mark_successful(); }
+};
 
 }  // namespace esphome::safe_mode
-
-#endif  // USE_SAFE_MODE_CALLBACK

esphome/components/safe_mode/safe_mode.cpp

@@ -63,18 +63,22 @@ void SafeModeComponent::dump_config() {
 
 float SafeModeComponent::get_setup_priority() const { return setup_priority::AFTER_WIFI; }
 
+void SafeModeComponent::mark_successful() {
+  this->clean_rtc();
+  this->boot_successful_ = true;
+#if defined(USE_ESP32) && defined(USE_OTA_ROLLBACK)
+  // Mark OTA partition as valid to prevent rollback
+  esp_ota_mark_app_valid_cancel_rollback();
+#endif
+  // Disable loop since we no longer need to check
+  this->disable_loop();
+}
+
 void SafeModeComponent::loop() {
   if (!this->boot_successful_ && (millis() - this->safe_mode_start_time_) > this->safe_mode_boot_is_good_after_) {
     // successful boot, reset counter
     ESP_LOGI(TAG, "Boot seems successful; resetting boot loop counter");
-    this->clean_rtc();
-    this->boot_successful_ = true;
-#if defined(USE_ESP32) && defined(USE_OTA_ROLLBACK)
-    // Mark OTA partition as valid to prevent rollback
-    esp_ota_mark_app_valid_cancel_rollback();
-#endif
-    // Disable loop since we no longer need to check
-    this->disable_loop();
+    this->mark_successful();
   }
 }
 

esphome/components/safe_mode/safe_mode.h

@@ -31,6 +31,8 @@ class SafeModeComponent final : public Component {
 
   void on_safe_shutdown() override;
 
+  void mark_successful();
+
 #ifdef USE_SAFE_MODE_CALLBACK
   void add_on_safe_mode_callback(std::function<void()> &&callback) {
     this->safe_mode_callback_.add(std::move(callback));

tests/components/safe_mode/common-enabled.yaml

@@ -16,3 +16,7 @@ button:
 switch:
   - platform: safe_mode
     name: Safe Mode Switch
+
+esphome:
+  on_boot:
+    - safe_mode.mark_successful