diff --git a/esphome/components/api/api.proto b/esphome/components/api/api.proto
index d25934c60b..18dac6a2d1 100644
--- a/esphome/components/api/api.proto
+++ b/esphome/components/api/api.proto
@@ -1155,9 +1155,11 @@ enum WaterHeaterCommandHasField {
   WATER_HEATER_COMMAND_HAS_NONE = 0;
   WATER_HEATER_COMMAND_HAS_MODE = 1;
   WATER_HEATER_COMMAND_HAS_TARGET_TEMPERATURE = 2;
-  WATER_HEATER_COMMAND_HAS_STATE = 4;
+  WATER_HEATER_COMMAND_HAS_STATE = 4 [deprecated=true];
   WATER_HEATER_COMMAND_HAS_TARGET_TEMPERATURE_LOW = 8;
   WATER_HEATER_COMMAND_HAS_TARGET_TEMPERATURE_HIGH = 16;
+  WATER_HEATER_COMMAND_HAS_ON_STATE = 32;
+  WATER_HEATER_COMMAND_HAS_AWAY_STATE = 64;
 }
 
 message WaterHeaterCommandRequest {
diff --git a/esphome/components/api/api_connection.cpp b/esphome/components/api/api_connection.cpp
index 39a4174cdf..14ccbd2248 100644
--- a/esphome/components/api/api_connection.cpp
+++ b/esphome/components/api/api_connection.cpp
@@ -1360,8 +1360,12 @@ void APIConnection::on_water_heater_command_request(const WaterHeaterCommandRequ
     call.set_target_temperature_low(msg.target_temperature_low);
   if (msg.has_fields & enums::WATER_HEATER_COMMAND_HAS_TARGET_TEMPERATURE_HIGH)
     call.set_target_temperature_high(msg.target_temperature_high);
-  if (msg.has_fields & enums::WATER_HEATER_COMMAND_HAS_STATE) {
+  if ((msg.has_fields & enums::WATER_HEATER_COMMAND_HAS_AWAY_STATE) ||
+      (msg.has_fields & enums::WATER_HEATER_COMMAND_HAS_STATE)) {
     call.set_away((msg.state & water_heater::WATER_HEATER_STATE_AWAY) != 0);
+  }
+  if ((msg.has_fields & enums::WATER_HEATER_COMMAND_HAS_ON_STATE) ||
+      (msg.has_fields & enums::WATER_HEATER_COMMAND_HAS_STATE)) {
     call.set_on((msg.state & water_heater::WATER_HEATER_STATE_ON) != 0);
   }
   call.perform();
diff --git a/esphome/components/api/api_pb2.h b/esphome/components/api/api_pb2.h
index 15819da172..d001f869c5 100644
--- a/esphome/components/api/api_pb2.h
+++ b/esphome/components/api/api_pb2.h
@@ -147,6 +147,8 @@ enum WaterHeaterCommandHasField : uint32_t {
   WATER_HEATER_COMMAND_HAS_STATE = 4,
   WATER_HEATER_COMMAND_HAS_TARGET_TEMPERATURE_LOW = 8,
   WATER_HEATER_COMMAND_HAS_TARGET_TEMPERATURE_HIGH = 16,
+  WATER_HEATER_COMMAND_HAS_ON_STATE = 32,
+  WATER_HEATER_COMMAND_HAS_AWAY_STATE = 64,
 };
 #ifdef USE_NUMBER
 enum NumberMode : uint32_t {
diff --git a/esphome/components/api/api_pb2_dump.cpp b/esphome/components/api/api_pb2_dump.cpp
index f1e3bdcafe..73690610ed 100644
--- a/esphome/components/api/api_pb2_dump.cpp
+++ b/esphome/components/api/api_pb2_dump.cpp
@@ -385,6 +385,10 @@ const char *proto_enum_to_string<enums::WaterHeaterCommandHasField>(enums::Water
       return "WATER_HEATER_COMMAND_HAS_TARGET_TEMPERATURE_LOW";
     case enums::WATER_HEATER_COMMAND_HAS_TARGET_TEMPERATURE_HIGH:
       return "WATER_HEATER_COMMAND_HAS_TARGET_TEMPERATURE_HIGH";
+    case enums::WATER_HEATER_COMMAND_HAS_ON_STATE:
+      return "WATER_HEATER_COMMAND_HAS_ON_STATE";
+    case enums::WATER_HEATER_COMMAND_HAS_AWAY_STATE:
+      return "WATER_HEATER_COMMAND_HAS_AWAY_STATE";
     default:
       return "UNKNOWN";
   }
diff --git a/esphome/components/ld2450/__init__.py b/esphome/components/ld2450/__init__.py
index bd6d697c90..5854a5794c 100644
--- a/esphome/components/ld2450/__init__.py
+++ b/esphome/components/ld2450/__init__.py
@@ -1,7 +1,8 @@
+from esphome import automation
 import esphome.codegen as cg
 from esphome.components import uart
 import esphome.config_validation as cv
-from esphome.const import CONF_ID, CONF_THROTTLE
+from esphome.const import CONF_ID, CONF_ON_DATA, CONF_THROTTLE, CONF_TRIGGER_ID
 
 AUTO_LOAD = ["ld24xx"]
 DEPENDENCIES = ["uart"]
@@ -11,6 +12,8 @@ MULTI_CONF = True
 ld2450_ns = cg.esphome_ns.namespace("ld2450")
 LD2450Component = ld2450_ns.class_("LD2450Component", cg.Component, uart.UARTDevice)
 
+LD2450DataTrigger = ld2450_ns.class_("LD2450DataTrigger", automation.Trigger.template())
+
 CONF_LD2450_ID = "ld2450_id"
 
 CONFIG_SCHEMA = cv.All(
@@ -20,6 +23,11 @@ CONFIG_SCHEMA = cv.All(
             cv.Optional(CONF_THROTTLE): cv.invalid(
                 f"{CONF_THROTTLE} has been removed; use per-sensor filters, instead"
             ),
+            cv.Optional(CONF_ON_DATA): automation.validate_automation(
+                {
+                    cv.GenerateID(CONF_TRIGGER_ID): cv.declare_id(LD2450DataTrigger),
+                }
+            ),
         }
     )
     .extend(uart.UART_DEVICE_SCHEMA)
@@ -45,3 +53,6 @@ async def to_code(config):
     var = cg.new_Pvariable(config[CONF_ID])
     await cg.register_component(var, config)
     await uart.register_uart_device(var, config)
+    for conf in config.get(CONF_ON_DATA, []):
+        trigger = cg.new_Pvariable(conf[CONF_TRIGGER_ID], var)
+        await automation.build_automation(trigger, [], conf)
diff --git a/esphome/components/ld2450/ld2450.cpp b/esphome/components/ld2450/ld2450.cpp
index b04b509a16..1ea5c18271 100644
--- a/esphome/components/ld2450/ld2450.cpp
+++ b/esphome/components/ld2450/ld2450.cpp
@@ -413,6 +413,10 @@ void LD2450Component::restart_and_read_all_info() {
   this->set_timeout(1500, [this]() { this->read_all_info(); });
 }
 
+void LD2450Component::add_on_data_callback(std::function<void()> &&callback) {
+  this->data_callback_.add(std::move(callback));
+}
+
 // Send command with values to LD2450
 void LD2450Component::send_command_(uint8_t command, const uint8_t *command_value, uint8_t command_value_len) {
   ESP_LOGV(TAG, "Sending COMMAND %02X", command);
@@ -613,6 +617,8 @@ void LD2450Component::handle_periodic_data_() {
     this->still_presence_millis_ = App.get_loop_component_start_time();
   }
 #endif
+
+  this->data_callback_.call();
 }
 
 bool LD2450Component::handle_ack_data_() {
diff --git a/esphome/components/ld2450/ld2450.h b/esphome/components/ld2450/ld2450.h
index b94c3cac37..fe69cd81d0 100644
--- a/esphome/components/ld2450/ld2450.h
+++ b/esphome/components/ld2450/ld2450.h
@@ -141,6 +141,9 @@ class LD2450Component : public Component, public uart::UARTDevice {
                       int32_t zone2_x1, int32_t zone2_y1, int32_t zone2_x2, int32_t zone2_y2, int32_t zone3_x1,
                       int32_t zone3_y1, int32_t zone3_x2, int32_t zone3_y2);
 
+  /// Add a callback that will be called after each successfully processed periodic data frame.
+  void add_on_data_callback(std::function<void()> &&callback);
+
  protected:
   void send_command_(uint8_t command_str, const uint8_t *command_value, uint8_t command_value_len);
   void set_config_mode_(bool enable);
@@ -190,6 +193,15 @@ class LD2450Component : public Component, public uart::UARTDevice {
 #ifdef USE_TEXT_SENSOR
   std::array<text_sensor::TextSensor *, 3> direction_text_sensors_{};
 #endif
+
+  LazyCallbackManager<void()> data_callback_;
+};
+
+class LD2450DataTrigger : public Trigger<> {
+ public:
+  explicit LD2450DataTrigger(LD2450Component *parent) {
+    parent->add_on_data_callback([this]() { this->trigger(); });
+  }
 };
 
 }  // namespace esphome::ld2450
diff --git a/esphome/components/template/water_heater/__init__.py b/esphome/components/template/water_heater/__init__.py
index 5f96155fbf..71f98c826a 100644
--- a/esphome/components/template/water_heater/__init__.py
+++ b/esphome/components/template/water_heater/__init__.py
@@ -3,6 +3,7 @@ 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_OPTIMISTIC,
@@ -18,6 +19,7 @@ from esphome.types import ConfigType
 from .. import template_ns
 
 CONF_CURRENT_TEMPERATURE = "current_temperature"
+CONF_IS_ON = "is_on"
 
 TemplateWaterHeater = template_ns.class_(
     "TemplateWaterHeater", cg.Component, water_heater.WaterHeater
@@ -51,6 +53,8 @@ CONFIG_SCHEMA = (
             cv.Optional(CONF_SUPPORTED_MODES): cv.ensure_list(
                 water_heater.validate_water_heater_mode
             ),
+            cv.Optional(CONF_AWAY): cv.returning_lambda,
+            cv.Optional(CONF_IS_ON): cv.returning_lambda,
         }
     )
     .extend(cv.COMPONENT_SCHEMA)
@@ -98,6 +102,22 @@ async def to_code(config: ConfigType) -> None:
     if CONF_SUPPORTED_MODES in config:
         cg.add(var.set_supported_modes(config[CONF_SUPPORTED_MODES]))
 
+    if CONF_AWAY in config:
+        template_ = await cg.process_lambda(
+            config[CONF_AWAY],
+            [],
+            return_type=cg.optional.template(bool),
+        )
+        cg.add(var.set_away_lambda(template_))
+
+    if CONF_IS_ON in config:
+        template_ = await cg.process_lambda(
+            config[CONF_IS_ON],
+            [],
+            return_type=cg.optional.template(bool),
+        )
+        cg.add(var.set_is_on_lambda(template_))
+
 
 @automation.register_action(
     "water_heater.template.publish",
@@ -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_AWAY): cv.templatable(cv.boolean),
+            cv.Optional(CONF_IS_ON): 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 CONF_AWAY in config:
+        template_ = await cg.templatable(config[CONF_AWAY], args, bool)
+        cg.add(var.set_away(template_))
+
+    if CONF_IS_ON in config:
+        template_ = await cg.templatable(config[CONF_IS_ON], args, bool)
+        cg.add(var.set_is_on(template_))
+
     return var
diff --git a/esphome/components/template/water_heater/automation.h b/esphome/components/template/water_heater/automation.h
index 3dad2b85ae..d19542db41 100644
--- a/esphome/components/template/water_heater/automation.h
+++ b/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, away)
+  TEMPLATABLE_VALUE(bool, is_on)
 
   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->away_.has_value() ||
+                      this->is_on_.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->away_.has_value()) {
+        call.set_away(this->away_.value(x...));
+      }
+      if (this->is_on_.has_value()) {
+        call.set_on(this->is_on_.value(x...));
+      }
       call.perform();
     } else {
       this->parent_->publish_state();
diff --git a/esphome/components/template/water_heater/template_water_heater.cpp b/esphome/components/template/water_heater/template_water_heater.cpp
index c354deee0e..57c76286a0 100644
--- a/esphome/components/template/water_heater/template_water_heater.cpp
+++ b/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->away_f_.has_value() && !this->is_on_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->away_f_.has_value()) {
+    traits.set_supports_away_mode(true);
+  }
+  if (this->is_on_f_.has_value()) {
+    traits.add_feature_flags(water_heater::WATER_HEATER_SUPPORTS_ON_OFF);
+  }
   return traits;
 }
 
