Update config flow and improve error handling

custom_components/irm_kmi/__init__.py

@@ -3,8 +3,8 @@
 # File inspired from https://github.com/ludeeus/integration_blueprint
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_ZONE
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryError
 
 from .const import DOMAIN, PLATFORMS
 from .coordinator import IrmKmiCoordinator
@@ -15,10 +15,14 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up this integration using UI."""
     hass.data.setdefault(DOMAIN, {})
 
-    hass.data[DOMAIN][entry.entry_id] = coordinator = IrmKmiCoordinator(hass, entry.data[CONF_ZONE])
+    hass.data[DOMAIN][entry.entry_id] = coordinator = IrmKmiCoordinator(hass, entry)
 
+    try:
         # https://developers.home-assistant.io/docs/integration_fetching_data#coordinated-single-api-poll-for-data-for-all-entities
         await coordinator.async_config_entry_first_refresh()
+    except ConfigEntryError:
+        # This happens when the zone is out of Benelux (no forecast available there)
+        return False
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     entry.async_on_unload(entry.add_update_listener(async_reload_entry))

custom_components/irm_kmi/camera.py

@@ -91,4 +91,3 @@ class IrmKmiRadar(CoordinatorEntity, Camera):
         """Return the camera state attributes."""
         attrs = {"hint": self.coordinator.data.get('animation', {}).get('hint')}
         return attrs
-

custom_components/irm_kmi/config_flow.py

@@ -6,9 +6,14 @@ from homeassistant.components.zone import DOMAIN as ZONE_DOMAIN
 from homeassistant.config_entries import ConfigFlow
 from homeassistant.const import CONF_ZONE
 from homeassistant.data_entry_flow import FlowResult
-from homeassistant.helpers.selector import EntitySelector, EntitySelectorConfig
+from homeassistant.helpers.selector import (EntitySelector,
+                                            EntitySelectorConfig,
+                                            SelectSelector,
+                                            SelectSelectorConfig,
+                                            SelectSelectorMode)
 
-from .const import DOMAIN
+from .const import (CONF_DARK_MODE, CONF_STYLE, CONF_STYLE_OPTIONS,
+                    CONF_STYLE_STD, DOMAIN)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -18,6 +23,7 @@ class IrmKmiConfigFlow(ConfigFlow, domain=DOMAIN):
 
     async def async_step_user(self, user_input: dict | None = None) -> FlowResult:
         """Define the user step of the configuration flow."""
+        print(f"IN CONFIG FLOW HERE with {user_input}")
         if user_input is not None:
             _LOGGER.debug(f"Provided config user is: {user_input}")
 
@@ -27,13 +33,21 @@ class IrmKmiConfigFlow(ConfigFlow, domain=DOMAIN):
             state = self.hass.states.get(user_input[CONF_ZONE])
             return self.async_create_entry(
                 title=state.name if state else "IRM KMI",
-                data={CONF_ZONE: user_input[CONF_ZONE]},
+                data={CONF_ZONE: user_input[CONF_ZONE],
+                      CONF_STYLE: user_input[CONF_STYLE],
+                      CONF_DARK_MODE: user_input[CONF_DARK_MODE]},
             )
 
         return self.async_show_form(
             step_id="user",
             data_schema=vol.Schema({
-                vol.Required(CONF_ZONE, description="Zone to use for weather forecast"):
+                vol.Required(CONF_ZONE):
                     EntitySelector(EntitySelectorConfig(domain=ZONE_DOMAIN)),
-            })
+
-        )
+                vol.Optional(CONF_STYLE, default=CONF_STYLE_STD):
+                    SelectSelector(SelectSelectorConfig(options=CONF_STYLE_OPTIONS,
+                                                        mode=SelectSelectorMode.DROPDOWN,
+                                                        translation_key=CONF_STYLE)),
+
+                vol.Optional(CONF_DARK_MODE, default=False): bool
+            }))

custom_components/irm_kmi/const.py

@@ -1,4 +1,5 @@
 """Constants for the IRM KMI integration."""
+from typing import Final
 
 from homeassistant.components.weather import (ATTR_CONDITION_CLEAR_NIGHT,
                                               ATTR_CONDITION_CLOUDY,
@@ -13,16 +14,32 @@ from homeassistant.components.weather import (ATTR_CONDITION_CLEAR_NIGHT,
                                               ATTR_CONDITION_SUNNY)
 from homeassistant.const import Platform
 
-DOMAIN = 'irm_kmi'
+DOMAIN: Final = 'irm_kmi'
-PLATFORMS: list[Platform] = [Platform.WEATHER, Platform.CAMERA]
+PLATFORMS: Final = [Platform.WEATHER, Platform.CAMERA]
-OUT_OF_BENELUX = ["außerhalb der Benelux (Brussels)",
+OUT_OF_BENELUX: Final = ["außerhalb der Benelux (Brussels)",
                          "Hors de Belgique (Bxl)",
                          "Outside the Benelux (Brussels)",
                          "Buiten de Benelux (Brussel)"]
-LANGS = ['en', 'fr', 'nl', 'de']
+LANGS: Final = ['en', 'fr', 'nl', 'de']
+
+
+CONF_STYLE_STD: Final = 'standard_style'
+CONF_STYLE_CONTRAST: Final = 'contrast_style'
+CONF_STYLE_YELLOW_RED: Final = 'yellow_red_style'
+CONF_STYLE_SATELLITE: Final = 'satellite_style'
+CONF_STYLE: Final = "style"
+
+CONF_STYLE_OPTIONS: Final = [
+    CONF_STYLE_STD,
+    CONF_STYLE_CONTRAST,
+    CONF_STYLE_YELLOW_RED,
+    CONF_STYLE_SATELLITE
+]
+
+CONF_DARK_MODE: Final = "dark_mode"
 
 # map ('ww', 'dayNight') tuple from IRM KMI to HA conditions
-IRM_KMI_TO_HA_CONDITION_MAP = {
+IRM_KMI_TO_HA_CONDITION_MAP: Final = {
     (0, 'd'): ATTR_CONDITION_SUNNY,
     (0, 'n'): ATTR_CONDITION_CLEAR_NIGHT,
     (1, 'd'): ATTR_CONDITION_SUNNY,

custom_components/irm_kmi/coordinator.py

@@ -7,20 +7,22 @@ from typing import Any, List, Tuple
 
 import async_timeout
 import pytz
+from PIL import Image, ImageDraw, ImageFont
 from homeassistant.components.weather import Forecast
-from homeassistant.components.zone import Zone
+from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_LATITUDE, ATTR_LONGITUDE
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryError
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.update_coordinator import (DataUpdateCoordinator,
                                                       UpdateFailed)
-from PIL import Image, ImageDraw, ImageFont
 
 from .api import IrmKmiApiClient, IrmKmiApiError
-from .const import IRM_KMI_TO_HA_CONDITION_MAP as CDT_MAP, LANGS
+from .const import IRM_KMI_TO_HA_CONDITION_MAP as CDT_MAP
-from .const import OUT_OF_BENELUX
+from .const import LANGS, OUT_OF_BENELUX
 from .data import (AnimationFrameData, CurrentWeatherData, IrmKmiForecast,
                    ProcessedCoordinatorData, RadarAnimationData)
+from .utils import disable_from_config, enable_from_config
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -28,7 +30,7 @@ _LOGGER = logging.getLogger(__name__)
 class IrmKmiCoordinator(DataUpdateCoordinator):
     """Coordinator to update data from IRM KMI"""
 
-    def __init__(self, hass: HomeAssistant, zone: Zone):
+    def __init__(self, hass: HomeAssistant, entry: ConfigEntry):
         """Initialize the coordinator."""
         super().__init__(
             hass,
@@ -39,7 +41,9 @@ class IrmKmiCoordinator(DataUpdateCoordinator):
             update_interval=timedelta(minutes=7),
         )
         self._api_client = IrmKmiApiClient(session=async_get_clientsession(hass))
-        self._zone = zone
+        self._zone = entry.data.get('zone')
+        self._config_entry = entry
+        self._disabled = False
 
     async def _async_update_data(self) -> ProcessedCoordinatorData:
         """Fetch data from API endpoint.
