Add On/Off and Away mode support to template platform

Implemented support for On/Off and Away modes in the template water heater
platform, including optimistic control and lambda-based state reporting.

Refactored the base 'WaterHeaterCall' to replace the 'state_' bitmask with
'optional<bool>' for 'on' and 'away' fields. This change was necessary to
enable partial (delta) updates. The previous bitmask implementation did not
distinguish between a field being "set to false" and "not set at all,"
causing unintended state resets (e.g., turning the device off when only
adjusting temperature).

esphome/components/template/water_heater/__init__.py

@@ -3,8 +3,10 @@ import esphome.codegen as cg
 from esphome.components import water_heater
 import esphome.config_validation as cv
 from esphome.const import (
+    CONF_AWAY,
     CONF_ID,
     CONF_MODE,
+    CONF_ON,
     CONF_OPTIMISTIC,
     CONF_RESTORE_MODE,
     CONF_SET_ACTION,
@@ -48,6 +50,8 @@ CONFIG_SCHEMA = (
             cv.Optional(CONF_CURRENT_TEMPERATURE): cv.returning_lambda,
             cv.Optional(CONF_TARGET_TEMPERATURE): cv.returning_lambda,
             cv.Optional(CONF_MODE): cv.returning_lambda,
+            cv.Optional(CONF_ON): cv.returning_lambda,
+            cv.Optional(CONF_AWAY): cv.returning_lambda,
             cv.Optional(CONF_SUPPORTED_MODES): cv.ensure_list(
                 water_heater.validate_water_heater_mode
             ),
@@ -95,6 +99,22 @@ async def to_code(config: ConfigType) -> None:
         )
         cg.add(var.set_mode_lambda(template_))
 
+    if CONF_ON in config:
+        template_ = await cg.process_lambda(
+            config[CONF_ON],
+            [],
+            return_type=cg.optional.template(cg.bool_),
+        )
+        cg.add(var.set_on_lambda(template_))
+
+    if CONF_AWAY in config:
+        template_ = await cg.process_lambda(
+            config[CONF_AWAY],
+            [],
+            return_type=cg.optional.template(cg.bool_),
+        )
+        cg.add(var.set_away_lambda(template_))
+
     if CONF_SUPPORTED_MODES in config:
         cg.add(var.set_supported_modes(config[CONF_SUPPORTED_MODES]))
 
@@ -110,6 +130,8 @@ async def to_code(config: ConfigType) -> None:
             cv.Optional(CONF_MODE): cv.templatable(
                 water_heater.validate_water_heater_mode
             ),
+            cv.Optional(CONF_ON): cv.templatable(cv.boolean),
+            cv.Optional(CONF_AWAY): cv.templatable(cv.boolean),
         }
     ),
 )
