Add ability to use hourly and twice daily forecasts as source

Limych · May 10, 2024 · c0089cc · c0089cc
1 parent 3bdf98f
commit c0089cc
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 26 deletions.
diff --git a/custom_components/snowtire/binary_sensor.py b/custom_components/snowtire/binary_sensor.py
@@ -167,19 +167,22 @@ async def async_update(
         tmpu = self.hass.config.units.temperature_unit
         temp = wstate.attributes.get(ATTR_WEATHER_TEMPERATURE)
 
-        wfeatures = wstate.attributes.get(ATTR_SUPPORTED_FEATURES)
-        if (
-            not isinstance(wfeatures, WeatherEntityFeature)
-            or (wfeatures & WeatherEntityFeature.FORECAST_DAILY) == 0
-        ):
-            raise HomeAssistantError("Weather entity doesn't support 'daily' forecast")
+        wfeatures = wstate.attributes.get(ATTR_SUPPORTED_FEATURES) or 0
+        if (wfeatures & WeatherEntityFeature.FORECAST_DAILY) != 0:
+            forecast_type = "daily"
+        elif (wfeatures & WeatherEntityFeature.FORECAST_TWICE_DAILY) != 0:
+            forecast_type = "twice_daily"
+        elif (wfeatures & WeatherEntityFeature.FORECAST_HOURLY) != 0:
+            forecast_type = "hourly"
+        else:
+            raise HomeAssistantError("Weather entity doesn't support any forecast")
 
         try:
             forecast = await self.hass.services.async_call(
                 WEATHER_DOMAIN,
                 SERVICE_GET_FORECASTS,
                 {
-                    CONF_TYPE: "daily",
+                    CONF_TYPE: forecast_type,
                     CONF_ENTITY_ID: self._weather_entity,
                 },
                 blocking=True,

diff --git a/tests/test_binary_sensor.py b/tests/test_binary_sensor.py
@@ -2,7 +2,7 @@
 
 # pylint: disable=redefined-outer-name
 from typing import Final
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import pytest
 from pytest import raises
@@ -30,9 +30,10 @@
 from homeassistant.const import (
     ATTR_SUPPORTED_FEATURES,
     CONF_PLATFORM,
+    CONF_TYPE,
     UnitOfTemperature,
 )
-from homeassistant.core import HomeAssistant, SupportsResponse
+from homeassistant.core import HomeAssistant, ServiceRegistry, SupportsResponse
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.util import dt as dt_util
 
@@ -124,7 +125,7 @@ async def test_async_update_forecast_fail(hass: HomeAssistant, default_sensor):
         },
     )
 
-    with raises(HomeAssistantError, match="doesn't support 'daily' forecast"):
+    with raises(HomeAssistantError, match="doesn't support any forecast"):
         await default_sensor.async_update()
 
     hass.states.async_set(
@@ -136,46 +137,63 @@ async def test_async_update_forecast_fail(hass: HomeAssistant, default_sensor):
         },
     )
 
-    with raises(HomeAssistantError, match="doesn't support 'daily' forecast"):
+    with raises(TypeError):
+        await default_sensor.async_update()
+
+
+async def test_async_update(hass: HomeAssistant, default_sensor):
+    """Test sensor update."""
+    hass.states.async_set(MOCK_WEATHER_ENTITY, None)
+
+    with raises(HomeAssistantError):
+        await default_sensor.async_update()
+
+    hass.states.async_set(MOCK_WEATHER_ENTITY, "State")
+
+    with raises(HomeAssistantError):
         await default_sensor.async_update()
 
     hass.states.async_set(
         MOCK_WEATHER_ENTITY,
         "State",
         attributes={
             ATTR_WEATHER_TEMPERATURE: -1,
-            ATTR_SUPPORTED_FEATURES: WeatherEntityFeature.FORECAST_HOURLY,
+            ATTR_SUPPORTED_FEATURES: WeatherEntityFeature.FORECAST_DAILY
+            | WeatherEntityFeature.FORECAST_TWICE_DAILY
+            | WeatherEntityFeature.FORECAST_HOURLY,
         },
     )
 
-    with raises(HomeAssistantError, match="doesn't support 'daily' forecast"):
+    with patch.object(ServiceRegistry, "async_call") as call:
         await default_sensor.async_update()
+        assert call.call_args.args[2][CONF_TYPE] == "daily"
 
     hass.states.async_set(
         MOCK_WEATHER_ENTITY,
         "State",
         attributes={
             ATTR_WEATHER_TEMPERATURE: -1,
-            ATTR_SUPPORTED_FEATURES: WeatherEntityFeature.FORECAST_HOURLY
-            | WeatherEntityFeature.FORECAST_DAILY,
+            ATTR_SUPPORTED_FEATURES: WeatherEntityFeature.FORECAST_TWICE_DAILY
+            | WeatherEntityFeature.FORECAST_HOURLY,
         },
     )
 
-    with raises(HomeAssistantError, match="Can't get forecast data!"):
-        await default_sensor.async_update()
-
-
-async def test_async_update(hass: HomeAssistant, default_sensor):
-    """Test sensor update."""
-    hass.states.async_set(MOCK_WEATHER_ENTITY, None)
-
-    with raises(HomeAssistantError):
+    with patch.object(ServiceRegistry, "async_call") as call:
         await default_sensor.async_update()
+        assert call.call_args.args[2][CONF_TYPE] == "twice_daily"
 
-    hass.states.async_set(MOCK_WEATHER_ENTITY, "State")
+    hass.states.async_set(
+        MOCK_WEATHER_ENTITY,
+        "State",
+        attributes={
+            ATTR_WEATHER_TEMPERATURE: -1,
+            ATTR_SUPPORTED_FEATURES: WeatherEntityFeature.FORECAST_HOURLY,
+        },
+    )
 
-    with raises(HomeAssistantError):
+    with patch.object(ServiceRegistry, "async_call") as call:
         await default_sensor.async_update()
+        assert call.call_args.args[2][CONF_TYPE] == "hourly"
 
     today = dt_util.start_of_local_day()
     today_ts = int(today.timestamp() * 1000)