@@ -64,10 +68,24 @@ class IrmKmiCoordinator(DataUpdateCoordinator):
             raise UpdateFailed(f"Error communicating with API: {err}")
 
         if api_data.get('cityName', None) in OUT_OF_BENELUX:
-            raise UpdateFailed(f"Zone '{self._zone}' is out of Benelux and forecast is only available in the Benelux")
+            if self.data is None:
+                error_text = f"Zone '{self._zone}' is out of Benelux and forecast is only available in the Benelux"
+                _LOGGER.error(error_text)
+                raise ConfigEntryError(error_text)
+            else:
+                # TODO create a repair when this triggers
+                _LOGGER.error(f"The zone {self._zone} is now out of Benelux and forecast is only available in Benelux."
+                              f"Associated device is now disabled.  Move the zone back in Benelux and re-enable to fix "
+                              f"this")
+                disable_from_config(self.hass, self._config_entry)
+                return ProcessedCoordinatorData()
 
         return await self.process_api_data(api_data)
 
+    async def async_refresh(self) -> None:
+        """Refresh data and log errors."""
+        await self._async_refresh(log_failures=True, raise_on_entry_error=True)
+
     async def _async_animation_data(self, api_data: dict) -> RadarAnimationData:
         """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."""
@@ -103,7 +121,7 @@ class IrmKmiCoordinator(DataUpdateCoordinator):
         )
 
     async def download_images_from_api(self,
-                                       animation_data: dict,
+                                       animation_data: list,
                                        country: str,
                                        localisation_layer_url: str) -> tuple[Any]:
         """Download a batch of images to create the radar frames."""
