From 627a9bb08bb1394e0e3cd9c9cfbd989278fb24b9 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@home-assistant.io>
Date: Sat, 24 Jan 2026 14:51:13 -1000
Subject: [PATCH] refactor restore_state_ to eliminate awkward setup delegation

---
 .../template/water_heater/template_water_heater.cpp        | 4 +---
 esphome/components/water_heater/water_heater.cpp           | 7 ++-----
 esphome/components/water_heater/water_heater.h             | 7 +++----
 3 files changed, 6 insertions(+), 12 deletions(-)

diff --git a/esphome/components/template/water_heater/template_water_heater.cpp b/esphome/components/template/water_heater/template_water_heater.cpp
index 65b1ba5478..e89c96ca48 100644
--- a/esphome/components/template/water_heater/template_water_heater.cpp
+++ b/esphome/components/template/water_heater/template_water_heater.cpp
@@ -8,11 +8,9 @@ static const char *const TAG = "template.water_heater";
 TemplateWaterHeater::TemplateWaterHeater() : set_trigger_(new Trigger<>()) {}
 
 void TemplateWaterHeater::setup() {
-  this->water_heater::WaterHeater::setup();
-
   if (this->restore_mode_ == TemplateWaterHeaterRestoreMode::WATER_HEATER_RESTORE ||
       this->restore_mode_ == TemplateWaterHeaterRestoreMode::WATER_HEATER_RESTORE_AND_CALL) {
-    auto restore = this->restore_state();
+    auto restore = this->restore_state_();
 
     if (restore.has_value()) {
       restore->perform();
diff --git a/esphome/components/water_heater/water_heater.cpp b/esphome/components/water_heater/water_heater.cpp
index 7b947057e1..fbb4181209 100644
--- a/esphome/components/water_heater/water_heater.cpp
+++ b/esphome/components/water_heater/water_heater.cpp
@@ -146,10 +146,6 @@ void WaterHeaterCall::validate_() {
   }
 }
 
-void WaterHeater::setup() {
-  this->pref_ = global_preferences->make_preference<SavedWaterHeaterState>(this->get_preference_hash());
-}
-
 void WaterHeater::publish_state() {
   auto traits = this->get_traits();
   ESP_LOGD(TAG,
@@ -188,7 +184,8 @@ void WaterHeater::publish_state() {
   this->pref_.save(&saved);
 }
 
-optional<WaterHeaterCall> WaterHeater::restore_state() {
+optional<WaterHeaterCall> WaterHeater::restore_state_() {
+  this->pref_ = global_preferences->make_preference<SavedWaterHeaterState>(this->get_preference_hash());
   SavedWaterHeaterState recovered{};
   if (!this->pref_.load(&recovered))
     return {};
diff --git a/esphome/components/water_heater/water_heater.h b/esphome/components/water_heater/water_heater.h
index 6a3fca8a6b..84fc46d208 100644
--- a/esphome/components/water_heater/water_heater.h
+++ b/esphome/components/water_heater/water_heater.h
@@ -204,16 +204,15 @@ class WaterHeater : public EntityBase {
 #endif
   virtual void control(const WaterHeaterCall &call) = 0;
 
-  void setup();
-
-  optional<WaterHeaterCall> restore_state();
-
  protected:
   virtual WaterHeaterTraits traits() = 0;
 
   /// Log the traits of this water heater for dump_config().
   void dump_traits_(const char *tag);
 
+  /// Restore the state of the water heater, call this from your setup() method.
+  optional<WaterHeaterCall> restore_state_();
+
   /// Set the mode of the water heater. Should only be called from control().
   void set_mode_(WaterHeaterMode mode) { this->mode_ = mode; }
   /// Set the target temperature of the water heater. Should only be called from control().