@@ -134,4 +156,12 @@ async def water_heater_template_publish_to_code(
         template_ = await cg.templatable(mode, args, water_heater.WaterHeaterMode)
         cg.add(var.set_mode(template_))
 
+    if on := config.get(CONF_ON):
+        template_ = await cg.templatable(on, args, bool)
+        cg.add(var.set_on(template_))
+
+    if away := config.get(CONF_AWAY):
+        template_ = await cg.templatable(away, args, bool)
+        cg.add(var.set_away(template_))
+
     return var

esphome/components/template/water_heater/automation.h

@@ -11,12 +11,15 @@ class TemplateWaterHeaterPublishAction : public Action<Ts...>, public Parented<T
   TEMPLATABLE_VALUE(float, current_temperature)
   TEMPLATABLE_VALUE(float, target_temperature)
   TEMPLATABLE_VALUE(water_heater::WaterHeaterMode, mode)
+  TEMPLATABLE_VALUE(bool, on)
+  TEMPLATABLE_VALUE(bool, away)
 
   void play(const Ts &...x) override {
     if (this->current_temperature_.has_value()) {
       this->parent_->set_current_temperature(this->current_temperature_.value(x...));
     }
-    bool needs_call = this->target_temperature_.has_value() || this->mode_.has_value();
+    bool needs_call = this->target_temperature_.has_value() || this->mode_.has_value() || this->on_.has_value() ||
+                      this->away_.has_value();
     if (needs_call) {
       auto call = this->parent_->make_call();
       if (this->target_temperature_.has_value()) {
@@ -25,6 +28,12 @@ class TemplateWaterHeaterPublishAction : public Action<Ts...>, public Parented<T
       if (this->mode_.has_value()) {
         call.set_mode(this->mode_.value(x...));
       }
+      if (this->on_.has_value()) {
+        call.set_on(this->on_.value(x...));
+      }
+      if (this->away_.has_value()) {
+        call.set_away(this->away_.value(x...));
+      }
       call.perform();
     } else {
       this->parent_->publish_state();

esphome/components/template/water_heater/template_water_heater.cpp

@@ -17,7 +17,7 @@ void TemplateWaterHeater::setup() {
     }
   }
   if (!this->current_temperature_f_.has_value() && !this->target_temperature_f_.has_value() &&
-      !this->mode_f_.has_value())
+      !this->mode_f_.has_value() && !this->on_f_.has_value() && !this->away_f_.has_value())
     this->disable_loop();
 }
 
@@ -32,6 +32,12 @@ water_heater::WaterHeaterTraits TemplateWaterHeater::traits() {
   if (this->target_temperature_f_.has_value()) {
     traits.add_feature_flags(water_heater::WATER_HEATER_SUPPORTS_TARGET_TEMPERATURE);
   }
+  if (this->on_f_.has_value()) {
+    traits.add_feature_flags(water_heater::WATER_HEATER_SUPPORTS_ON_OFF);
+  }
+  if (this->away_f_.has_value()) {
+    traits.add_feature_flags(water_heater::WATER_HEATER_SUPPORTS_AWAY_MODE);
+  }
   return traits;
 }
 
@@ -62,6 +68,22 @@ void TemplateWaterHeater::loop() {
     }
   }
 
+  auto on = this->on_f_.call();
+  if (on.has_value()) {
+    if (*on != this->is_on()) {
+      this->set_state_flag_(water_heater::WATER_HEATER_STATE_ON, *on);
+      changed = true;
+    }
+  }
+
+  auto away = this->away_f_.call();
+  if (away.has_value()) {
+    if (*away != this->is_away()) {
+      this->set_state_flag_(water_heater::WATER_HEATER_STATE_AWAY, *away);
+      changed = true;
+    }
+  }
+
   if (changed) {
     this->publish_state();
   }
@@ -89,6 +111,14 @@ void TemplateWaterHeater::control(const water_heater::WaterHeaterCall &call) {
       this->target_temperature_ = call.get_target_temperature();
     }
   }
+  if (this->optimistic_) {
+    if (call.get_on().has_value()) {
+      this->set_state_flag_(water_heater::WATER_HEATER_STATE_ON, *call.get_on());
+    }
+    if (call.get_away().has_value()) {
+      this->set_state_flag_(water_heater::WATER_HEATER_STATE_AWAY, *call.get_away());
+    }
+  }
 
   this->set_trigger_.trigger();
 

esphome/components/template/water_heater/template_water_heater.h

@@ -24,6 +24,8 @@ class TemplateWaterHeater : public Component, public water_heater::WaterHeater {
     this->target_temperature_f_.set(std::forward<F>(f));
   }
   template<typename F> void set_mode_lambda(F &&f) { this->mode_f_.set(std::forward<F>(f)); }
+  template<typename F> void set_on_lambda(F &&f) { this->on_f_.set(std::forward<F>(f)); }
+  template<typename F> void set_away_lambda(F &&f) { this->away_f_.set(std::forward<F>(f)); }
 
   void set_optimistic(bool optimistic) { this->optimistic_ = optimistic; }
   void set_restore_mode(TemplateWaterHeaterRestoreMode restore_mode) { this->restore_mode_ = restore_mode; }
