diff --git a/custom_components/plum_ecomax/entity.py b/custom_components/plum_ecomax/entity.py
index 9d5be2b..868dfb5 100644
--- a/custom_components/plum_ecomax/entity.py
+++ b/custom_components/plum_ecomax/entity.py
@@ -25,6 +25,7 @@ async def async_added_to_hass(self):
         """Subscribe to events."""
 
         async def async_set_available(_=None) -> None:
+            """Mark entity as available."""
             self._attr_available = True
 
         func = self.entity_description.filter_fn(self.async_update)
diff --git a/custom_components/plum_ecomax/sensor.py b/custom_components/plum_ecomax/sensor.py
index 60b17fe..6a44864 100644
--- a/custom_components/plum_ecomax/sensor.py
+++ b/custom_components/plum_ecomax/sensor.py
@@ -567,7 +567,8 @@ async def async_update(self, value: dict[str, Any]) -> None:
     async def async_added_to_hass(self):
         """Subscribe to regdata event."""
 
-        async def async_set_available(regdata: dict[str, Any]) -> None:
+        async def async_set_available(regdata: dict[int, Any]) -> None:
+            """Mark entity as available."""
             if self.entity_description.key in regdata:
                 self._attr_available = True