diff --git a/homeassistant/components/forecast_solar/__init__.py b/homeassistant/components/forecast_solar/__init__.py index b00e5f1c4ce..b20a0befb96 100644 --- a/homeassistant/components/forecast_solar/__init__.py +++ b/homeassistant/components/forecast_solar/__init__.py @@ -5,10 +5,12 @@ 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 +from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.update_coordinator import DataUpdateCoordinator from .const import ( @@ -55,7 +57,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: ) await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {}) + 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.config_entries.async_setup_platforms(entry, PLATFORMS) @@ -77,3 +81,20 @@ 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] = { + timestamp.isoformat(): val + for timestamp, val in coordinator.data.watts.items() + } + + connection.send_result(msg["id"], forecasts) diff --git a/tests/components/forecast_solar/test_init.py b/tests/components/forecast_solar/test_init.py index 719041aaf58..7544f7d352b 100644 --- a/tests/components/forecast_solar/test_init.py +++ b/tests/components/forecast_solar/test_init.py @@ -1,4 +1,5 @@ """Tests for the Forecast.Solar integration.""" +from datetime import datetime, timezone from unittest.mock import MagicMock, patch from forecast_solar import ForecastSolarConnectionError @@ -14,14 +15,37 @@ 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.watts = { + 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 + # 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: { + "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()