Add an energy solar platform for solar forecasts (#54576)

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

homeassistant/components/energy/types.py

@@ -0,0 +1,27 @@
+"""Types for the energy platform."""
+from __future__ import annotations
+
+from typing import Awaitable, Callable, TypedDict
+
+from homeassistant.core import HomeAssistant
+
+
+class SolarForecastType(TypedDict):
+    """Return value for solar forecast."""
+
+    wh_hours: dict[str, float | int]
+
+
+GetSolarForecastType = Callable[
+    [HomeAssistant, str], Awaitable["SolarForecastType | None"]
+]
+
+
+class EnergyPlatform:
+    """This class represents the methods we expect on the energy platforms."""
+
+    @staticmethod
+    async def async_get_solar_forecast(
+        hass: HomeAssistant, config_entry_id: str
+    ) -> SolarForecastType | None:
+        """Get forecast for solar production for specific config entry ID."""

homeassistant/components/energy/websocket_api.py

@@ -3,12 +3,17 @@ from __future__ import annotations
 
 import asyncio
 import functools
-from typing import Any, Awaitable, Callable, Dict, cast
+from types import ModuleType
+from typing import Any, Awaitable, Callable, cast
 
 import voluptuous as vol
 
 from homeassistant.components import websocket_api
 from homeassistant.core import HomeAssistant, callback
+from homeassistant.helpers.integration_platform import (
+    async_process_integration_platforms,
+)
+from homeassistant.helpers.singleton import singleton
 
 from .const import DOMAIN
 from .data import (
@@ -18,14 +23,15 @@ from .data import (
     EnergyPreferencesUpdate,
     async_get_manager,
 )
+from .types import EnergyPlatform, GetSolarForecastType
 from .validate import async_validate
 
 EnergyWebSocketCommandHandler = Callable[
-    [HomeAssistant, websocket_api.ActiveConnection, Dict[str, Any], "EnergyManager"],
+    [HomeAssistant, websocket_api.ActiveConnection, "dict[str, Any]", "EnergyManager"],
     None,
 ]
 AsyncEnergyWebSocketCommandHandler = Callable[
-    [HomeAssistant, websocket_api.ActiveConnection, Dict[str, Any], "EnergyManager"],
+    [HomeAssistant, websocket_api.ActiveConnection, "dict[str, Any]", "EnergyManager"],
     Awaitable[None],
 ]
 
@@ -37,6 +43,28 @@ def async_setup(hass: HomeAssistant) -> None:
     websocket_api.async_register_command(hass, ws_save_prefs)
     websocket_api.async_register_command(hass, ws_info)
     websocket_api.async_register_command(hass, ws_validate)
+    websocket_api.async_register_command(hass, ws_solar_forecast)
+
+
+@singleton("energy_platforms")
+async def async_get_energy_platforms(
+    hass: HomeAssistant,
+) -> dict[str, GetSolarForecastType]:
+    """Get energy platforms."""
+    platforms: dict[str, GetSolarForecastType] = {}
+
+    async def _process_energy_platform(
+        hass: HomeAssistant, domain: str, platform: ModuleType
+    ) -> None:
+        """Process energy platforms."""
+        if not hasattr(platform, "async_get_solar_forecast"):
+            return
+
+        platforms[domain] = cast(EnergyPlatform, platform).async_get_solar_forecast
+
+    await async_process_integration_platforms(hass, DOMAIN, _process_energy_platform)
+
+    return platforms
 
 
 def _ws_with_manager(
@@ -107,14 +135,21 @@ async def ws_save_prefs(
         vol.Required("type"): "energy/info",
     }
 )
-@callback
-def ws_info(
+@websocket_api.async_response
+async def ws_info(
     hass: HomeAssistant,
     connection: websocket_api.ActiveConnection,
     msg: dict,
 ) -> None:
     """Handle get info command."""
-    connection.send_result(msg["id"], hass.data[DOMAIN])
+    forecast_platforms = await async_get_energy_platforms(hass)
+    connection.send_result(
+        msg["id"],
+        {
+            "cost_sensors": hass.data[DOMAIN]["cost_sensors"],
+            "solar_forecast_domains": list(forecast_platforms),
+        },
+    )
 
 
 @websocket_api.websocket_command(
@@ -130,3 +165,56 @@ async def ws_validate(
 ) -> None:
     """Handle validate command."""
     connection.send_result(msg["id"], (await async_validate(hass)).as_dict())
+
+
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "energy/solar_forecast",
+    }
+)
+@_ws_with_manager
+async def ws_solar_forecast(
+    hass: HomeAssistant,
+    connection: websocket_api.ActiveConnection,
+    msg: dict,
+    manager: EnergyManager,
+) -> None:
+    """Handle solar forecast command."""
+    if manager.data is None:
+        connection.send_result(msg["id"], {})
+        return
+
+    config_entries: dict[str, str | None] = {}
+
+    for source in manager.data["energy_sources"]:
+        if (
+            source["type"] != "solar"
+            or source.get("config_entry_solar_forecast") is None
+        ):
+            continue
+
+        # typing is not catching the above guard for config_entry_solar_forecast being none
+        for config_entry in source["config_entry_solar_forecast"]:  # type: ignore[union-attr]
+            config_entries[config_entry] = None
+
+    if not config_entries:
+        connection.send_result(msg["id"], {})
+        return
+
+    forecasts = {}
+
+    forecast_platforms = await async_get_energy_platforms(hass)
+
+    for config_entry_id in config_entries:
+        config_entry = hass.config_entries.async_get_entry(config_entry_id)
+        # Filter out non-existing config entries or unsupported domains
+
+        if config_entry is None or config_entry.domain not in forecast_platforms:
+            continue
+
+        forecast = await forecast_platforms[config_entry.domain](hass, config_entry_id)
+
+        if forecast is not None:
+            forecasts[config_entry_id] = forecast
+
+    connection.send_result(msg["id"], forecasts)

homeassistant/components/forecast_solar/__init__.py

@@ -5,12 +5,10 @@ from datetime import timedelta
 import logging
 
 from forecast_solar import ForecastSolar
-import voluptuous as vol
 
-from homeassistant.components import websocket_api
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
@@ -60,10 +58,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
     await coordinator.async_config_entry_first_refresh()
 
-    if DOMAIN not in hass.data:
-        hass.data[DOMAIN] = {}
-        websocket_api.async_register_command(hass, ws_list_forecasts)
-    hass.data[DOMAIN][entry.entry_id] = coordinator
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
 
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
@@ -84,22 +79,3 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 async def async_update_options(hass: HomeAssistant, entry: ConfigEntry) -> None:
     """Update options."""
     await hass.config_entries.async_reload(entry.entry_id)
-
-
-@websocket_api.websocket_command({vol.Required("type"): "forecast_solar/forecasts"})
-@callback
-def ws_list_forecasts(
-    hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict
-) -> None:
-    """Return a list of available forecasts."""
-    forecasts = {}
-
-    for config_entry_id, coordinator in hass.data[DOMAIN].items():
-        forecasts[config_entry_id] = {
-            "wh_hours": {
-                timestamp.isoformat(): val
-                for timestamp, val in coordinator.data.wh_hours.items()
-            }
-        }
-
-    connection.send_result(msg["id"], forecasts)

homeassistant/components/forecast_solar/energy.py

@@ -0,0 +1,23 @@
+"""Energy platform."""
+from __future__ import annotations
+
+from homeassistant.core import HomeAssistant
+
+from .const import DOMAIN
+
+
+async def async_get_solar_forecast(
+    hass: HomeAssistant, config_entry_id: str
+) -> dict[str, dict[str, float | int]] | None:
+    """Get solar forecast for a config entry ID."""
+    coordinator = hass.data[DOMAIN].get(config_entry_id)
+
+    if coordinator is None:
+        return None
+
+    return {
+        "wh_hours": {
+            timestamp.isoformat(): val
+            for timestamp, val in coordinator.data.wh_hours.items()
+        }
+    }

homeassistant/helpers/integration_platform.py

@@ -23,7 +23,7 @@ async def async_process_integration_platforms(
     """Process a specific platform for all current and future loaded integrations."""
 
     async def _process(component_name: str) -> None:
-        """Process the intents of a component."""
+        """Process component being loaded."""
         if "." in component_name:
             return
 

tests/components/energy/test_websocket_api.py

@@ -1,10 +1,12 @@
 """Test the Energy websocket API."""
+from unittest.mock import AsyncMock, Mock
+
 import pytest
 
 from homeassistant.components.energy import data, is_configured
 from homeassistant.setup import async_setup_component
 
-from tests.common import flush_store
+from tests.common import MockConfigEntry, flush_store, mock_platform
 
 
 @pytest.fixture(autouse=True)
@@ -15,6 +17,26 @@ async def setup_integration(hass):
     )
 
 
+@pytest.fixture
+def mock_energy_platform(hass):
+    """Mock an energy platform."""
+    hass.config.components.add("some_domain")
+    mock_platform(
+        hass,
+        "some_domain.energy",
+        Mock(
+            async_get_solar_forecast=AsyncMock(
+                return_value={
+                    "wh_hours": {
+                        "2021-06-27T13:00:00+00:00": 12,
+                        "2021-06-27T14:00:00+00:00": 8,
+                    }
+                }
+            )
+        ),
+    )
+
+
 async def test_get_preferences_no_data(hass, hass_ws_client) -> None:
     """Test we get error if no preferences set."""
     client = await hass_ws_client(hass)
@@ -46,7 +68,9 @@ async def test_get_preferences_default(hass, hass_ws_client, hass_storage) -> No
     assert msg["result"] == data.EnergyManager.default_preferences()
 
 
-async def test_save_preferences(hass, hass_ws_client, hass_storage) -> None:
+async def test_save_preferences(
+    hass, hass_ws_client, hass_storage, mock_energy_platform
+) -> None:
     """Test we can save preferences."""
     client = await hass_ws_client(hass)
 
@@ -140,7 +164,8 @@ async def test_save_preferences(hass, hass_ws_client, hass_storage) -> None:
         "cost_sensors": {
             "sensor.heat_pump_meter_2": "sensor.heat_pump_meter_2_cost",
             "sensor.return_to_grid_offpeak": "sensor.return_to_grid_offpeak_compensation",
-        }
+        },
+        "solar_forecast_domains": ["some_domain"],
     }
 
     # Prefs with limited options
