diff --git a/custom_components/ohme/__init__.py b/custom_components/ohme/__init__.py
index f21b8db..c60ada8 100644
--- a/custom_components/ohme/__init__.py
+++ b/custom_components/ohme/__init__.py
@@ -1,8 +1,17 @@
+"""The ohme integration."""
+
 import logging
 from homeassistant import core
 from homeassistant.helpers.entity_registry import RegistryEntry, async_migrate_entries
-from .const import *
-from .utils import get_option
+from .const import (
+    DOMAIN,
+    CONFIG_VERSION,
+    ENTITY_TYPES,
+    DATA_CLIENT,
+    DATA_COORDINATORS,
+    DATA_OPTIONS,
+    LEGACY_MAPPING,
+)
 from ohme import OhmeApiClient
 from .coordinator import (
     OhmeChargeSessionsCoordinator,
@@ -21,7 +30,7 @@ async def async_setup(hass: core.HomeAssistant, config: dict) -> bool:
 
 
 async def async_setup_dependencies(hass, entry):
-    """Instantiate client and refresh session"""
+    """Instantiate client and refresh session."""
     client = OhmeApiClient(entry.data["email"], entry.data["password"])
     account_id = entry.data["email"]
 
@@ -41,7 +50,7 @@ async def async_update_listener(hass, entry):
 
 
 async def async_setup_entry(hass, entry):
-    """This is called from the config flow."""
+    """Called from the config flow."""
 
     def _update_unique_id(entry: RegistryEntry) -> dict[str, str] | None:
         """Update unique IDs from old format."""
@@ -49,10 +58,7 @@ def _update_unique_id(entry: RegistryEntry) -> dict[str, str] | None:
             parts = entry.unique_id.split("_")
             legacy_id = "_".join(parts[2:])
 
-            if legacy_id in LEGACY_MAPPING:
-                new_id = LEGACY_MAPPING[legacy_id]
-            else:
-                new_id = legacy_id
+            new_id = LEGACY_MAPPING.get(legacy_id, legacy_id)
 
             new_id = f"{parts[1]}_{new_id}"
 
@@ -90,7 +96,7 @@ def _update_unique_id(entry: RegistryEntry) -> dict[str, str] | None:
         # Catch failures if this is an 'optional' coordinator
         try:
             await coordinator.async_config_entry_first_refresh()
-        except ConfigEntryNotReady as ex:
+        except ConfigEntryNotReady:
             allow_failure = False
             for optional in coordinators_optional:
                 allow_failure = (
@@ -99,10 +105,11 @@ def _update_unique_id(entry: RegistryEntry) -> dict[str, str] | None:
 
             if allow_failure:
                 _LOGGER.error(
-                    f"{coordinator.__class__.__name__} failed to setup. This coordinator is optional so the integration will still function, but please raise an issue if this persists."
+                    "%s failed to setup. This coordinator is optional so the integration will still function, but please raise an issue if this persists.",
+                    coordinator.__class__.__name__,
                 )
             else:
-                raise ex
+                raise
 
     hass.data[DOMAIN][account_id][DATA_COORDINATORS] = coordinators
 
diff --git a/custom_components/ohme/base.py b/custom_components/ohme/base.py
index 65fe7bd..c8ac9c8 100644
--- a/custom_components/ohme/base.py
+++ b/custom_components/ohme/base.py
@@ -1,3 +1,5 @@
+"""Base class for entities."""
+
 from homeassistant.helpers.entity import Entity
 from homeassistant.core import callback
 
diff --git a/custom_components/ohme/binary_sensor.py b/custom_components/ohme/binary_sensor.py
index 183790c..c5238f0 100644
--- a/custom_components/ohme/binary_sensor.py
+++ b/custom_components/ohme/binary_sensor.py
@@ -1,4 +1,4 @@
-"""Platform for sensor integration."""
+"""Platform for binary_sensor."""
 
 from __future__ import annotations
 import logging
@@ -6,9 +6,8 @@
     BinarySensorDeviceClass,
     BinarySensorEntity,
 )
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.entity import generate_entity_id
 from homeassistant.util.dt import utcnow
 from .const import (
     DOMAIN,
@@ -18,6 +17,7 @@
     COORDINATOR_ADVANCED,
     DATA_CLIENT,
 )
+from .coordinator import OhmeChargeSessionsCoordinator
 from .utils import in_slot
 from .base import OhmeEntity
 
@@ -25,8 +25,8 @@
 
 
 async def async_setup_entry(
-    hass: core.HomeAssistant,
-    config_entry: config_entries.ConfigEntry,
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
     async_add_entities,
 ):
     """Setup sensors and configure coordinator."""
@@ -59,6 +59,8 @@ class ConnectedBinarySensor(OhmeEntity, BinarySensorEntity):
 
     @property
     def is_on(self) -> bool:
+        """Calculate state."""
+
         if self.coordinator.data is None:
             self._state = False
         else:
@@ -88,10 +90,13 @@ def __init__(
 
     @property
     def is_on(self) -> bool:
+        """Return state."""
+
         return self._state
 
     def _calculate_state(self) -> bool:
         """Some trickery to get the charge state to update quickly."""
+
         power = self.coordinator.data["power"]["watt"]
 
         # If no last reading or no batterySoc/power, fallback to power > 0
@@ -176,6 +181,7 @@ def _calculate_state(self) -> bool:
     @callback
     def _handle_coordinator_update(self) -> None:
         """Update data."""
+
         # Don't accept updates if 5s hasnt passed
         # State calculations use deltas that may be unreliable to check if requests are too often
         if self._last_updated and (
@@ -229,11 +235,7 @@ class CurrentSlotBinarySensor(OhmeEntity, BinarySensorEntity):
     def extra_state_attributes(self):
         """Attributes of the sensor."""
         now = utcnow()
-        slots = (
-            self._hass.data[DOMAIN][self._client.email][DATA_SLOTS]
-            if DATA_SLOTS in self._hass.data[DOMAIN][self._client.email]
-            else []
-        )
+        slots = self._hass.data[DOMAIN][self._client.email].get(DATA_SLOTS, [])
 
         return {
             "planned_dispatches": [x for x in slots if not x["end"] or x["end"] > now],
@@ -242,6 +244,8 @@ def extra_state_attributes(self):
 
     @property
     def is_on(self) -> bool:
+        """Return state."""
+
         return self._state
 
     @callback
@@ -268,6 +272,8 @@ class ChargerOnlineBinarySensor(OhmeEntity, BinarySensorEntity):
 
     @property
     def is_on(self) -> bool:
+        """Calculate state."""
+
         if self.coordinator.data and self.coordinator.data["online"]:
             return True
         elif self.coordinator.data:
diff --git a/custom_components/ohme/button.py b/custom_components/ohme/button.py
index 789a795..40bb869 100644
--- a/custom_components/ohme/button.py
+++ b/custom_components/ohme/button.py
@@ -1,9 +1,11 @@
+"""Platform for button."""
+
 from __future__ import annotations
 import logging
 import asyncio
 
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.entity import generate_entity_id
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.components.button import ButtonEntity
 
 from .const import DOMAIN, DATA_CLIENT, DATA_COORDINATORS, COORDINATOR_CHARGESESSIONS
@@ -13,7 +15,7 @@
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, config_entry: config_entries.ConfigEntry, async_add_entities
+    hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities
 ):
     """Setup switches."""
     account_id = config_entry.data["email"]
diff --git a/custom_components/ohme/config_flow.py b/custom_components/ohme/config_flow.py
index 347d5cf..eea6296 100644
--- a/custom_components/ohme/config_flow.py
+++ b/custom_components/ohme/config_flow.py
@@ -1,3 +1,5 @@
+"""UI configuration flow."""
+
 import voluptuous as vol
 from homeassistant.config_entries import ConfigFlow, OptionsFlow
 from .const import (
@@ -20,6 +22,8 @@ class OhmeConfigFlow(ConfigFlow, domain=DOMAIN):
     VERSION = CONFIG_VERSION
 
     async def async_step_user(self, info):
+        """First config step."""
+
         errors = {}
 
         if info is not None:
@@ -35,17 +39,21 @@ async def async_step_user(self, info):
             step_id="user", data_schema=USER_SCHEMA, errors=errors
         )
 
-    def async_get_options_flow(entry):
-        return OhmeOptionsFlow(entry)
+    def async_get_options_flow(self):
+        """Return options flow."""
+        return OhmeOptionsFlow(self)
 
 
 class OhmeOptionsFlow(OptionsFlow):
     """Options flow."""
 
     def __init__(self, entry) -> None:
+        """Initialize options flow and store config entry."""
         self._config_entry = entry
 
     async def async_step_init(self, options):
+        """First step of options flow."""
+
         errors = {}
         # If form filled
         if options is not None:
diff --git a/custom_components/ohme/const.py b/custom_components/ohme/const.py
index 9a715a8..ca837f9 100644
--- a/custom_components/ohme/const.py
+++ b/custom_components/ohme/const.py
@@ -1,8 +1,6 @@
-"""Component constants"""
+"""Component constants."""
 
 DOMAIN = "ohme"
-USER_AGENT = "dan-r-homeassistant-ohme"
-INTEGRATION_VERSION = "1.1.0"
 CONFIG_VERSION = 1
 ENTITY_TYPES = ["sensor", "binary_sensor", "switch", "button", "number", "time"]
 
diff --git a/custom_components/ohme/coordinator.py b/custom_components/ohme/coordinator.py
index 02e5ea2..be59562 100644
--- a/custom_components/ohme/coordinator.py
+++ b/custom_components/ohme/coordinator.py
@@ -1,7 +1,10 @@
+"""Ohme coordinators."""
+
 from datetime import timedelta
 import logging
 
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+from ohme import ApiException
 
 from .const import (
     DOMAIN,
@@ -41,8 +44,8 @@ async def _async_update_data(self):
         try:
             return await self._client.async_get_charge_sessions()
 
-        except BaseException:
-            raise UpdateFailed("Error communicating with API")
+        except ApiException as e:
+            raise UpdateFailed("Error communicating with API") from e
 
 
 class OhmeAccountInfoCoordinator(DataUpdateCoordinator):
@@ -70,8 +73,8 @@ async def _async_update_data(self):
         try:
             return await self._client.async_get_account_info()
 
-        except BaseException:
-            raise UpdateFailed("Error communicating with API")
+        except ApiException as e:
+            raise UpdateFailed("Error communicating with API") from e
 
 
 class OhmeAdvancedSettingsCoordinator(DataUpdateCoordinator):
@@ -96,8 +99,8 @@ async def _async_update_data(self):
         try:
             return await self._client.async_get_advanced_settings()
 
-        except BaseException:
-            raise UpdateFailed("Error communicating with API")
+        except ApiException as e:
+            raise UpdateFailed("Error communicating with API") from e
 
 
 class OhmeChargeSchedulesCoordinator(DataUpdateCoordinator):
@@ -122,5 +125,5 @@ async def _async_update_data(self):
         try:
             return await self._client.async_get_schedule()
 
-        except BaseException:
-            raise UpdateFailed("Error communicating with API")
+        except ApiException as e:
+            raise UpdateFailed("Error communicating with API") from e
diff --git a/custom_components/ohme/number.py b/custom_components/ohme/number.py
index dd2b873..436d86d 100644
--- a/custom_components/ohme/number.py
+++ b/custom_components/ohme/number.py
@@ -1,9 +1,11 @@
+"""Platform for number."""
+
 from __future__ import annotations
 import asyncio
 from homeassistant.components.number import NumberEntity, NumberDeviceClass
 from homeassistant.components.number.const import NumberMode, PERCENTAGE
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import UnitOfTime
-from homeassistant.helpers.entity import generate_entity_id
 from homeassistant.core import callback, HomeAssistant
 from .const import (
     DOMAIN,
@@ -18,7 +20,7 @@
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, config_entry: config_entries.ConfigEntry, async_add_entities
+    hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities
 ):
     """Setup switches and configure coordinator."""
     account_id = config_entry.data["email"]
@@ -59,6 +61,7 @@ class TargetPercentNumber(OhmeEntity, NumberEntity):
     _attr_suggested_display_precision = 0
 
     def __init__(self, coordinator, coordinator_schedules, hass: HomeAssistant, client):
+        """Initialise the entity and set up a second coordinator."""
         super().__init__(coordinator, hass, client)
         self.coordinator_schedules = coordinator_schedules
 
@@ -85,7 +88,7 @@ async def async_set_native_value(self, value: float) -> None:
 
     @callback
     def _handle_coordinator_update(self) -> None:
-        """Get value from data returned from API by coordinator"""
+        """Get value from data returned from API by coordinator."""
         # Set with the same logic as reading
         if session_in_progress(self.hass, self._client.email, self.coordinator.data):
             target = round(self.coordinator.data["appliedRule"]["targetPercent"])
@@ -96,6 +99,7 @@ def _handle_coordinator_update(self) -> None:
 
     @property
     def native_value(self):
+        """Return the state of the entity."""
         return self._state
 
 
@@ -111,6 +115,7 @@ class PreconditioningNumber(OhmeEntity, NumberEntity):
     _attr_native_max_value = 60
 
     def __init__(self, coordinator, coordinator_schedules, hass: HomeAssistant, client):
+        """Initialise the entity and set up a second coordinator."""
         super().__init__(coordinator, hass, client)
         self.coordinator_schedules = coordinator_schedules
 
@@ -147,7 +152,7 @@ async def async_set_native_value(self, value: float) -> None:
 
     @callback
     def _handle_coordinator_update(self) -> None:
-        """Get value from data returned from API by coordinator"""
+        """Get value from data returned from API by coordinator."""
         precondition = None
         # Set with the same logic as reading
         if session_in_progress(self.hass, self._client.email, self.coordinator.data):
@@ -175,6 +180,7 @@ def _handle_coordinator_update(self) -> None:
 
     @property
     def native_value(self):
+        """Return the state of the entity."""
         return self._state
 
 
@@ -206,16 +212,17 @@ def native_unit_of_measurement(self):
 
     @callback
     def _handle_coordinator_update(self) -> None:
-        """Get value from data returned from API by coordinator"""
+        """Get value from data returned from API by coordinator."""
         if self.coordinator.data is not None:
             try:
                 self._state = self.coordinator.data["userSettings"]["chargeSettings"][
                     0
                 ]["value"]
-            except:
+            except (IndexError, KeyError):
                 self._state = None
         self.async_write_ha_state()
 
     @property
     def native_value(self):
+        """Return the state of the entity."""
         return self._state
diff --git a/custom_components/ohme/sensor.py b/custom_components/ohme/sensor.py
index 694551e..efd3c67 100644
--- a/custom_components/ohme/sensor.py
+++ b/custom_components/ohme/sensor.py
@@ -8,7 +8,7 @@
 )
 import math
 import logging
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     UnitOfPower,
     UnitOfEnergy,
@@ -17,7 +17,6 @@
     PERCENTAGE,
 )
 from homeassistant.core import HomeAssistant, callback
-from homeassistant.helpers.entity import generate_entity_id
 from homeassistant.util.dt import utcnow
 from .const import (
     DOMAIN,
@@ -27,16 +26,15 @@
     COORDINATOR_CHARGESESSIONS,
     COORDINATOR_ADVANCED,
 )
-from .coordinator import OhmeChargeSessionsCoordinator, OhmeAdvancedSettingsCoordinator
-from .utils import next_slot, get_option, slot_list, slot_list_str
+from .utils import next_slot, slot_list, slot_list_str
 from .base import OhmeEntity
 
 _LOGGER = logging.getLogger(__name__)
 
 
 async def async_setup_entry(
-    hass: core.HomeAssistant,
-    config_entry: config_entries.ConfigEntry,
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
     async_add_entities,
 ):
     """Setup sensors and configure coordinator."""
@@ -73,7 +71,7 @@ class PowerDrawSensor(OhmeEntity, SensorEntity):
 
     @property
     def native_value(self):
-        """Get value from data returned from API by coordinator"""
+        """Get value from data returned from API by coordinator."""
         if self.coordinator.data and self.coordinator.data["power"]:
             return self.coordinator.data["power"]["watt"]
         return 0
@@ -89,7 +87,7 @@ class CurrentDrawSensor(OhmeEntity, SensorEntity):
 
     @property
     def native_value(self):
-        """Get value from data returned from API by coordinator"""
+        """Get value from data returned from API by coordinator."""
         if self.coordinator.data and self.coordinator.data["power"]:
             return self.coordinator.data["power"]["amp"]
         return 0
@@ -105,7 +103,7 @@ class VoltageSensor(OhmeEntity, SensorEntity):
 
     @property
     def native_value(self):
-        """Get value from data returned from API by coordinator"""
+        """Get value from data returned from API by coordinator."""
         if self.coordinator.data and self.coordinator.data["power"]:
             return self.coordinator.data["power"]["volt"]
         return None
@@ -121,7 +119,7 @@ class CTSensor(OhmeEntity, SensorEntity):
 
     @property
     def native_value(self):
-        """Get value from data returned from API by coordinator"""
+        """Get value from data returned from API by coordinator."""
         return self.coordinator.data["clampAmps"]
 
 
@@ -144,7 +142,7 @@ def _handle_coordinator_update(self) -> None:
             new_state = 0
             try:
                 new_state = self.coordinator.data["chargeGraph"]["now"]["y"]
-            except BaseException:
+            except KeyError:
                 _LOGGER.debug(
                     "EnergyUsageSensor: ChargeGraph reading failed, falling back to batterySoc"
                 )
@@ -172,6 +170,7 @@ def _handle_coordinator_update(self) -> None:
 
     @property
     def native_value(self):
+        """Return the state of the entity."""
         return self._state
 
 
@@ -289,7 +288,7 @@ def icon(self):
 
     @callback
     def _handle_coordinator_update(self) -> None:
-        """Get value from data returned from API by coordinator"""
+        """Get value from data returned from API by coordinator."""
         if (
             self.coordinator.data
             and self.coordinator.data["car"]
@@ -309,4 +308,5 @@ def _handle_coordinator_update(self) -> None:
 
     @property
     def native_value(self):
+        """Return the state of the entity."""
         return self._state
diff --git a/custom_components/ohme/switch.py b/custom_components/ohme/switch.py
index e2c9855..04cec82 100644
--- a/custom_components/ohme/switch.py
+++ b/custom_components/ohme/switch.py
@@ -1,11 +1,11 @@
+"""Platform for switch integration."""
+
 from __future__ import annotations
 import logging
 import asyncio
 
 from homeassistant.core import callback, HomeAssistant
-from homeassistant.helpers.entity import generate_entity_id
-
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.components.switch import SwitchEntity
 from homeassistant.util.dt import utcnow
 
@@ -22,7 +22,7 @@
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, config_entry: config_entries.ConfigEntry, async_add_entities
+    hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities
 ):
     """Setup switches and configure coordinator."""
     account_id = config_entry.data["email"]
@@ -89,6 +89,7 @@ class OhmePauseChargeSwitch(OhmeEntity, SwitchEntity):
     @callback
     def _handle_coordinator_update(self) -> None:
         """Determine if charge is paused.
+
         We handle this differently to the sensors as the state of this switch
         is evaluated only when new data is fetched to stop the switch flicking back then forth."""
         if self.coordinator.data is None:
@@ -144,6 +145,7 @@ async def async_turn_on(self):
 
     async def async_turn_off(self):
         """Stop max charging.
+
         We are not changing anything, just applying the last rule. No need to supply anything."""
         await self._client.async_max_charge(False)
 
@@ -163,6 +165,7 @@ def __init__(
         icon,
         config_key,
     ):
+        """Initialise switch."""
         self._attr_icon = f"mdi:{icon}"
         self._attr_translation_key = translation_key
         self._config_key = config_key
diff --git a/custom_components/ohme/time.py b/custom_components/ohme/time.py
index 105aa9f..b783cd9 100644
--- a/custom_components/ohme/time.py
+++ b/custom_components/ohme/time.py
@@ -1,8 +1,10 @@
+"""Platform for time integration."""
+
 from __future__ import annotations
 import asyncio
 import logging
 from homeassistant.components.time import TimeEntity
-from homeassistant.helpers.entity import generate_entity_id
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import callback, HomeAssistant
 from .const import (
     DOMAIN,
@@ -19,7 +21,7 @@
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, config_entry: config_entries.ConfigEntry, async_add_entities
+    hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities
 ):
     """Setup switches and configure coordinator."""
     account_id = config_entry.data["email"]
@@ -47,6 +49,7 @@ class TargetTime(OhmeEntity, TimeEntity):
     _attr_icon = "mdi:alarm-check"
 
     def __init__(self, coordinator, coordinator_schedules, hass: HomeAssistant, client):
+        """Initialise target time sensor."""
         super().__init__(coordinator, hass, client)
 
         self.coordinator_schedules = coordinator_schedules
@@ -78,7 +81,7 @@ async def async_set_value(self, value: dt_time) -> None:
 
     @callback
     def _handle_coordinator_update(self) -> None:
-        """Get value from data returned from API by coordinator"""
+        """Get value from data returned from API by coordinator."""
         # Read with the same logic as setting
         target = None
         if session_in_progress(self.hass, self._client.email, self.coordinator.data):
@@ -94,4 +97,5 @@ def _handle_coordinator_update(self) -> None:
 
     @property
     def native_value(self):
+        """Return the state of the entity."""
         return self._state
diff --git a/custom_components/ohme/utils.py b/custom_components/ohme/utils.py
index 86fa706..2f368d4 100644
--- a/custom_components/ohme/utils.py
+++ b/custom_components/ohme/utils.py
@@ -1,9 +1,9 @@
+"""Common utility functions."""
+
 from functools import reduce
-from datetime import datetime, timedelta
+import datetime
 from .const import DOMAIN, DATA_OPTIONS
 import pytz
-# import logging
-# _LOGGER = logging.getLogger(__name__)
 
 
 def next_slot(hass, account_id, data):
@@ -17,12 +17,12 @@ def next_slot(hass, account_id, data):
     # Loop through slots
     for slot in slots:
         # Only take the first slot start/end that matches. These are in order.
-        if end is None and slot["end"] > datetime.now().astimezone():
+        if end is None and slot["end"] > datetime.datetime.now().astimezone():
             end = slot["end"]
 
         if (
             start is None
-            and slot["start"] > datetime.now().astimezone()
+            and slot["start"] > datetime.datetime.now().astimezone()
             and slot["start"] != end
         ):
             start = slot["start"]
@@ -52,10 +52,14 @@ def slot_list(data):
     for slot in session_slots:
         slots.append(
             {
-                "start": datetime.utcfromtimestamp(slot["startTimeMs"] / 1000)
+                "start": datetime.datetime.fromtimestamp(
+                    slot["startTimeMs"] / 1000, tz=datetime.UTC
+                )
                 .replace(tzinfo=pytz.utc, microsecond=0)
                 .astimezone(),
-                "end": datetime.utcfromtimestamp(slot["endTimeMs"] / 1000)
+                "end": datetime.datetime.fromtimestamp(
+                    slot["endTimeMs"] / 1000, tz=datetime.UTC
+                )
                 .replace(tzinfo=pytz.utc, microsecond=0)
                 .astimezone(),
                 "charge_in_kwh": -(
@@ -99,15 +103,15 @@ def slot_list_str(hass, account_id, slots):
 
 
 def in_slot(data):
-    """Are we currently in a charge slot?"""
+    """Are we currently in a charge slot."""
     slots = slot_list(data)
 
     # Loop through slots
     for slot in slots:
         # If we are in one
         if (
-            slot["start"] < datetime.now().astimezone()
-            and slot["end"] > datetime.now().astimezone()
+            slot["start"] < datetime.datetime.now().astimezone()
+            and slot["end"] > datetime.datetime.now().astimezone()
         ):
             return True
 
@@ -116,16 +120,17 @@ def in_slot(data):
 
 def time_next_occurs(hour, minute):
     """Find when this time next occurs."""
-    current = datetime.now()
+    current = datetime.datetime.now()
     target = current.replace(hour=hour, minute=minute, second=0, microsecond=0)
-    if target <= datetime.now():
-        target = target + timedelta(days=1)
+    if target <= datetime.datetime.now():
+        target = target + datetime.timedelta(days=1)
 
     return target
 
 
 def session_in_progress(hass, account_id, data):
     """Is there a session in progress?
+
     Used to check if we should update the current session rather than the first schedule."""
     # If config option set, never update session specific schedule
     if get_option(hass, account_id, "never_session_specific"):
diff --git a/tests/conftest.py b/tests/conftest.py
index 5e7d8c2..72601bb 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -11,4 +11,5 @@ def auto_enable_custom_integrations(enable_custom_integrations):
 
 
 def enable_external_sockets():
+    """Enable external sockets for custom integrations."""
     pytest_socket.enable_socket()
diff --git a/tests/test_binary_sensor.py b/tests/test_binary_sensor.py
index c440c10..9a0f2a9 100644
--- a/tests/test_binary_sensor.py
+++ b/tests/test_binary_sensor.py
@@ -1,12 +1,12 @@
+"""Tests for binary sensors."""
+
 import pytest
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import MagicMock
 from homeassistant.core import HomeAssistant
-from homeassistant.util.dt import utcnow
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 from custom_components.ohme.const import (
     DOMAIN,
     DATA_CLIENT,
-    DATA_SLOTS,
     DATA_COORDINATORS,
     COORDINATOR_CHARGESESSIONS,
     COORDINATOR_ADVANCED,
@@ -16,13 +16,13 @@
     ConnectedBinarySensor,
     ChargingBinarySensor,
     PendingApprovalBinarySensor,
-    CurrentSlotBinarySensor,
     ChargerOnlineBinarySensor,
 )
 
 
 @pytest.fixture
 def mock_hass():
+    """Fixture for creating a mock Home Assistant instance."""
     hass = MagicMock(spec=HomeAssistant)
     hass.data = {
         DOMAIN: {
@@ -40,17 +40,20 @@ def mock_hass():
 
 @pytest.fixture
 def mock_coordinator():
+    """Fixture for creating a mock coordinator."""
     return MagicMock(spec=DataUpdateCoordinator)
 
 
 @pytest.fixture
 def mock_client():
+    """Fixture for creating a mock client."""
     mock = MagicMock()
     mock.email = "test_account"
     return mock
 
 
 def test_connected_binary_sensor(mock_hass, mock_coordinator, mock_client):
+    """Test ConnectedBinarySensor."""
     sensor = ConnectedBinarySensor(mock_coordinator, mock_hass, mock_client)
     mock_coordinator.data = {"mode": "CONNECTED"}
     assert sensor.is_on is True
@@ -60,6 +63,7 @@ def test_connected_binary_sensor(mock_hass, mock_coordinator, mock_client):
 
 
 def test_charging_binary_sensor(mock_hass, mock_coordinator, mock_client):
+    """Test ChargingBinarySensor."""
     sensor = ChargingBinarySensor(mock_coordinator, mock_hass, mock_client)
     mock_coordinator.data = {
         "power": {"watt": 100},
@@ -72,6 +76,7 @@ def test_charging_binary_sensor(mock_hass, mock_coordinator, mock_client):
 
 
 def test_pending_approval_binary_sensor(mock_hass, mock_coordinator, mock_client):
+    """Test PendingApprovalBinarySensor."""
     sensor = PendingApprovalBinarySensor(mock_coordinator, mock_hass, mock_client)
     mock_coordinator.data = {"mode": "PENDING_APPROVAL"}
     assert sensor.is_on is True
@@ -81,6 +86,7 @@ def test_pending_approval_binary_sensor(mock_hass, mock_coordinator, mock_client
 
 
 def test_charger_online_binary_sensor(mock_hass, mock_coordinator, mock_client):
+    """Test ChargerOnlineBinarySensor."""
     sensor = ChargerOnlineBinarySensor(mock_coordinator, mock_hass, mock_client)
     mock_coordinator.data = {"online": True}
     assert sensor.is_on is True
diff --git a/tests/test_button.py b/tests/test_button.py
index aed7f59..d046d5b 100644
--- a/tests/test_button.py
+++ b/tests/test_button.py
@@ -1,8 +1,7 @@
+"""Tests for buttons."""
+
 import pytest
-from unittest.mock import AsyncMock, MagicMock, patch
-from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry as er
-from homeassistant.setup import async_setup_component
+from unittest.mock import AsyncMock, MagicMock
 from custom_components.ohme.button import async_setup_entry, OhmeApproveChargeButton
 from custom_components.ohme.const import (
     DOMAIN,
@@ -14,6 +13,7 @@
 
 @pytest.fixture
 def mock_hass():
+    """Fixture for creating a mock Home Assistant instance."""
     hass = MagicMock()
     hass.data = {DOMAIN: {"test_account": {}}}
     return hass
@@ -21,11 +21,13 @@ def mock_hass():
 
 @pytest.fixture
 def mock_config_entry():
+    """Fixture for creating a mock config entry."""
     return AsyncMock(data={"email": "test@example.com"})
 
 
 @pytest.fixture
 def mock_client():
+    """Fixture for creating a mock client."""
     client = AsyncMock()
     client.is_capable.return_value = True
     client.async_approve_charge = AsyncMock()
@@ -34,6 +36,7 @@ def mock_client():
 
 @pytest.fixture
 def mock_coordinator():
+    """Fixture for creating a mock coordinator."""
     coordinator = AsyncMock()
     coordinator.async_refresh = AsyncMock()
     return coordinator
@@ -41,6 +44,7 @@ def mock_coordinator():
 
 @pytest.fixture
 def setup_hass(mock_hass, mock_config_entry, mock_client, mock_coordinator):
+    """Fixture for setting up Home Assistant."""
     mock_hass.data = {
         DOMAIN: {
             "test@example.com": {
@@ -54,6 +58,7 @@ def setup_hass(mock_hass, mock_config_entry, mock_client, mock_coordinator):
 
 @pytest.mark.asyncio
 async def test_async_setup_entry(setup_hass, mock_config_entry):
+    """Test async_setup_entry."""
     async_add_entities = AsyncMock()
     await async_setup_entry(setup_hass, mock_config_entry, async_add_entities)
     assert async_add_entities.call_count == 1
@@ -61,6 +66,7 @@ async def test_async_setup_entry(setup_hass, mock_config_entry):
 
 @pytest.mark.asyncio
 async def test_ohme_approve_charge_button(setup_hass, mock_client, mock_coordinator):
+    """Test OhmeApproveChargeButton."""
     button = OhmeApproveChargeButton(mock_coordinator, setup_hass, mock_client)
     await button.async_press()
     mock_client.async_approve_charge.assert_called_once()
diff --git a/tests/test_config_flow.py b/tests/test_config_flow.py
index cbc8711..ce80e04 100644
--- a/tests/test_config_flow.py
+++ b/tests/test_config_flow.py
@@ -1,12 +1,8 @@
 """Tests for the config flow."""
 
 from unittest import mock
-from homeassistant.const import CONF_ACCESS_TOKEN, CONF_NAME, CONF_PATH
-import pytest
-from pytest_homeassistant_custom_component.common import MockConfigEntry
 
 from custom_components.ohme import config_flow
-from custom_components.ohme.const import DOMAIN
 
 
 async def test_step_account(hass):
diff --git a/tests/test_number.py b/tests/test_number.py
index f8bb713..8ff70fb 100644
--- a/tests/test_number.py
+++ b/tests/test_number.py
@@ -1,6 +1,7 @@
+"""Tests for number entities."""
+
 import pytest
 from unittest.mock import AsyncMock, MagicMock, patch
-from homeassistant.core import HomeAssistant
 from custom_components.ohme.const import (
     DOMAIN,
     DATA_CLIENT,
@@ -20,7 +21,8 @@
 
 @pytest.fixture
 def mock_hass():
-    hass = MagicMock(
+    """Fixture for creating a mock Home Assistant instance."""
+    return MagicMock(
         data={
             DOMAIN: {
                 "test@example.com": {
@@ -35,27 +37,30 @@ def mock_hass():
             }
         }
     )
-    return hass
 
 
 @pytest.fixture
 def mock_config_entry():
+    """Fixture for creating a mock config entry."""
     return AsyncMock(data={"email": "test@example.com"})
 
 
 @pytest.fixture
 def mock_async_add_entities():
+    """Fixture for creating a mock async_add_entities."""
     return AsyncMock()
 
 
 @pytest.mark.asyncio
 async def test_async_setup_entry(mock_hass, mock_config_entry, mock_async_add_entities):
+    """Test async_setup_entry."""
     await async_setup_entry(mock_hass, mock_config_entry, mock_async_add_entities)
     assert mock_async_add_entities.call_count == 1
 
 
 @pytest.mark.asyncio
 async def test_target_percent_number(mock_hass):
+    """Test TargetPercentNumber."""
     coordinator = mock_hass.data[DOMAIN]["test@example.com"][DATA_COORDINATORS][
         COORDINATOR_CHARGESESSIONS
     ]
@@ -75,6 +80,7 @@ async def test_target_percent_number(mock_hass):
 
 @pytest.mark.asyncio
 async def test_preconditioning_number(mock_hass):
+    """Test PreconditioningNumber."""
     coordinator = mock_hass.data[DOMAIN]["test@example.com"][DATA_COORDINATORS][
         COORDINATOR_CHARGESESSIONS
     ]
@@ -96,6 +102,7 @@ async def test_preconditioning_number(mock_hass):
 
 @pytest.mark.asyncio
 async def test_price_cap_number(mock_hass):
+    """Test PriceCapNumber."""
     coordinator = mock_hass.data[DOMAIN]["test@example.com"][DATA_COORDINATORS][
         COORDINATOR_ACCOUNTINFO
     ]
diff --git a/tests/test_sensor.py b/tests/test_sensor.py
index 6271d63..8ca9836 100644
--- a/tests/test_sensor.py
+++ b/tests/test_sensor.py
@@ -1,3 +1,5 @@
+"""Tests for sensor entities."""
+
 import pytest
 from unittest.mock import MagicMock
 from custom_components.ohme.sensor import VoltageSensor
@@ -6,8 +8,7 @@
 @pytest.fixture
 def mock_coordinator():
     """Fixture for creating a mock coordinator."""
-    coordinator = MagicMock()
-    return coordinator
+    return MagicMock()
 
 
 @pytest.fixture
diff --git a/tests/test_switch.py b/tests/test_switch.py
index 062b1ca..28add1d 100644
--- a/tests/test_switch.py
+++ b/tests/test_switch.py
@@ -1,5 +1,7 @@
+"""Tests for switch entities."""
+
 import pytest
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 from custom_components.ohme.const import (
@@ -22,11 +24,13 @@
 
 @pytest.fixture
 def mock_hass():
+    """Fixture for creating a mock Home Assistant instance."""
     return AsyncMock(spec=HomeAssistant)
 
 
 @pytest.fixture
 def mock_client():
+    """Fixture for creating a mock client."""
     client = AsyncMock()
     client.cap_available.return_value = True
     client.solar_capable.return_value = True
@@ -40,16 +44,19 @@ def mock_client():
 
 @pytest.fixture
 def mock_coordinator():
+    """Fixture for creating a mock coordinator."""
     return AsyncMock(spec=DataUpdateCoordinator)
 
 
 @pytest.fixture
 def mock_config_entry():
+    """Fixture for creating a mock config entry."""
     return AsyncMock(data={"email": "test@example.com"})
 
 
 @pytest.fixture
 def setup_hass_data(mock_hass, mock_client, mock_coordinator):
+    """Fixture for setting up Home Assistant data."""
     mock_hass.data = {
         DOMAIN: {
             "test@example.com": {
@@ -65,6 +72,7 @@ def setup_hass_data(mock_hass, mock_client, mock_coordinator):
 
 @pytest.mark.asyncio
 async def test_async_setup_entry(mock_hass, mock_config_entry, setup_hass_data):
+    """Test async_setup_entry."""
     async_add_entities = AsyncMock()
     await async_setup_entry(mock_hass, mock_config_entry, async_add_entities)
     assert async_add_entities.call_count == 1
@@ -73,6 +81,7 @@ async def test_async_setup_entry(mock_hass, mock_config_entry, setup_hass_data):
 
 @pytest.mark.asyncio
 async def test_ohme_pause_charge_switch(mock_hass, mock_client, mock_coordinator):
+    """Test OhmePauseChargeSwitch."""
     switch = OhmePauseChargeSwitch(mock_coordinator, mock_hass, mock_client)
     await switch.async_turn_on()
     mock_client.async_pause_charge.assert_called_once()
@@ -82,6 +91,7 @@ async def test_ohme_pause_charge_switch(mock_hass, mock_client, mock_coordinator
 
 @pytest.mark.asyncio
 async def test_ohme_max_charge_switch(mock_hass, mock_client, mock_coordinator):
+    """Test OhmeMaxChargeSwitch."""
     switch = OhmeMaxChargeSwitch(mock_coordinator, mock_hass, mock_client)
     await switch.async_turn_on()
     mock_client.async_max_charge.assert_called_once_with(True)
@@ -92,6 +102,7 @@ async def test_ohme_max_charge_switch(mock_hass, mock_client, mock_coordinator):
 
 @pytest.mark.asyncio
 async def test_ohme_configuration_switch(mock_hass, mock_client, mock_coordinator):
+    """Test OhmeConfigurationSwitch."""
     switch = OhmeConfigurationSwitch(
         mock_coordinator,
         mock_hass,
@@ -113,6 +124,7 @@ async def test_ohme_configuration_switch(mock_hass, mock_client, mock_coordinato
 
 @pytest.mark.asyncio
 async def test_ohme_solar_boost_switch(mock_hass, mock_client, mock_coordinator):
+    """Test OhmeSolarBoostSwitch."""
     switch = OhmeSolarBoostSwitch(mock_coordinator, mock_hass, mock_client)
     await switch.async_turn_on()
     mock_client.async_set_configuration_value.assert_called_once_with(
@@ -127,6 +139,7 @@ async def test_ohme_solar_boost_switch(mock_hass, mock_client, mock_coordinator)
 
 @pytest.mark.asyncio
 async def test_ohme_price_cap_switch(mock_hass, mock_client, mock_coordinator):
+    """Test OhmePrice"""
     switch = OhmePriceCapSwitch(mock_coordinator, mock_hass, mock_client)
     await switch.async_turn_on()
     mock_client.async_change_price_cap.assert_called_once_with(enabled=True)
diff --git a/tests/test_time.py b/tests/test_time.py
index 713116d..f08a403 100644
--- a/tests/test_time.py
+++ b/tests/test_time.py
@@ -1,8 +1,8 @@
+"""Tests for time entity."""
+
 import pytest
 from datetime import time as dt_time
 from unittest.mock import AsyncMock, MagicMock, patch
-from homeassistant.core import HomeAssistant
-from homeassistant.helpers.entity import Entity
 from custom_components.ohme.time import async_setup_entry, TargetTime
 from custom_components.ohme.const import (
     DOMAIN,
@@ -15,6 +15,7 @@
 
 @pytest.fixture
 def mock_hass():
+    """Fixture for creating a mock Home Assistant instance."""
     hass = MagicMock()
     hass.data = {
         DOMAIN: {
@@ -37,22 +38,26 @@ def mock_hass():
 
 @pytest.fixture
 def mock_config_entry():
+    """Fixture for creating a mock config entry."""
     return AsyncMock(data={"email": "test@example.com"})
 
 
 @pytest.fixture
 def mock_async_add_entities():
+    """Fixture for creating a mock async_add_entities."""
     return AsyncMock()
 
 
 @pytest.mark.asyncio
 async def test_async_setup_entry(mock_hass, mock_config_entry, mock_async_add_entities):
+    """Test async_setup_entry."""
     await async_setup_entry(mock_hass, mock_config_entry, mock_async_add_entities)
     assert mock_async_add_entities.called
 
 
 @pytest.fixture
 def target_time_entity(mock_hass):
+    """Fixture for creating a target time entity."""
     coordinator = mock_hass.data[DOMAIN]["test@example.com"][DATA_COORDINATORS][
         COORDINATOR_CHARGESESSIONS
     ]
@@ -65,6 +70,7 @@ def target_time_entity(mock_hass):
 
 @pytest.mark.asyncio
 async def test_async_added_to_hass(target_time_entity):
+    """Test async_added_to_hass."""
     with patch.object(
         target_time_entity.coordinator_schedules,
         "async_add_listener",
@@ -76,6 +82,7 @@ async def test_async_added_to_hass(target_time_entity):
 
 @pytest.mark.asyncio
 async def test_async_set_value(target_time_entity):
+    """Test async_set_value."""
     with patch("custom_components.ohme.time.session_in_progress", return_value=True):
         await target_time_entity.async_set_value(dt_time(12, 30))
         assert target_time_entity._client.async_apply_session_rule.called
@@ -86,5 +93,6 @@ async def test_async_set_value(target_time_entity):
 
 
 def test_native_value(target_time_entity):
+    """Test native_value."""
     target_time_entity._state = dt_time(12, 30)
     assert target_time_entity.native_value == dt_time(12, 30)
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 3bc9409..58e54e7 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1,7 +1 @@
 """Tests for the utils."""
-
-from unittest import mock
-import random
-from time import time
-
-from custom_components.ohme import utils