Separate concerns: rain graph and api

custom_components/irm_kmi/coordinator.py

@@ -112,10 +112,14 @@ class IrmKmiCoordinator(TimestampDataUpdateCoordinator):
                 if self.data is not None else PollenParser.get_unavailable_data()
 
         try:
-            radar_animation, image_path, bg_size = self._api.get_animation_data(tz, lang, self._style,
+            radar_animation = self._api.get_animation_data(tz, lang, self._style, self._dark_mode)
-                                                                                      self._dark_mode)
+            animation = await RainGraph(radar_animation,
-            animation = await RainGraph(radar_animation, image_path, bg_size, tz=tz, dark_mode=self._dark_mode,
+                                        country=self._api.get_country(),
-                                        api_client=self._api).build()
+                                        style=self._style,
+                                        tz=tz,
+                                        dark_mode=self._dark_mode,
+                                        api_client=self._api
+                                        ).build()
         except ValueError:
             animation = None
 

custom_components/irm_kmi/irm_kmi_api/api.py

@@ -417,8 +417,7 @@ class IrmKmiApiClientHa(IrmKmiApiClient):
             )
         return forecast
 
-    def get_animation_data(self, tz: ZoneInfo, lang: str, style: str, dark_mode: bool) -> (RadarAnimationData,
+    def get_animation_data(self, tz: ZoneInfo, lang: str, style: str, dark_mode: bool) -> RadarAnimationData:
-                                                                                                 str, Tuple[int, int]):
         """From the API data passed in, call the API to get all the images and create the radar animation data object.
         Frames from the API are merged with the background map and the location marker to create each frame."""
         animation_data = self._api_data.get('animation', {}).get('sequence')
@@ -443,11 +442,9 @@ class IrmKmiApiClientHa(IrmKmiApiClient):
         r = self._get_rain_graph_data(
             radar_animation,
             animation_data,
-            country,
             images_from_api,
-            tz,
+            tz
-            style,
+        )
-            dark_mode)
 
         return r
 
@@ -520,12 +517,9 @@ class IrmKmiApiClientHa(IrmKmiApiClient):
     @staticmethod
     def _get_rain_graph_data(radar_animation: RadarAnimationData,
                              api_animation_data: List[dict],
-                             country: str | None,
                              images_from_api: list[str],
                              tz: ZoneInfo,
-                             style: str,
+                             ) -> RadarAnimationData:
-                             dark_mode: bool
-                             ) -> (RadarAnimationData, str, Tuple[int, int]):
         """Create a RainGraph object that is ready to output animated and still SVG images"""
         sequence: List[AnimationFrameData] = list()
 
@@ -549,14 +543,4 @@ class IrmKmiApiClientHa(IrmKmiApiClient):
         radar_animation['sequence'] = sequence
         radar_animation['most_recent_image_idx'] = most_recent_frame
 
-        satellite_mode = style == OPTION_STYLE_SATELLITE
+        return radar_animation
-
-        if country == 'NL':
-            image_path = "custom_components/irm_kmi/resources/nl.png"
-            bg_size = (640, 600)
-        else:
-            image_path = (f"custom_components/irm_kmi/resources/be_"
-                          f"{'satellite' if satellite_mode else 'black' if dark_mode else 'white'}.png")
-            bg_size = (640, 490)
-
-        return radar_animation, image_path, bg_size

custom_components/irm_kmi/irm_kmi_api/rain_graph.py

@@ -7,12 +7,12 @@ import logging
 from typing import List, Self, Any
 
 import async_timeout
-from homeassistant.util import dt
 from svgwrite import Drawing
 from svgwrite.animate import Animate
 from svgwrite.container import FONT_TEMPLATE
 
 from .api import IrmKmiApiClient
+from .const import OPTION_STYLE_SATELLITE
 from .data import AnimationFrameData, RadarAnimationData
 from .resources import be_black, be_satellite, be_white, nl, roboto
 