@@ -232,3 +257,35 @@ async def test_validate(hass, hass_ws_client) -> None:
         "energy_sources": [],
         "device_consumption": [],
     }
+
+
+async def test_get_solar_forecast(hass, hass_ws_client, mock_energy_platform) -> None:
+    """Test we get preferences."""
+    entry = MockConfigEntry(domain="some_domain")
+    entry.add_to_hass(hass)
+
+    manager = await data.async_get_manager(hass)
+    manager.data = data.EnergyManager.default_preferences()
+    manager.data["energy_sources"].append(
+        {
+            "type": "solar",
+            "stat_energy_from": "my_solar_production",
+            "config_entry_solar_forecast": [entry.entry_id],
+        }
+    )
+    client = await hass_ws_client(hass)
+
+    await client.send_json({"id": 5, "type": "energy/solar_forecast"})
+
+    msg = await client.receive_json()
+
+    assert msg["id"] == 5
+    assert msg["success"]
+    assert msg["result"] == {
+        entry.entry_id: {
+            "wh_hours": {
+                "2021-06-27T13:00:00+00:00": 12,
+                "2021-06-27T14:00:00+00:00": 8,
+            }
+        }
+    }

tests/components/forecast_solar/test_energy.py

@@ -0,0 +1,34 @@
+"""Test forecast solar energy platform."""
+from datetime import datetime, timezone
+from unittest.mock import MagicMock
+
+from homeassistant.components.forecast_solar import energy
+from homeassistant.config_entries import ConfigEntryState
+from homeassistant.core import HomeAssistant
+
+from tests.common import MockConfigEntry
+
+
+async def test_energy_solar_forecast(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_forecast_solar: MagicMock,
+) -> None:
+    """Test the Forecast.Solar energy platform solar forecast."""
+    mock_forecast_solar.estimate.return_value.wh_hours = {
+        datetime(2021, 6, 27, 13, 0, tzinfo=timezone.utc): 12,
+        datetime(2021, 6, 27, 14, 0, tzinfo=timezone.utc): 8,
+    }
+
+    mock_config_entry.add_to_hass(hass)
+    await hass.config_entries.async_setup(mock_config_entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert mock_config_entry.state == ConfigEntryState.LOADED
+
+    assert await energy.async_get_solar_forecast(hass, mock_config_entry.entry_id) == {
+        "wh_hours": {
+            "2021-06-27T13:00:00+00:00": 12,
+            "2021-06-27T14:00:00+00:00": 8,
+        }
+    }

tests/components/forecast_solar/test_init.py

@@ -1,5 +1,4 @@
 """Tests for the Forecast.Solar integration."""
-from datetime import datetime, timezone
 from unittest.mock import MagicMock, patch
 
 from forecast_solar import ForecastSolarConnectionError
@@ -7,6 +6,7 @@ from forecast_solar import ForecastSolarConnectionError
 from homeassistant.components.forecast_solar.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
+from homeassistant.setup import async_setup_component
 
 from tests.common import MockConfigEntry
 
@@ -15,39 +15,13 @@ async def test_load_unload_config_entry(
     hass: HomeAssistant,
     mock_config_entry: MockConfigEntry,
     mock_forecast_solar: MagicMock,
-    hass_ws_client,
 ) -> None:
     """Test the Forecast.Solar configuration entry loading/unloading."""
-    mock_forecast_solar.estimate.return_value.wh_hours = {
-        datetime(2021, 6, 27, 13, 0, tzinfo=timezone.utc): 12,
-        datetime(2021, 6, 27, 14, 0, tzinfo=timezone.utc): 8,
-    }
-
     mock_config_entry.add_to_hass(hass)
-    await hass.config_entries.async_setup(mock_config_entry.entry_id)
-    await hass.async_block_till_done()
+    await async_setup_component(hass, "forecast_solar", {})
 
     assert mock_config_entry.state == ConfigEntryState.LOADED
 
-    # Test WS API set up
-    client = await hass_ws_client(hass)
-    await client.send_json(
-        {
-            "id": 5,
-            "type": "forecast_solar/forecasts",
-        }
-    )
-    result = await client.receive_json()
-    assert result["success"]
-    assert result["result"] == {
-        mock_config_entry.entry_id: {
-            "wh_hours": {
-                "2021-06-27T13:00:00+00:00": 12,
-                "2021-06-27T14:00:00+00:00": 8,
-            }
-        }
-    }
-
     await hass.config_entries.async_unload(mock_config_entry.entry_id)
     await hass.async_block_till_done()