diff --git a/custom_components/melcloud_custom/binary_sensor.py b/custom_components/melcloud_custom/binary_sensor.py
index 1891b8c..27f4008 100644
--- a/custom_components/melcloud_custom/binary_sensor.py
+++ b/custom_components/melcloud_custom/binary_sensor.py
@@ -1,4 +1,5 @@
 """Support for MelCloud device binary sensors."""
+
 from __future__ import annotations
 
 from dataclasses import dataclass
@@ -12,6 +13,9 @@
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import MelCloudDevice
@@ -46,7 +50,9 @@ class MelcloudBinarySensorEntityDescription(
 _LOGGER = logging.getLogger(__name__)
 
 
-async def async_setup_entry(hass, entry, async_add_entities):
+async def async_setup_entry(
+    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+):
     """Set up MELCloud device binary sensors based on config_entry."""
     entry_config = hass.data[DOMAIN][entry.entry_id]
 
diff --git a/custom_components/melcloud_custom/climate.py b/custom_components/melcloud_custom/climate.py
index 79ff6db..e80f07f 100644
--- a/custom_components/melcloud_custom/climate.py
+++ b/custom_components/melcloud_custom/climate.py
@@ -25,7 +25,8 @@
 )
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_TEMPERATURE, UnitOfTemperature
-from homeassistant.helpers.typing import HomeAssistantType
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import MelCloudDevice
@@ -94,7 +95,7 @@
 
 
 async def async_setup_entry(
-    hass: HomeAssistantType, entry: ConfigEntry, async_add_entities
+    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ):
     """Set up MelCloud device climate based on config_entry."""
     mel_devices = hass.data[DOMAIN][entry.entry_id].get(MEL_DEVICES)
diff --git a/custom_components/melcloud_custom/sensor.py b/custom_components/melcloud_custom/sensor.py
index 10b1396..9eabd92 100644
--- a/custom_components/melcloud_custom/sensor.py
+++ b/custom_components/melcloud_custom/sensor.py
@@ -20,6 +20,9 @@
     UnitOfEnergy,
     UnitOfTemperature,
 )
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import MelCloudDevice
@@ -138,7 +141,9 @@ class MelcloudSensorEntityDescription(
 _LOGGER = logging.getLogger(__name__)
 
 
-async def async_setup_entry(hass, entry, async_add_entities):
+async def async_setup_entry(
+    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+):
     """Set up MELCloud device sensors based on config_entry."""
     entry_config = hass.data[DOMAIN][entry.entry_id]
 
diff --git a/custom_components/melcloud_custom/water_heater.py b/custom_components/melcloud_custom/water_heater.py
index aacda60..c412bf3 100644
--- a/custom_components/melcloud_custom/water_heater.py
+++ b/custom_components/melcloud_custom/water_heater.py
@@ -1,4 +1,5 @@
 """Platform for water_heater integration."""
+
 from __future__ import annotations
 
 from typing import Any
@@ -45,6 +46,7 @@ class AtwWaterHeater(CoordinatorEntity, WaterHeaterEntity):
 
     _attr_supported_features = (
         WaterHeaterEntityFeature.TARGET_TEMPERATURE
+        | WaterHeaterEntityFeature.ON_OFF
         | WaterHeaterEntityFeature.OPERATION_MODE
     )
     _attr_temperature_unit = UnitOfTemperature.CELSIUS
@@ -59,11 +61,11 @@ def __init__(self, api: MelCloudDevice, device: AtwDevice) -> None:
         self._attr_unique_id = f"{device.serial}-WH"
         self._attr_device_info = api.device_info
 
-    async def async_turn_on(self) -> None:
+    async def async_turn_on(self, **kwargs: Any) -> None:
         """Turn the entity on."""
         await self._api.async_set({PROPERTY_POWER: True})
 
-    async def async_turn_off(self) -> None:
+    async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn the entity off."""
         await self._api.async_set({PROPERTY_POWER: False})