@@ -62,6 +68,22 @@ void TemplateWaterHeater::loop() {
     }
   }
 
+  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;
+    }
+  }
+
+  auto is_on = this->is_on_f_.call();
+  if (is_on.has_value()) {
+    if (*is_on != this->is_on()) {
+      this->set_state_flag_(water_heater::WATER_HEATER_STATE_ON, *is_on);
+      changed = true;
+    }
+  }
+
   if (changed) {
     this->publish_state();
   }
@@ -90,6 +112,17 @@ void TemplateWaterHeater::control(const water_heater::WaterHeaterCall &call) {
     }
   }
 
+  if (call.get_away().has_value()) {
+    if (this->optimistic_) {
+      this->set_state_flag_(water_heater::WATER_HEATER_STATE_AWAY, *call.get_away());
+    }
+  }
+  if (call.get_on().has_value()) {
+    if (this->optimistic_) {
+      this->set_state_flag_(water_heater::WATER_HEATER_STATE_ON, *call.get_on());
+    }
+  }
+
   this->set_trigger_.trigger();
 
   if (this->optimistic_) {
diff --git a/esphome/components/template/water_heater/template_water_heater.h b/esphome/components/template/water_heater/template_water_heater.h
index 22173209aa..045a142e40 100644
--- a/esphome/components/template/water_heater/template_water_heater.h
+++ b/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_away_lambda(F &&f) { this->away_f_.set(std::forward<F>(f)); }
+  template<typename F> void set_is_on_lambda(F &&f) { this->is_on_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> away_f_;
+  TemplateLambda<bool> is_on_f_;
   TemplateWaterHeaterRestoreMode restore_mode_{WATER_HEATER_NO_RESTORE};
   water_heater::WaterHeaterModeMask supported_modes_;
   bool optimistic_{true};
diff --git a/requirements.txt b/requirements.txt
index 1771867535..b8adc22013 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -12,7 +12,7 @@ platformio==6.1.19
 esptool==5.1.0
 click==8.1.7
 esphome-dashboard==20260110.0
-aioesphomeapi==43.14.0
+aioesphomeapi==44.0.0
 zeroconf==0.148.0
 puremagic==1.30
 ruamel.yaml==0.19.1 # dashboard_import
diff --git a/tests/components/ld2450/common.yaml b/tests/components/ld2450/common.yaml
index cfa3c922fc..617228ca34 100644
--- a/tests/components/ld2450/common.yaml
+++ b/tests/components/ld2450/common.yaml
@@ -1,5 +1,8 @@
 ld2450:
   - id: ld2450_radar
+    on_data:
+      then:
+        - logger.log: "LD2450 Radar Data Received"
 
 button:
   - platform: ld2450
diff --git a/tests/components/template/common-base.yaml b/tests/components/template/common-base.yaml
index b8742f8c7b..e9ddfcf43e 100644
--- a/tests/components/template/common-base.yaml
+++ b/tests/components/template/common-base.yaml
@@ -13,6 +13,8 @@ esphome:
         id: template_water_heater
         target_temperature: 50.0
         mode: ECO
+        away: false
+        is_on: true
 
     # Templated
     - water_heater.template.publish:
@@ -20,6 +22,8 @@ esphome:
         current_temperature: !lambda "return 45.0;"
         target_temperature: !lambda "return 55.0;"
         mode: !lambda "return water_heater::WATER_HEATER_MODE_GAS;"
+        away: !lambda "return true;"
+        is_on: !lambda "return false;"
 
     # Test C++ API: set_template() with stateless lambda (no captures)
     # NOTE: set_template() is not intended to be a public API, but we test it to ensure it doesn't break.
@@ -414,6 +418,8 @@ water_heater:
     current_temperature: !lambda "return 42.0f;"
     target_temperature: !lambda "return 60.0f;"
     mode: !lambda "return water_heater::WATER_HEATER_MODE_ECO;"
+    away: !lambda "return false;"
+    is_on: !lambda "return true;"
     supported_modes:
       - "OFF"
       - ECO
diff --git a/tests/integration/fixtures/water_heater_template.yaml b/tests/integration/fixtures/water_heater_template.yaml
index 1aaded1991..0c82ff68ce 100644
--- a/tests/integration/fixtures/water_heater_template.yaml
+++ b/tests/integration/fixtures/water_heater_template.yaml
@@ -4,6 +4,14 @@ host:
 api:
 logger:
 
+globals:
+  - id: global_away
+    type: bool
+    initial_value: "false"
+  - id: global_is_on
+    type: bool
+    initial_value: "true"
+
 water_heater:
   - platform: template
     id: test_boiler
@@ -11,6 +19,8 @@ water_heater:
     optimistic: true
     current_temperature: !lambda "return 45.0f;"
     target_temperature: !lambda "return 60.0f;"
+    away: !lambda "return id(global_away);"
+    is_on: !lambda "return id(global_is_on);"
     # Note: No mode lambda - we want optimistic mode changes to stick
     # A mode lambda would override mode changes in loop()
     supported_modes:
@@ -22,3 +32,8 @@ water_heater:
       min_temperature: 30.0
       max_temperature: 85.0
       target_temperature_step: 0.5
+    set_action:
+      - lambda: |-
+          // Sync optimistic state back to globals so lambdas reflect the change
+          id(global_away) = id(test_boiler).is_away();
+          id(global_is_on) = id(test_boiler).is_on();
diff --git a/tests/integration/test_water_heater_template.py b/tests/integration/test_water_heater_template.py
index 6b4a685d0d..096d4c8461 100644
--- a/tests/integration/test_water_heater_template.py
+++ b/tests/integration/test_water_heater_template.py
@@ -5,7 +5,13 @@ from __future__ import annotations
 import asyncio
 
 import aioesphomeapi
-from aioesphomeapi import WaterHeaterInfo, WaterHeaterMode, WaterHeaterState
+from aioesphomeapi import (
+    WaterHeaterFeature,
+    WaterHeaterInfo,
+    WaterHeaterMode,
+    WaterHeaterState,
+    WaterHeaterStateFlag,
+)
 import pytest
 
 from .state_utils import InitialStateHelper
@@ -22,18 +28,25 @@ async def test_water_heater_template(
     loop = asyncio.get_running_loop()
     async with run_compiled(yaml_config), api_client_connected() as client:
         states: dict[int, aioesphomeapi.EntityState] = {}
-        gas_mode_future: asyncio.Future[WaterHeaterState] = loop.create_future()
-        eco_mode_future: asyncio.Future[WaterHeaterState] = loop.create_future()
+        state_future: asyncio.Future[WaterHeaterState] | None = None
 
         def on_state(state: aioesphomeapi.EntityState) -> None:
             states[state.key] = state
-            if isinstance(state, WaterHeaterState):
-                # Wait for GAS mode
-                if state.mode == WaterHeaterMode.GAS and not gas_mode_future.done():
-                    gas_mode_future.set_result(state)
-                # Wait for ECO mode (we start at OFF, so test transitioning to ECO)
-                elif state.mode == WaterHeaterMode.ECO and not eco_mode_future.done():
-                    eco_mode_future.set_result(state)
+            if (
+                isinstance(state, WaterHeaterState)
+                and state_future is not None
+                and not state_future.done()
+            ):
+                state_future.set_result(state)
+
+        async def wait_for_state(timeout: float = 5.0) -> WaterHeaterState:
+            """Wait for next water heater state change."""
+            nonlocal state_future
+            state_future = loop.create_future()
+            try:
+                return await asyncio.wait_for(state_future, timeout)
+            finally:
+                state_future = None
 
         # Get entities and set up state synchronization
         entities, services = await client.list_entities_services()
@@ -89,24 +102,52 @@ async def test_water_heater_template(
             f"Expected target temp 60.0, got {initial_state.target_temperature}"
         )
 
+        # Verify supported features: away mode and on/off (fixture has away + is_on lambdas)
+        assert (
+            test_water_heater.supported_features & WaterHeaterFeature.SUPPORTS_AWAY_MODE
+        ) != 0, "Expected SUPPORTS_AWAY_MODE in supported_features"
+        assert (
+            test_water_heater.supported_features & WaterHeaterFeature.SUPPORTS_ON_OFF
+        ) != 0, "Expected SUPPORTS_ON_OFF in supported_features"
+
+        # Verify initial state: on (is_on lambda returns true), not away (away lambda returns false)
+        assert (initial_state.state & WaterHeaterStateFlag.ON) != 0, (
+            "Expected initial state to include ON flag"
+        )
+        assert (initial_state.state & WaterHeaterStateFlag.AWAY) == 0, (
+            "Expected initial state to not include AWAY flag"
+        )
+
+        # Test turning on away mode
+        client.water_heater_command(test_water_heater.key, away=True)
+        away_on_state = await wait_for_state()
+        assert (away_on_state.state & WaterHeaterStateFlag.AWAY) != 0
+        # ON flag should still be set (is_on lambda returns true)
+        assert (away_on_state.state & WaterHeaterStateFlag.ON) != 0
+
+        # Test turning off away mode
+        client.water_heater_command(test_water_heater.key, away=False)
+        away_off_state = await wait_for_state()
+        assert (away_off_state.state & WaterHeaterStateFlag.AWAY) == 0
+        assert (away_off_state.state & WaterHeaterStateFlag.ON) != 0
+
+        # Test turning off (on=False)
+        client.water_heater_command(test_water_heater.key, on=False)
+        off_state = await wait_for_state()
+        assert (off_state.state & WaterHeaterStateFlag.ON) == 0
+        assert (off_state.state & WaterHeaterStateFlag.AWAY) == 0
+
+        # Test turning back on (on=True)
+        client.water_heater_command(test_water_heater.key, on=True)
+        on_state = await wait_for_state()
+        assert (on_state.state & WaterHeaterStateFlag.ON) != 0
+
         # Test changing to GAS mode
         client.water_heater_command(test_water_heater.key, mode=WaterHeaterMode.GAS)
-
-        try:
-            gas_state = await asyncio.wait_for(gas_mode_future, timeout=5.0)
-        except TimeoutError:
-            pytest.fail("GAS mode change not received within 5 seconds")
-
-        assert isinstance(gas_state, WaterHeaterState)
+        gas_state = await wait_for_state()
         assert gas_state.mode == WaterHeaterMode.GAS
 
         # Test changing to ECO mode (from GAS)
         client.water_heater_command(test_water_heater.key, mode=WaterHeaterMode.ECO)
-
-        try:
-            eco_state = await asyncio.wait_for(eco_mode_future, timeout=5.0)
-        except TimeoutError:
-            pytest.fail("ECO mode change not received within 5 seconds")
-
-        assert isinstance(eco_state, WaterHeaterState)
+        eco_state = await wait_for_state()
         assert eco_state.mode == WaterHeaterMode.ECO