@@ -172,7 +190,6 @@ class IrmKmiCoordinator(DataUpdateCoordinator):
             if most_recent_frame is None and current_time < time_image:
                 recent_idx = idx - 1 if idx > 0 else idx
                 most_recent_frame = sequence[recent_idx].get('image', None)
-                _LOGGER.debug(f"Most recent frame is at {sequence[recent_idx].get('time')}")
 
         background.close()
         most_recent_frame = most_recent_frame if most_recent_frame is not None else sequence[-1].get('image')

custom_components/irm_kmi/manifest.json

@@ -3,7 +3,7 @@
   "name": "IRM KMI Weather Belgium",
   "codeowners": ["@jdejaegh"],
   "config_flow": true,
-  "dependencies": [],
+  "dependencies": ["zone"],
   "documentation": "https://github.com/jdejaegh/irm-kmi-ha/",
   "integration_type": "service",
   "iot_class": "cloud_polling",

custom_components/irm_kmi/translations/en.json

@@ -3,9 +3,22 @@
   "config": {
     "step": {
       "user": {
-        "title": "Select a zone",
+        "title": "Configuration",
         "data": {
-          "zone": "Zone"
+          "zone": "Zone",
+          "style": "Style of the radar",
+          "dark_mode": "Radar dark mode"
+        }
+      }
+    }
+  },
+  "selector": {
+    "style": {
+      "options": {
+        "standard_style": "Standard",
+        "contrast_style": "High contrast",
+        "yellow_red_style": "Yellow-Red",
+        "satellite_style": "Satellite"
       }
     }
   },
@@ -14,4 +27,3 @@
     "unknown": "Unknown error occurred"
   }
 }
-}

custom_components/irm_kmi/utils.py

@@ -0,0 +1,25 @@
+import logging
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers import device_registry, entity_registry
+
+_LOGGER = logging.getLogger(__name__)
+
+
+def disable_from_config(hass: HomeAssistant, config_entry: ConfigEntry):
+    modify_from_config(hass, config_entry, False)
+
+
+def enable_from_config(hass: HomeAssistant, config_entry: ConfigEntry):
+    modify_from_config(hass, config_entry, True)
+
+
+def modify_from_config(hass: HomeAssistant, config_entry: ConfigEntry, enable: bool):
+    dr = device_registry.async_get(hass)
+    devices = device_registry.async_entries_for_config_entry(dr, config_entry.entry_id)
+    _LOGGER.info(f"Trying to {'enable' if enable else 'disable'} {config_entry.entry_id}: {len(devices)} device(s)")
+    for device in devices:
+        _LOGGER.info(f"Disabling device {device.name} because it is out of Benelux")
+        dr.async_update_device(device_id=device.id,
+                               disabled_by=None if enable else device_registry.DeviceEntryDisabler.INTEGRATION)