@@ -22,10 +22,10 @@ _LOGGER = logging.getLogger(__name__)
 class RainGraph:
     def __init__(self,
                  animation_data: RadarAnimationData,
-                 background_image_path: str,
+                 country: str,
-                 background_size: (int, int),
+                 style: str,
                  dark_mode: bool = False,
-                 tz: datetime.tzinfo = dt.get_default_time_zone(),
+                 tz: datetime.tzinfo = None,
                  svg_width: float = 640,
                  inset: float = 20,
                  graph_height: float = 150,
@@ -37,20 +37,26 @@ class RainGraph:
                  ):
 
         self._animation_data: RadarAnimationData = animation_data
-        self._background_image_path: str = background_image_path
+        self._country: str = country
-        self._background_size: (int, int) = background_size
+
+        if self._country == 'NL':
+            self._background_size: (int, int) = (640, 600)
+        else:
+            self._background_size: (int, int) = (640, 490)
+
+        self._style = style
         self._dark_mode: bool = dark_mode
         self._tz = tz
         self._svg_width: float = svg_width
         self._inset: float = inset
         self._graph_height: float = graph_height
-        self._top_text_space: float = top_text_space + background_size[1]
+        self._top_text_space: float = top_text_space + self._background_size[1]
-        self._top_text_y_pos: float = top_text_y_pos + background_size[1]
+        self._top_text_y_pos: float = top_text_y_pos + self._background_size[1]
         self._bottom_text_space: float = bottom_text_space
-        self._bottom_text_y_pos: float = bottom_text_y_pos + background_size[1]
+        self._bottom_text_y_pos: float = bottom_text_y_pos + self._background_size[1]
         self._api_client = api_client
 
-        self._frame_count: int = len(self._animation_data['sequence'])
+        self._frame_count: int = max(len(self._animation_data['sequence']), 1)
         self._graph_width: float = self._svg_width - 2 * self._inset
         self._graph_bottom: float = self._top_text_space + self._graph_height
         self._svg_height: float = self._graph_height + self._top_text_space + self._bottom_text_space
@@ -404,13 +410,13 @@ class RainGraph:
         return copy.deepcopy(self._dwg)
 
     def get_background_png_b64(self):
-        _LOGGER.debug(f"Get b64 for {self._background_image_path}")
+        _LOGGER.debug(f"Get b64 for {self._country} {self._style} {'dark' if self._dark_mode else 'light'} mode")
-        if self._background_image_path.endswith('be_black.png'):
+        if self._country == 'NL':
-            return be_black.be_black_b64
-        elif self._background_image_path.endswith('be_white.png'):
-            return be_white.be_white_b64
-        elif self._background_image_path.endswith('be_satellite.png'):
-            return be_satellite.be_satelitte_b64
-        elif self._background_image_path.endswith('nl.png'):
             return nl.nl_b64
-        return None
+        elif self._style == OPTION_STYLE_SATELLITE:
+            return be_satellite.be_satelitte_b64
+        elif self._dark_mode:
+            return be_black.be_black_b64
+        else:
+            return be_white.be_white_b64
+

tests/conftest.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 
 import json
 from typing import Generator
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, patch, AsyncMock
 
 import pytest
 from homeassistant.const import CONF_ZONE
@@ -13,6 +13,7 @@ from custom_components.irm_kmi import OPTION_STYLE_STD
 from custom_components.irm_kmi.const import (CONF_DARK_MODE, CONF_LANGUAGE_OVERRIDE, CONF_STYLE,
                                              CONF_USE_DEPRECATED_FORECAST, DOMAIN, OPTION_DEPRECATED_FORECAST_NOT_USED,
                                              OPTION_DEPRECATED_FORECAST_TWICE_DAILY, IRM_KMI_TO_HA_CONDITION_MAP)
+from custom_components.irm_kmi.data import ProcessedCoordinatorData
 from custom_components.irm_kmi.irm_kmi_api.api import IrmKmiApiError, IrmKmiApiParametersError, IrmKmiApiClientHa
 
 

