diff --git a/tests/integration/test_water_heater_template.py b/tests/integration/test_water_heater_template.py
index 3b6863766f..acb884c9d9 100644
--- a/tests/integration/test_water_heater_template.py
+++ b/tests/integration/test_water_heater_template.py
@@ -6,28 +6,18 @@ import asyncio
 
 import aioesphomeapi
 from aioesphomeapi import (
+    WaterHeaterFeature,
     WaterHeaterInfo,
     WaterHeaterMode,
     WaterHeaterState,
     WaterHeaterStateFlag,
 )
-from aioesphomeapi.model import APIIntEnum
 import pytest
 
 from .state_utils import InitialStateHelper
 from .types import APIClientConnectedFactory, RunCompiledFunction
 
 
-class WaterHeaterFeature(APIIntEnum):
-    """ESPHome water heater feature flags (WaterHeaterFeature)."""
-
-    SUPPORTS_CURRENT_TEMPERATURE = 1 << 0
-    SUPPORTS_TARGET_TEMPERATURE = 1 << 1
-    SUPPORTS_OPERATION_MODE = 1 << 2
-    SUPPORTS_AWAY_MODE = 1 << 3
-    SUPPORTS_ON_OFF = 1 << 4
-
-
 @pytest.mark.asyncio
 async def test_water_heater_template(
     yaml_config: str,