diff --git a/homeassistant/components/energy/data.py b/homeassistant/components/energy/data.py index c053dea4741..9196694953a 100644 --- a/homeassistant/components/energy/data.py +++ b/homeassistant/components/energy/data.py @@ -79,7 +79,16 @@ class SolarSourceType(TypedDict): config_entry_solar_forecast: list[str] | None -SourceType = Union[GridSourceType, SolarSourceType] +class BatterySourceType(TypedDict): + """Dictionary holding the source of battery storage.""" + + type: Literal["battery"] + + stat_energy_from: str + stat_energy_to: str + + +SourceType = Union[GridSourceType, SolarSourceType, BatterySourceType] class DeviceConsumption(TypedDict): @@ -177,6 +186,13 @@ SOLAR_SOURCE_SCHEMA = vol.Schema( vol.Optional("config_entry_solar_forecast"): vol.Any([str], None), } ) +BATTERY_SOURCE_SCHEMA = vol.Schema( + { + vol.Required("type"): "battery", + vol.Required("stat_energy_from"): str, + vol.Required("stat_energy_to"): str, + } +) def check_type_limits(value: list[SourceType]) -> list[SourceType]: @@ -197,6 +213,7 @@ ENERGY_SOURCE_SCHEMA = vol.All( { "grid": GRID_SOURCE_SCHEMA, "solar": SOLAR_SOURCE_SCHEMA, + "battery": BATTERY_SOURCE_SCHEMA, }, ) ] diff --git a/tests/components/energy/test_websocket_api.py b/tests/components/energy/test_websocket_api.py index a14a8d0986e..60ac82108bc 100644 --- a/tests/components/energy/test_websocket_api.py +++ b/tests/components/energy/test_websocket_api.py @@ -104,6 +104,11 @@ async def test_save_preferences(hass, hass_ws_client, hass_storage) -> None: "stat_energy_from": "my_solar_production", "config_entry_solar_forecast": ["predicted_config_entry"], }, + { + "type": "battery", + "stat_energy_from": "my_battery_draining", + "stat_energy_to": "my_battery_charging", + }, ], "device_consumption": [{"stat_consumption": "some_device_usage"}], }