From b0ef9623f931a840e6b616b39a748ee11b7a8463 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Fri, 8 Dec 2023 16:51:07 +0100
Subject: [PATCH] Add test for energy cost sensor for late price sensor
 (#105312)

---
 tests/components/energy/test_sensor.py | 108 +++++++++++++++++++++++++
 1 file changed, 108 insertions(+)

diff --git a/tests/components/energy/test_sensor.py b/tests/components/energy/test_sensor.py
index f4a1f661f9b..fe3e6d08653 100644
--- a/tests/components/energy/test_sensor.py
+++ b/tests/components/energy/test_sensor.py
@@ -877,6 +877,114 @@ async def test_cost_sensor_handle_price_units(
     assert state.state == "20.0"
 
 
+async def test_cost_sensor_handle_late_price_sensor(
+    setup_integration,
+    hass: HomeAssistant,
+    hass_storage: dict[str, Any],
+) -> None:
+    """Test energy cost where the price sensor is not immediately available."""
+    energy_attributes = {
+        ATTR_UNIT_OF_MEASUREMENT: UnitOfEnergy.KILO_WATT_HOUR,
+        ATTR_STATE_CLASS: SensorStateClass.TOTAL_INCREASING,
+    }
+    price_attributes = {
+        ATTR_UNIT_OF_MEASUREMENT: f"EUR/{UnitOfEnergy.KILO_WATT_HOUR}",
+        ATTR_STATE_CLASS: SensorStateClass.MEASUREMENT,
+    }
+    energy_data = data.EnergyManager.default_preferences()
+    energy_data["energy_sources"].append(
+        {
+            "type": "grid",
+            "flow_from": [
+                {
+                    "stat_energy_from": "sensor.energy_consumption",
+                    "stat_cost": None,
+                    "entity_energy_price": "sensor.energy_price",
+                    "number_energy_price": None,
+                }
+            ],
+            "flow_to": [],
+            "cost_adjustment_day": 0,
+        }
+    )
+
+    hass_storage[data.STORAGE_KEY] = {
+        "version": 1,
+        "data": energy_data,
+    }
+
+    # Initial state: 10kWh, price sensor not yet available
+    hass.states.async_set("sensor.energy_price", "unknown", price_attributes)
+    hass.states.async_set(
+        "sensor.energy_consumption",
+        10,
+        energy_attributes,
+    )
+
+    await setup_integration(hass)
+
+    state = hass.states.get("sensor.energy_consumption_cost")
+    assert state.state == "unknown"
+
+    # Energy use bumped by 10 kWh, price sensor still not yet available
+    hass.states.async_set(
+        "sensor.energy_consumption",
+        20,
+        energy_attributes,
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.energy_consumption_cost")
+    assert state.state == "unknown"
+
+    # Energy use bumped by 10 kWh, price sensor now available
+    hass.states.async_set("sensor.energy_price", "1", price_attributes)
+    hass.states.async_set(
+        "sensor.energy_consumption",
+        30,
+        energy_attributes,
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.energy_consumption_cost")
+    assert state.state == "0.0"
+
+    # Energy use bumped by 10 kWh, price sensor available
+    hass.states.async_set(
+        "sensor.energy_consumption",
+        40,
+        energy_attributes,
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.energy_consumption_cost")
+    assert state.state == "10.0"
+
+    # Energy use bumped by 10 kWh, price sensor no longer available
+    hass.states.async_set("sensor.energy_price", "unknown", price_attributes)
+    hass.states.async_set(
+        "sensor.energy_consumption",
+        50,
+        energy_attributes,
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.energy_consumption_cost")
+    assert state.state == "10.0"
+
+    # Energy use bumped by 10 kWh, price sensor again available
+    hass.states.async_set("sensor.energy_price", "2", price_attributes)
+    hass.states.async_set(
+        "sensor.energy_consumption",
+        60,
+        energy_attributes,
+    )
+    await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.energy_consumption_cost")
+    assert state.state == "50.0"
+
+
 @pytest.mark.parametrize(
     "unit",
     (UnitOfVolume.CUBIC_FEET, UnitOfVolume.CUBIC_METERS),