custom_components/irm_kmi/weather.py

@@ -23,9 +23,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry, async_add_e
     """Set up the weather entry."""
 
     coordinator = hass.data[DOMAIN][entry.entry_id]
-    async_add_entities(
+    async_add_entities([IrmKmiWeather(coordinator, entry)])
-        [IrmKmiWeather(coordinator, entry)]
-    )
 
 
 class IrmKmiWeather(CoordinatorEntity, WeatherEntity):
@@ -59,11 +57,11 @@ class IrmKmiWeather(CoordinatorEntity, WeatherEntity):
 
     @property
     def condition(self) -> str | None:
-        return self.coordinator.data.get('current_weather').get('condition')
+        return self.coordinator.data.get('current_weather', {}).get('condition')
 
     @property
     def native_temperature(self) -> float | None:
-        return self.coordinator.data.get('current_weather').get('temperature')
+        return self.coordinator.data.get('current_weather', {}).get('temperature')
 
     @property
     def native_temperature_unit(self) -> str | None:
@@ -75,15 +73,15 @@ class IrmKmiWeather(CoordinatorEntity, WeatherEntity):
 
     @property
     def native_wind_speed(self) -> float | None:
-        return self.coordinator.data.get('current_weather').get('wind_speed')
+        return self.coordinator.data.get('current_weather', {}).get('wind_speed')
 
     @property
     def native_wind_gust_speed(self) -> float | None:
-        return self.coordinator.data.get('current_weather').get('wind_gust_speed')
+        return self.coordinator.data.get('current_weather', {}).get('wind_gust_speed')
 
     @property
     def wind_bearing(self) -> float | str | None:
-        return self.coordinator.data.get('current_weather').get('wind_bearing')
+        return self.coordinator.data.get('current_weather', {}).get('wind_bearing')
 
     @property
     def native_precipitation_unit(self) -> str | None:
@@ -91,7 +89,7 @@ class IrmKmiWeather(CoordinatorEntity, WeatherEntity):
 
     @property
     def native_pressure(self) -> float | None:
-        return self.coordinator.data.get('current_weather').get('pressure')
+        return self.coordinator.data.get('current_weather', {}).get('pressure')
 
     @property
     def native_pressure_unit(self) -> str | None:
@@ -99,7 +97,7 @@ class IrmKmiWeather(CoordinatorEntity, WeatherEntity):
 
     @property
     def uv_index(self) -> float | None:
-        return self.coordinator.data.get('current_weather').get('uv_index')
+        return self.coordinator.data.get('current_weather', {}).get('uv_index')
 
     async def async_forecast_twice_daily(self) -> List[Forecast] | None:
         return self.coordinator.data.get('daily_forecast')

tests/conftest.py

@@ -11,7 +11,7 @@ from pytest_homeassistant_custom_component.common import (MockConfigEntry,
                                                           load_fixture)
 
 from custom_components.irm_kmi.api import IrmKmiApiParametersError
-from custom_components.irm_kmi.const import DOMAIN
+from custom_components.irm_kmi.const import DOMAIN, CONF_STYLE, CONF_STYLE_STD, CONF_DARK_MODE
 
 
 @pytest.fixture(autouse=True)
@@ -25,7 +25,9 @@ def mock_config_entry() -> MockConfigEntry:
     return MockConfigEntry(
         title="Home",
         domain=DOMAIN,
-        data={CONF_ZONE: "zone.home"},
+        data={CONF_ZONE: "zone.home",
+              CONF_STYLE: CONF_STYLE_STD,
+              CONF_DARK_MODE: True},
         unique_id="zone.home",
     )
 
