From 92fe57778840cb7627adf63aaaf23dfbf838ddae Mon Sep 17 00:00:00 2001
From: Andrey Khrolenok <andrey@khrolenok.ru>
Date: Fri, 10 May 2024 13:13:35 +0300
Subject: [PATCH] Add ability to use hourly and twice daily forecasts as source

---
 custom_components/car_wash/binary_sensor.py | 17 +++---
 tests/test_binary_sensor.py                 | 60 +++++++++++++--------
 2 files changed, 49 insertions(+), 28 deletions(-)

diff --git a/custom_components/car_wash/binary_sensor.py b/custom_components/car_wash/binary_sensor.py
index 4ac4533..2c31dce 100644
--- a/custom_components/car_wash/binary_sensor.py
+++ b/custom_components/car_wash/binary_sensor.py
@@ -162,19 +162,22 @@ async def async_update(self):
         temp = wstate.attributes.get(ATTR_WEATHER_TEMPERATURE)
         cond = wstate.state
 
-        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
index 24f2489..073a4da 100644
--- 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
@@ -29,9 +29,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
 
@@ -132,7 +133,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(
@@ -144,7 +145,22 @@ 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 component update."""
+    assert default_sensor.is_on is None
+
+    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(
@@ -152,40 +168,42 @@ async def test_async_update_forecast_fail(hass: HomeAssistant, default_sensor):
         "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 component update."""
-    assert default_sensor.is_on is None
-
-    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)