actually fix tests?

agittins · Nov 2, 2024 · bcbfad3 · bcbfad3
1 parent d2e2aa4
commit bcbfad3
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 10 deletions.
diff --git a/custom_components/bermuda/__init__.py b/custom_components/bermuda/__init__.py
@@ -87,7 +87,6 @@ async def async_unload_entry(hass: HomeAssistant, entry: BermudaConfigEntry) ->
     """Handle removal of an entry."""
     if unload_result := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
         _LOGGER.debug("Unloaded platforms.")
-        hass.data[DOMAIN].pop(entry.entry_id)
     return unload_result
 
 

diff --git a/custom_components/bermuda/diagnostics.py b/custom_components/bermuda/diagnostics.py
@@ -9,14 +9,13 @@
 from .const import DOMAIN
 
 if TYPE_CHECKING:
-    from homeassistant.config_entries import ConfigEntry
-
+    from . import BermudaConfigEntry
     from .coordinator import BermudaDataUpdateCoordinator
 
 
-async def async_get_config_entry_diagnostics(hass: HomeAssistant, entry: ConfigEntry) -> dict[str, Any]:
+async def async_get_config_entry_diagnostics(hass: HomeAssistant, entry: BermudaConfigEntry) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: BermudaDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator: BermudaDataUpdateCoordinator = entry.runtime_data.coordintor
 
     # We can call this with our own config_entry because the diags step doesn't
     # actually use it.

diff --git a/custom_components/bermuda/entity.py b/custom_components/bermuda/entity.py
@@ -21,10 +21,8 @@
 )
 
 if TYPE_CHECKING:
-    from homeassistant.config_entries import ConfigEntry
-
+    from . import BermudaConfigEntry
     from .coordinator import BermudaDataUpdateCoordinator
-
     # from . import BermudaDevice
 
 
@@ -39,7 +37,7 @@ class BermudaEntity(CoordinatorEntity):
     def __init__(
         self,
         coordinator: BermudaDataUpdateCoordinator,
-        config_entry: ConfigEntry,
+        config_entry: BermudaConfigEntry,
         address: str,
     ) -> None:
         super().__init__(coordinator)
@@ -161,7 +159,7 @@ class BermudaGlobalEntity(CoordinatorEntity):
     def __init__(
         self,
         coordinator: BermudaDataUpdateCoordinator,
-        config_entry: ConfigEntry,
+        config_entry: BermudaConfigEntry,
     ) -> None:
         super().__init__(coordinator)
         self.coordinator = coordinator