tests/test_rain_graph.py

@@ -1,5 +1,6 @@
 import base64
-from datetime import datetime, timedelta
+from datetime import datetime as dt, timedelta
+import datetime
 
 from custom_components.irm_kmi.irm_kmi_api.data import AnimationFrameData, RadarAnimationData
 from custom_components.irm_kmi.irm_kmi_api.rain_graph import RainGraph
@@ -13,7 +14,7 @@ def get_radar_animation_data() -> RadarAnimationData:
 
     sequence = [
         AnimationFrameData(
-            time=datetime.fromisoformat("2023-12-26T18:30:00+00:00") + timedelta(minutes=10 * i),
+            time=dt.fromisoformat("2023-12-26T18:30:00+00:00") + timedelta(minutes=10 * i),
             image=image_data,
             value=2,
             position=.5,
@@ -36,8 +37,8 @@ async def test_svg_frame_setup():
     data = get_radar_animation_data()
     rain_graph = RainGraph(
         animation_data=data,
-        background_image_path="custom_components/irm_kmi/irm_kmi_api/resources/be_white.png",
+        country='BE',
-        background_size=(640, 490),
+        style='STD'
     )
 
     await rain_graph.draw_svg_frame()
@@ -56,8 +57,8 @@ def test_svg_hint():
     data = get_radar_animation_data()
     rain_graph = RainGraph(
         animation_data=data,
-        background_image_path="custom_components/irm_kmi/irm_kmi_api/resources/be_white.png",
+        country='BE',
-        background_size=(640, 490),
+        style='STD'
     )
 
     rain_graph.write_hint()
@@ -71,8 +72,9 @@ def test_svg_time_bars():
     data = get_radar_animation_data()
     rain_graph = RainGraph(
         animation_data=data,
-        background_image_path="custom_components/irm_kmi/irm_kmi_api/resources/be_white.png",
+        country='BE',
-        background_size=(640, 490),
+        style='STD',
+        tz=datetime.UTC
     )
 
     rain_graph.draw_hour_bars()
@@ -90,8 +92,9 @@ def test_draw_chances_path():
     data = get_radar_animation_data()
     rain_graph = RainGraph(
         animation_data=data,
-        background_image_path="custom_components/irm_kmi/irm_kmi_api/resources/be_white.png",
+        country='BE',
-        background_size=(640, 490),
+        style='STD',
+        tz=datetime.UTC
     )
 
     rain_graph.draw_chances_path()
@@ -108,8 +111,8 @@ def test_draw_data_line():
     data = get_radar_animation_data()
     rain_graph = RainGraph(
         animation_data=data,
-        background_image_path="custom_components/irm_kmi/irm_kmi_api/resources/be_white.png",
+        country='BE',
-        background_size=(640, 490),
+        style='STD'
     )
 
     rain_graph.draw_data_line()
@@ -126,8 +129,8 @@ async def test_insert_background():
     data = get_radar_animation_data()
     rain_graph = RainGraph(
         animation_data=data,
-        background_image_path="custom_components/irm_kmi/irm_kmi_api/resources/be_white.png",
+        country='BE',
-        background_size=(640, 490),
+        style='STD'
     )
 
     await rain_graph.insert_background()
@@ -149,8 +152,8 @@ def test_draw_current_frame_line_moving():
     data = get_radar_animation_data()
     rain_graph = RainGraph(
         animation_data=data,
-        background_image_path="custom_components/irm_kmi/irm_kmi_api/resources/be_white.png",
+        country='BE',
-        background_size=(640, 490),
+        style='STD'
     )
 
     rain_graph.draw_current_fame_line()
@@ -177,8 +180,8 @@ def test_draw_current_frame_line_index():
     data = get_radar_animation_data()
     rain_graph = RainGraph(
         animation_data=data,
-        background_image_path="custom_components/irm_kmi/irm_kmi_api/resources/be_white.png",
+        country='BE',
-        background_size=(640, 490),
+        style='STD'
     )
 
     rain_graph.draw_current_fame_line(0)
@@ -205,8 +208,9 @@ def test_draw_description_text():
     data = get_radar_animation_data()
     rain_graph = RainGraph(
         animation_data=data,
-        background_image_path="custom_components/irm_kmi/irm_kmi_api/resources/be_white.png",
+        country='BE',
-        background_size=(640, 490),
+        style='STD',
+        tz=datetime.UTC
     )
 
     rain_graph.draw_description_text()
@@ -232,8 +236,8 @@ def test_draw_cloud_layer():
     data = get_radar_animation_data()
     rain_graph = RainGraph(
         animation_data=data,
-        background_image_path="custom_components/irm_kmi/irm_kmi_api/resources/be_white.png",
+        country='BE',
-        background_size=(640, 490),
+        style='STD'
     )
 
     rain_graph.insert_cloud_layer()
@@ -252,8 +256,8 @@ async def test_draw_location_layer():
     data = get_radar_animation_data()
     rain_graph = RainGraph(
         animation_data=data,
-        background_image_path="custom_components/irm_kmi/irm_kmi_api/resources/be_white.png",
+        country='BE',
-        background_size=(640, 490),
+        style='STD'
     )
 
     await rain_graph.draw_location()

tests/test_sensors.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, MagicMock
 
 from freezegun import freeze_time
 from homeassistant.core import HomeAssistant
@@ -10,6 +10,7 @@ from custom_components.irm_kmi.binary_sensor import IrmKmiWarning
 from custom_components.irm_kmi.const import CONF_LANGUAGE_OVERRIDE
 from custom_components.irm_kmi.sensor import IrmKmiNextSunMove, IrmKmiNextWarning
 from tests.conftest import get_api_with_data
+from tests.test_rain_graph import get_radar_animation_data
 
 
 @freeze_time(datetime.fromisoformat('2024-01-12T07:55:00+01:00'))
@@ -45,7 +46,7 @@ async def test_warning_data_unknown_lang(
     coordinator = IrmKmiCoordinator(hass, mock_config_entry)
 
     api.get_pollen = AsyncMock()
-    api.get_animation_data = AsyncMock()
+    api.get_animation_data = MagicMock(return_value=get_radar_animation_data())
     coordinator._api = api
 
 
@@ -76,7 +77,7 @@ async def test_next_warning_when_data_available(
     coordinator = IrmKmiCoordinator(hass, mock_config_entry)
 
     api.get_pollen = AsyncMock()
-    api.get_animation_data = AsyncMock()
+    api.get_animation_data = MagicMock(return_value=get_radar_animation_data())
     coordinator._api = api
 
     result = await coordinator.process_api_data()
@@ -105,7 +106,7 @@ async def test_next_warning_none_when_only_active_warnings(
     coordinator = IrmKmiCoordinator(hass, mock_config_entry)
 
     api.get_pollen = AsyncMock()
-    api.get_animation_data = AsyncMock()
+    api.get_animation_data = MagicMock(return_value=get_radar_animation_data())
     coordinator._api = api
 
     result = await coordinator.process_api_data()
@@ -170,7 +171,7 @@ async def test_next_sunrise_sunset(
 
     coordinator = IrmKmiCoordinator(hass, mock_config_entry)
     api.get_pollen = AsyncMock()
-    api.get_animation_data = AsyncMock()
+    api.get_animation_data = MagicMock(return_value=get_radar_animation_data())
     coordinator._api = api
 
     result = await coordinator.process_api_data()
@@ -199,7 +200,7 @@ async def test_next_sunrise_sunset_bis(
 
     coordinator = IrmKmiCoordinator(hass, mock_config_entry)
     api.get_pollen = AsyncMock()
-    api.get_animation_data = AsyncMock()
+    api.get_animation_data = MagicMock(return_value=get_radar_animation_data())
     coordinator._api = api
 
     result = await coordinator.process_api_data()