@@ -91,7 +93,7 @@ def mock_image_irm_kmi_api(request: pytest.FixtureRequest) -> Generator[None, Ma
         elif "getLocalizationLayerBE" in url:
             file_name = "tests/fixtures/loc_layer_be_n.png"
         elif "getLocalizationLayerNL" in url:
-            file_name = "tests/fixtures/loc_layer_nl_d.png"
+            file_name = "tests/fixtures/loc_layer_nl.png"
         else:
             raise ValueError("Not a valid parameter for the mock")
 

tests/test_config_flow.py

@@ -8,7 +8,7 @@ from homeassistant.const import CONF_ZONE
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 
-from custom_components.irm_kmi.const import DOMAIN
+from custom_components.irm_kmi.const import DOMAIN, CONF_STYLE, CONF_STYLE_STD, CONF_DARK_MODE
 
 
 async def test_full_user_flow(
@@ -25,9 +25,13 @@ async def test_full_user_flow(
 
     result2 = await hass.config_entries.flow.async_configure(
         result["flow_id"],
-        user_input={CONF_ZONE: ENTITY_ID_HOME},
+        user_input={CONF_ZONE: ENTITY_ID_HOME,
+                    CONF_STYLE: CONF_STYLE_STD,
+                    CONF_DARK_MODE: False},
     )
-
+    print(result2)
     assert result2.get("type") == FlowResultType.CREATE_ENTRY
-    assert result2.get("title") == "IRM KMI"
+    assert result2.get("title") == "test home"
-    assert result2.get("data") == {CONF_ZONE: ENTITY_ID_HOME}
+    assert result2.get("data") == {CONF_ZONE: ENTITY_ID_HOME,
+                                   CONF_STYLE: CONF_STYLE_STD,
+                                   CONF_DARK_MODE: False}

tests/test_coordinator.py

@@ -11,7 +11,7 @@ from homeassistant.components.weather import (ATTR_CONDITION_CLOUDY,
 from homeassistant.components.zone import Zone
 from homeassistant.core import HomeAssistant
 from PIL import Image, ImageDraw, ImageFont
-from pytest_homeassistant_custom_component.common import load_fixture
+from pytest_homeassistant_custom_component.common import load_fixture, MockConfigEntry
 
 from custom_components.irm_kmi.coordinator import IrmKmiCoordinator
 from custom_components.irm_kmi.data import CurrentWeatherData, IrmKmiForecast
@@ -111,9 +111,10 @@ def test_hourly_forecast() -> None:
 @freeze_time(datetime.fromisoformat("2023-12-28T15:30:00+01:00"))
 async def test_get_image_nl(
         hass: HomeAssistant,
-        mock_image_irm_kmi_api: AsyncMock) -> None:
+        mock_image_irm_kmi_api: AsyncMock,
+        mock_config_entry: MockConfigEntry) -> None:
     api_data = get_api_data("forecast_nl.json")
-    coordinator = IrmKmiCoordinator(hass, Zone({}))
+    coordinator = IrmKmiCoordinator(hass, mock_config_entry)
 
     result = await coordinator._async_animation_data(api_data)
 
@@ -121,7 +122,7 @@ async def test_get_image_nl(
     tz = pytz.timezone(hass.config.time_zone)
     background = Image.open("custom_components/irm_kmi/resources/nl.png").convert('RGBA')
     layer = Image.open("tests/fixtures/clouds_nl.png").convert('RGBA')
-    localisation = Image.open("tests/fixtures/loc_layer_nl_d.png").convert('RGBA')
+    localisation = Image.open("tests/fixtures/loc_layer_nl.png").convert('RGBA')
     temp = Image.alpha_composite(background, layer)
     expected = Image.alpha_composite(temp, localisation)
     draw = ImageDraw.Draw(expected)
@@ -145,9 +146,10 @@ async def test_get_image_nl(
 async def test_get_image_be(
         hass: HomeAssistant,
         mock_image_irm_kmi_api: AsyncMock,
+        mock_config_entry: MockConfigEntry
 ) -> None:
     api_data = get_api_data("forecast.json")
-    coordinator = IrmKmiCoordinator(hass, Zone({}))
+    coordinator = IrmKmiCoordinator(hass, mock_config_entry)
 
     result = await coordinator._async_animation_data(api_data)
 

tests/test_init.py

@@ -98,5 +98,5 @@ async def test_zone_out_of_benelux(
     await hass.config_entries.async_setup(mock_config_entry.entry_id)
     await hass.async_block_till_done()
 
-    assert mock_config_entry.state is ConfigEntryState.SETUP_RETRY
+    assert mock_config_entry.state is ConfigEntryState.SETUP_ERROR
     assert "Zone 'zone.london' is out of Benelux" in caplog.text