will need to test if this blocks when there's >1 core

tcarwash · Apr 18, 2024 · 25ed2e3 · 25ed2e3
1 parent ee88ebd
commit 25ed2e3
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 40 deletions.
diff --git a/custom_components/noaa_space_weather/api.py b/custom_components/noaa_space_weather/api.py
@@ -29,7 +29,6 @@ async def async_get_data(self) -> dict:
     async def async_get_first_frame(self, product) -> bytes:
         response_json = await self.swpc.get_data_method(product)
         first_frame_url = response_json[0].get("url")
-        _LOGGER.debug(first_frame_url)
         return await self.swpc.get_bytes_method(first_frame_url)
 
     async def async_load_animation(self, product) -> bytes:

diff --git a/custom_components/noaa_space_weather/entity.py b/custom_components/noaa_space_weather/entity.py
@@ -7,25 +7,6 @@
 from .const import DOMAIN
 
 
-class NoaaSpaceWeatherAnimationEntity(ImageEntity):
-    """NOAA Space Weather Image Entity"""
-
-    def __init__(self, coordinator, config_entry):
-        super().__init__(coordinator.hass, config_entry)
-        self.coordinator = coordinator
-        self.config_entry = config_entry
-
-    @property
-    def extra_state_attributes(self):
-        """Return the state attributes."""
-        return {
-            "attribution": ATTRIBUTION,
-            "id": str(self.coordinator.data.get("id")),
-            "integration": DOMAIN,
-            "last_updated": self.image_last_updated,
-        }
-
-
 class NoaaSpaceWeatherImageEntity(CoordinatorEntity, ImageEntity):
     """NOAA Space Weather Image Entity"""
 

diff --git a/custom_components/noaa_space_weather/image.py b/custom_components/noaa_space_weather/image.py
@@ -3,9 +3,9 @@
 from .const import DOMAIN, ICON
 import logging
 import asyncio
-from .entity import NoaaSpaceWeatherImageEntity, NoaaSpaceWeatherAnimationEntity
+from .entity import NoaaSpaceWeatherImageEntity
 from homeassistant.core import callback
-from datetime import datetime, timedelta
+from datetime import datetime
 
 
 _LOGGER: logging.Logger = logging.getLogger(__package__)
@@ -78,11 +78,10 @@ async def async_setup_entry(hass, entry, async_add_devices):
     )
     async_add_devices(
         [NoaaSpaceWeatherAnimation(coordinator, entry, image=i) for i in animationmap],
-        update_before_add=True,
     )
 
 
-class NoaaSpaceWeatherAnimation(NoaaSpaceWeatherAnimationEntity):
+class NoaaSpaceWeatherAnimation(NoaaSpaceWeatherImageEntity):
     """noaa_space_weather Image class."""
 
     def __init__(self, coordinator, entry, image):
@@ -122,32 +121,32 @@ def icon(self):
 
     async def async_update(self):
         if not self._cached_image:
-            _LOGGER.debug("returning still")
+            _LOGGER.debug(f"Returning still for {self.name}")
             image_bytes = await self.coordinator.api.async_get_first_frame(
                 self.image_data["product"]
             )
-            self._cached_image = image_bytes
-            self.image_last_updated = datetime.now() - timedelta(days=1)
-            loop = asyncio.get_event_loop()
-            loop.create_task(self.async_update())
+            asyncio.run_coroutine_threadsafe(self.async_update(), self.hass.loop)
         else:
-            self.image_last_updated = datetime.now()
-            self._attr_image_last_updated = self.image_last_updated
+            _LOGGER.debug(f"Returning still for {self.name}")
             image_bytes = await self.coordinator.api.async_load_animation(
                 self.image_data["product"]
             )
             _LOGGER.debug(f"Updated animation for {self.name}, caching image")
-            self._cached_image = image_bytes
+        self._cached_image = image_bytes
+        self.image_last_updated = datetime.now()
+        self._attr_image_last_updated = self.image_last_updated
+        _LOGGER.debug(f"Writing {self.name} state")
+        return image_bytes
+
+    @callback
+    def _handle_coordinator_update(self):
+        self.image_last_updated = datetime.now()
+        self._attr_image_last_updated = self.image_last_updated
+        asyncio.run_coroutine_threadsafe(self.async_update(), self.hass.loop)
+        self.async_write_ha_state()
 
     async def async_image(self):
-        if not self._cached_image or (
-            self.image_last_updated < datetime.now() - timedelta(minutes=5)
-        ):
-            loop = asyncio.get_event_loop()
-            loop.create_task(self.async_update())
-        else:
-            _LOGGER.debug(f"returning cached image for: {self.name}")
-        return self._cached_image
+        return self._cached_image or await self.async_update()
 
 
 class NoaaSpaceWeatherImage(NoaaSpaceWeatherImageEntity):