291 files

custom_components/spook/ectoplasms/zone/services/create.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING
 
 from homeassistant.components.zone import CREATE_FIELDS, DOMAIN, ZoneStorageCollection
 
+from ....helper_collections import async_get_storage_collection
 from ....services import AbstractSpookAdminService
 
 if TYPE_CHECKING:
@@ -21,15 +22,8 @@ class SpookService(AbstractSpookAdminService):
 
     async def async_handle_service(self, call: ServiceCall) -> None:
         """Handle the service call."""
-        collection: ZoneStorageCollection
-        if DOMAIN in self.hass.data:
-            collection = self.hass.data[DOMAIN]
-        else:
-            # Home zone is set in YAML, as a result Home Assistant doesn't
-            # set the storage collection into hass data.
-            # Major hack to get around this. 👻
-            collection = self.hass.data["websocket_api"]["zone/list"][
-                0
-            ].__self__.storage_collection
+        collection: ZoneStorageCollection = async_get_storage_collection(
+            self.hass, DOMAIN
+        )
 
         await collection.async_create_item(call.data.copy())

custom_components/spook/ectoplasms/zone/services/delete.py

@@ -13,6 +13,7 @@ from homeassistant.helpers.entity_component import DATA_INSTANCES, EntityCompone
 
 from ....const import DOMAIN as SPOOK_DOMAIN
 from ....errors import entity_not_found
+from ....helper_collections import async_get_storage_collection
 from ....services import AbstractSpookAdminService
 
 if TYPE_CHECKING:
@@ -32,16 +33,9 @@ class SpookService(AbstractSpookAdminService):
         """Handle the service call."""
         entity_component: EntityComponent[Zone] = self.hass.data[DATA_INSTANCES][DOMAIN]
 
-        collection: ZoneStorageCollection
-        if DOMAIN in self.hass.data:
-            collection = self.hass.data[DOMAIN]
-        else:
-            # Home zone is set in YAML, as a result Home Assistant doesn't
-            # set the storage collection into hass data.
-            # Major hack to get around this. 👻
-            collection = self.hass.data["websocket_api"]["zone/list"][
-                0
-            ].__self__.storage_collection
+        collection: ZoneStorageCollection = async_get_storage_collection(
+            self.hass, DOMAIN
+        )
 
         # Everything is looked up before anything is written. A bad one late
         # in the list should not leave the first ones changed behind an error.

custom_components/spook/ectoplasms/zone/services/update.py

@@ -20,6 +20,7 @@ from homeassistant.helpers.entity_component import DATA_INSTANCES, EntityCompone
 
 from ....const import DOMAIN as SPOOK_DOMAIN
 from ....errors import entity_not_found
+from ....helper_collections import async_get_storage_collection
 from ....services import AbstractSpookAdminService
 
 if TYPE_CHECKING:
@@ -55,16 +56,9 @@ class SpookService(AbstractSpookAdminService):
         """Handle the service call."""
         entity_component: EntityComponent[Zone] = self.hass.data[DATA_INSTANCES][DOMAIN]
 
-        collection: ZoneStorageCollection
-        if DOMAIN in self.hass.data:
-            collection = self.hass.data[DOMAIN]
-        else:
-            # Home zone is set in YAML, as a result Home Assistant doesn't
-            # set the storage collection into hass data.
-            # Major hack to get around this. 👻
-            collection = self.hass.data["websocket_api"]["zone/list"][
-                0
-            ].__self__.storage_collection
+        collection: ZoneStorageCollection = async_get_storage_collection(
+            self.hass, DOMAIN
+        )
 
         if not (entity := entity_component.get_entity(call.data["entity_id"])):
             raise entity_not_found(call.data["entity_id"])