@@ -49,6 +51,8 @@ class TemplateWaterHeater : public Component, public water_heater::WaterHeater {
   TemplateLambda<float> current_temperature_f_;
   TemplateLambda<float> target_temperature_f_;
   TemplateLambda<water_heater::WaterHeaterMode> mode_f_;
+  TemplateLambda<bool> on_f_;
+  TemplateLambda<bool> away_f_;
   TemplateWaterHeaterRestoreMode restore_mode_{WATER_HEATER_NO_RESTORE};
   water_heater::WaterHeaterModeMask supported_modes_;
   bool optimistic_{true};

esphome/components/water_heater/water_heater.cpp

@@ -60,20 +60,12 @@ WaterHeaterCall &WaterHeaterCall::set_target_temperature_high(float temperature)
 }
 
 WaterHeaterCall &WaterHeaterCall::set_away(bool away) {
-  if (away) {
-    this->state_ |= WATER_HEATER_STATE_AWAY;
-  } else {
-    this->state_ &= ~WATER_HEATER_STATE_AWAY;
-  }
+  this->away_ = away;
   return *this;
 }
 
 WaterHeaterCall &WaterHeaterCall::set_on(bool on) {
-  if (on) {
-    this->state_ |= WATER_HEATER_STATE_ON;
-  } else {
-    this->state_ &= ~WATER_HEATER_STATE_ON;
-  }
+  this->on_ = on;
   return *this;
 }
 
