Extract Forecast.Solar DataUpdateCoordinator into module (#83859)

homeassistant/components/forecast_solar/__init__.py

@@ -1,66 +1,19 @@
 """The Forecast.Solar integration."""
 from __future__ import annotations
 
-from datetime import timedelta
-import logging
-
-from forecast_solar import Estimate, ForecastSolar
-
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE, Platform
+from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .const import (
+from .const import DOMAIN
-    CONF_AZIMUTH,
+from .coordinator import ForecastSolarDataUpdateCoordinator
-    CONF_DAMPING,
-    CONF_DECLINATION,
-    CONF_INVERTER_SIZE,
-    CONF_MODULES_POWER,
-    DOMAIN,
-)
 
 PLATFORMS = [Platform.SENSOR]
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Forecast.Solar from a config entry."""
-    # Our option flow may cause it to be an empty string,
+    coordinator = ForecastSolarDataUpdateCoordinator(hass, entry)
-    # this if statement is here to catch that.
-    api_key = entry.options.get(CONF_API_KEY) or None
-
-    if (
-        inverter_size := entry.options.get(CONF_INVERTER_SIZE)
-    ) is not None and inverter_size > 0:
-        inverter_size = inverter_size / 1000
-
-    session = async_get_clientsession(hass)
-    forecast = ForecastSolar(
-        api_key=api_key,
-        session=session,
-        latitude=entry.data[CONF_LATITUDE],
-        longitude=entry.data[CONF_LONGITUDE],
-        declination=entry.options[CONF_DECLINATION],
-        azimuth=(entry.options[CONF_AZIMUTH] - 180),
-        kwp=(entry.options[CONF_MODULES_POWER] / 1000),
-        damping=entry.options.get(CONF_DAMPING, 0),
-        inverter=inverter_size,
-    )
-
-    # Free account have a resolution of 1 hour, using that as the default
-    # update interval. Using a higher value for accounts with an API key.
-    update_interval = timedelta(hours=1)
-    if api_key is not None:
-        update_interval = timedelta(minutes=30)
-
-    coordinator: DataUpdateCoordinator[Estimate] = DataUpdateCoordinator(
-        hass,
-        logging.getLogger(__name__),
-        name=DOMAIN,
-        update_method=forecast.estimate,
-        update_interval=update_interval,
-    )
     await coordinator.async_config_entry_first_refresh()
 
     hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator

homeassistant/components/forecast_solar/const.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 from datetime import timedelta
+import logging
 
 from homeassistant.components.sensor import SensorDeviceClass, SensorStateClass
 from homeassistant.const import UnitOfEnergy, UnitOfPower
@@ -9,6 +10,7 @@ from homeassistant.const import UnitOfEnergy, UnitOfPower
 from .models import ForecastSolarSensorEntityDescription
 
 DOMAIN = "forecast_solar"
+LOGGER = logging.getLogger(__package__)
 
 CONF_DECLINATION = "declination"
 CONF_AZIMUTH = "azimuth"

homeassistant/components/forecast_solar/coordinator.py

@@ -0,0 +1,65 @@
+"""DataUpdateCoordinator for the Forecast.Solar integration."""
+from __future__ import annotations
+
+from datetime import timedelta
+
+from forecast_solar import Estimate, ForecastSolar
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
+
+from .const import (
+    CONF_AZIMUTH,
+    CONF_DAMPING,
+    CONF_DECLINATION,
+    CONF_INVERTER_SIZE,
+    CONF_MODULES_POWER,
+    DOMAIN,
+    LOGGER,
+)
+
+
+class ForecastSolarDataUpdateCoordinator(DataUpdateCoordinator[Estimate]):
+    """The Forecast.Solar Data Update Coordinator."""
+
+    config_entry: ConfigEntry
+
+    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
+        """Initialize the Forecast.Solar coordinator."""
+        self.config_entry = entry
+
+        # Our option flow may cause it to be an empty string,
+        # this if statement is here to catch that.
+        api_key = entry.options.get(CONF_API_KEY) or None
+
+        if (
+            inverter_size := entry.options.get(CONF_INVERTER_SIZE)
+        ) is not None and inverter_size > 0:
+            inverter_size = inverter_size / 1000
+
+        self.forecast = ForecastSolar(
+            api_key=api_key,
+            session=async_get_clientsession(hass),
+            latitude=entry.data[CONF_LATITUDE],
+            longitude=entry.data[CONF_LONGITUDE],
+            declination=entry.options[CONF_DECLINATION],
+            azimuth=(entry.options[CONF_AZIMUTH] - 180),
+            kwp=(entry.options[CONF_MODULES_POWER] / 1000),
+            damping=entry.options.get(CONF_DAMPING, 0),
+            inverter=inverter_size,
+        )
+
+        # Free account have a resolution of 1 hour, using that as the default
+        # update interval. Using a higher value for accounts with an API key.
+        update_interval = timedelta(hours=1)
+        if api_key is not None:
+            update_interval = timedelta(minutes=30)
+
+        super().__init__(hass, LOGGER, name=DOMAIN, update_interval=update_interval)
+
+    async def _async_update_data(self) -> Estimate:
+        """Fetch Forecast.Solar estimates."""
+        return await self.forecast.estimate()

homeassistant/components/forecast_solar/sensor.py

@@ -10,12 +10,10 @@ from homeassistant.helpers.device_registry import DeviceEntryType
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
-from homeassistant.helpers.update_coordinator import (
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-)
 
 from .const import DOMAIN, SENSORS
+from .coordinator import ForecastSolarDataUpdateCoordinator
 from .models import ForecastSolarSensorEntityDescription
 
 
@@ -23,7 +21,7 @@ async def async_setup_entry(
     hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
 ) -> None:
     """Defer sensor setup to the shared sensor module."""
-    coordinator: DataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator: ForecastSolarDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
 
     async_add_entities(
         ForecastSolarSensorEntity(
@@ -35,7 +33,9 @@ async def async_setup_entry(
     )
 
 
-class ForecastSolarSensorEntity(CoordinatorEntity, SensorEntity):
+class ForecastSolarSensorEntity(
+    CoordinatorEntity[ForecastSolarDataUpdateCoordinator], SensorEntity
+):
     """Defines a Forecast.Solar sensor."""
 
     entity_description: ForecastSolarSensorEntityDescription
@@ -45,7 +45,7 @@ class ForecastSolarSensorEntity(CoordinatorEntity, SensorEntity):
         self,
         *,
         entry_id: str,
-        coordinator: DataUpdateCoordinator,
+        coordinator: ForecastSolarDataUpdateCoordinator,
         entity_description: ForecastSolarSensorEntityDescription,
     ) -> None:
         """Initialize Forecast.Solar sensor."""

tests/components/forecast_solar/conftest.py

@@ -60,7 +60,8 @@ def mock_forecast_solar(hass) -> Generator[None, MagicMock, None]:
     hass fixture included because it sets the time zone.
     """
     with patch(
-        "homeassistant.components.forecast_solar.ForecastSolar", autospec=True
+        "homeassistant.components.forecast_solar.coordinator.ForecastSolar",
+        autospec=True,
     ) as forecast_solar_mock:
         forecast_solar = forecast_solar_mock.return_value
         now = datetime(2021, 6, 27, 6, 0, tzinfo=dt_util.DEFAULT_TIME_ZONE)

tests/components/forecast_solar/test_init.py

@@ -29,7 +29,7 @@ async def test_load_unload_config_entry(
 
 
 @patch(
-    "homeassistant.components.forecast_solar.ForecastSolar.estimate",
+    "homeassistant.components.forecast_solar.coordinator.ForecastSolar.estimate",
     side_effect=ForecastSolarConnectionError,
 )
 async def test_config_entry_not_ready(