@@ -92,11 +84,11 @@ void WaterHeaterCall::perform() {
   if (!std::isnan(this->target_temperature_high_)) {
     ESP_LOGD(TAG, "  Target Temperature High: %.2f", this->target_temperature_high_);
   }
-  if (this->state_ & WATER_HEATER_STATE_AWAY) {
-    ESP_LOGD(TAG, "  Away: YES");
+  if (this->away_.has_value()) {
+    ESP_LOGD(TAG, "  Away: %s", YESNO(*this->away_));
   }
-  if (this->state_ & WATER_HEATER_STATE_ON) {
-    ESP_LOGD(TAG, "  On: YES");
+  if (this->on_.has_value()) {
+    ESP_LOGD(TAG, "  On: %s", YESNO(*this->on_));
   }
   this->parent_->control(*this);
 }
@@ -137,13 +129,13 @@ void WaterHeaterCall::validate_() {
       this->target_temperature_high_ = NAN;
     }
   }
-  if ((this->state_ & WATER_HEATER_STATE_AWAY) && !traits.get_supports_away_mode()) {
+  if (this->away_.has_value() && *this->away_ && !traits.get_supports_away_mode()) {
     ESP_LOGW(TAG, "'%s' - Away mode not supported", this->parent_->get_name().c_str());
-    this->state_ &= ~WATER_HEATER_STATE_AWAY;
+    this->away_.reset();
   }
   // If ON/OFF not supported, device is always on - clear the flag silently
-  if (!traits.has_feature_flags(WATER_HEATER_SUPPORTS_ON_OFF)) {
-    this->state_ &= ~WATER_HEATER_STATE_ON;
+  if (this->on_.has_value() && !traits.has_feature_flags(WATER_HEATER_SUPPORTS_ON_OFF)) {
+    this->on_.reset();
   }
 }
 

esphome/components/water_heater/water_heater.h

@@ -89,8 +89,8 @@ class WaterHeaterCall {
   float get_target_temperature() const { return this->target_temperature_; }
   float get_target_temperature_low() const { return this->target_temperature_low_; }
   float get_target_temperature_high() const { return this->target_temperature_high_; }
-  /// Get state flags value
-  uint32_t get_state() const { return this->state_; }
+  const optional<bool> &get_away() const { return this->away_; }
+  const optional<bool> &get_on() const { return this->on_; }
 
  protected:
   void validate_();
@@ -99,7 +99,8 @@ class WaterHeaterCall {
   float target_temperature_{NAN};
   float target_temperature_low_{NAN};
   float target_temperature_high_{NAN};
-  uint32_t state_{0};
+  optional<bool> away_;
+  optional<bool> on_;
 };
 
 struct WaterHeaterCallInternal : public WaterHeaterCall {
@@ -110,7 +111,8 @@ struct WaterHeaterCallInternal : public WaterHeaterCall {
     this->target_temperature_ = restore.target_temperature_;
     this->target_temperature_low_ = restore.target_temperature_low_;
     this->target_temperature_high_ = restore.target_temperature_high_;
-    this->state_ = restore.state_;
+    this->away_ = restore.away_;
+    this->on_ = restore.on_;
     return *this;
   }
 };

tests/components/template/common-base.yaml

@@ -414,6 +414,8 @@ water_heater:
     current_temperature: !lambda "return 42.0f;"
     target_temperature: !lambda "return 60.0f;"
     mode: !lambda "return water_heater::WATER_HEATER_MODE_ECO;"
+    on: !lambda "return true;"
+    away: !lambda "return false;"
     supported_modes:
       - "OFF"
       - ECO
@@ -424,6 +426,14 @@ water_heater:
       - PERFORMANCE
     set_action:
       - logger.log: "set_action"
+      - water_heater.template.publish:
+          id: template_water_heater
+          on: false
+          away: true
+      - water_heater.template.publish:
+          id: template_water_heater
+          on: !lambda "return true;"
+          away: !lambda "return false;"
 
 datetime:
   - platform: template

tests/integration/fixtures/water_heater_template.yaml

@@ -11,6 +11,8 @@ water_heater:
     optimistic: true
     current_temperature: !lambda "return 45.0f;"
     target_temperature: !lambda "return 60.0f;"
+    on: !lambda "return true;"
+    away: !lambda "return false;"
     # Note: No mode lambda - we want optimistic mode changes to stick
     # A mode lambda would override mode changes in loop()
     supported_modes:

tests/integration/test_water_heater_template.py

@@ -64,6 +64,12 @@ async def test_water_heater_template(
             f"Expected 4 supported modes, got {len(supported_modes)}: {supported_modes}"
         )
 
+        # Verify supported features
+        # WATER_HEATER_SUPPORTS_AWAY_MODE (1 << 3) = 8
+        # WATER_HEATER_SUPPORTS_ON_OFF (1 << 4) = 16
+        assert test_water_heater.supported_features & 8
+        assert test_water_heater.supported_features & 16
+
         # Subscribe with the wrapper that filters initial states
         client.subscribe_states(initial_state_helper.on_state_wrapper(on_state))
 
@@ -88,6 +94,9 @@ async def test_water_heater_template(
         assert initial_state.target_temperature == 60.0, (
             f"Expected target temp 60.0, got {initial_state.target_temperature}"
         )
+        # Verify On/Away (from lambdas in fixture)
+        assert initial_state.is_on is True
+        assert initial_state.away is False
 
         # Test changing to GAS mode
         client.water_heater_command(test_water_heater.key, mode=WaterHeaterMode.GAS)
@@ -100,6 +109,38 @@ async def test_water_heater_template(
         assert isinstance(gas_state, WaterHeaterState)
         assert gas_state.mode == WaterHeaterMode.GAS
 
+        # Test changing away mode and power (optimistic)
+        away_future: asyncio.Future[WaterHeaterState] = loop.create_future()
+        off_future: asyncio.Future[WaterHeaterState] = loop.create_future()
+
+        def on_state_update(state: aioesphomeapi.EntityState) -> None:
+            if (
+                isinstance(state, WaterHeaterState)
+                and state.key == test_water_heater.key
+            ):
+                if state.away is True and not away_future.done():
+                    away_future.set_result(state)
+                if state.is_on is False and not off_future.done():
+                    off_future.set_result(state)
+
+        client.subscribe_states(on_state_update)
+
+        # Change away mode
+        client.water_heater_command(test_water_heater.key, away=True)
+        try:
+            away_state = await asyncio.wait_for(away_future, timeout=5.0)
+        except TimeoutError:
+            pytest.fail("Away mode change not received within 5 seconds")
+        assert away_state.away is True
+
+        # Change power
+        client.water_heater_command(test_water_heater.key, is_on=False)
+        try:
+            off_state = await asyncio.wait_for(off_future, timeout=5.0)
+        except TimeoutError:
+            pytest.fail("Power off change not received within 5 seconds")
+        assert off_state.is_on is False
+
         # Test changing to ECO mode (from GAS)
         client.water_heater_command(test_water_heater.key, mode=WaterHeaterMode.ECO)