diff --git a/homeassistant/components/energy/sensor.py b/homeassistant/components/energy/sensor.py index 85fd93d59fda..3d3bac397414 100644 --- a/homeassistant/components/energy/sensor.py +++ b/homeassistant/components/energy/sensor.py @@ -397,11 +397,12 @@ class EnergyCostSensor(SensorEntity): self._attr_state_class = SensorStateClass.TOTAL self._config = config self._last_energy_sensor_state: State | None = None - # add_finished is set when either of async_added_to_hass - # or add_to_platform_abort is called + # SensorManager awaits add_finished; async_on_remove resolves it on the + # abort path too, since add_to_platform_abort fires on-remove callbacks. self.add_finished: asyncio.Future[None] = ( asyncio.get_running_loop().create_future() ) + self.async_on_remove(lambda: _set_result_unless_done(self.add_finished)) def _reset(self, energy_state: State) -> None: """Reset the cost sensor.""" @@ -612,13 +613,6 @@ class EnergyCostSensor(SensorEntity): self._update_cost() self.async_write_ha_state() - @callback - @override - def add_to_platform_abort(self) -> None: - """Abort adding an entity to a platform.""" - _set_result_unless_done(self.add_finished) - super().add_to_platform_abort() - @override async def async_will_remove_from_hass(self) -> None: """Handle removing from hass.""" @@ -697,11 +691,12 @@ class EnergyPowerSensor(SensorEntity): config["stat_rate_to"], ] - # add_finished is set when either async_added_to_hass or add_to_platform_abort - # is called + # SensorManager awaits add_finished; async_on_remove resolves it on the + # abort path too, since add_to_platform_abort fires on-remove callbacks. self.add_finished: asyncio.Future[None] = ( asyncio.get_running_loop().create_future() ) + self.async_on_remove(lambda: _set_result_unless_done(self.add_finished)) @property @override @@ -847,10 +842,3 @@ class EnergyPowerSensor(SensorEntity): """Handle source sensor state changes.""" self._update_state() self.async_write_ha_state() - - @callback - @override - def add_to_platform_abort(self) -> None: - """Abort adding an entity to a platform.""" - _set_result_unless_done(self.add_finished) - super().add_to_platform_abort() diff --git a/tests/components/energy/test_sensor.py b/tests/components/energy/test_sensor.py index 6ce13d86fc2a..dc418857dc71 100644 --- a/tests/components/energy/test_sensor.py +++ b/tests/components/energy/test_sensor.py @@ -2758,53 +2758,52 @@ async def test_missing_price_entity( assert state.state == "150.0" -async def test_energy_cost_sensor_add_to_platform_abort( - recorder_mock: Recorder, hass: HomeAssistant -) -> None: - """Test EnergyCostSensor.add_to_platform_abort sets the future.""" - adapter = SourceAdapter( - source_type="grid", - flow_type="flow_from", - stat_energy_key="stat_energy_from", - total_money_key="stat_cost", - name_suffix="Cost", - entity_id_suffix="cost", +def _make_cost_sensor() -> EnergyCostSensor: + """Build an EnergyCostSensor for the abort test.""" + return EnergyCostSensor( + SourceAdapter( + source_type="grid", + flow_type="flow_from", + stat_energy_key="stat_energy_from", + total_money_key="stat_cost", + name_suffix="Cost", + entity_id_suffix="cost", + ), + { + "stat_energy_from": "sensor.energy", + "stat_cost": None, + "entity_energy_price": "sensor.price", + "number_energy_price": None, + }, ) - config = { - "stat_energy_from": "sensor.energy", - "stat_cost": None, - "entity_energy_price": "sensor.price", - "number_energy_price": None, - } - - sensor = EnergyCostSensor(adapter, config) - - # Future should not be done yet - assert not sensor.add_finished.done() - - # Call abort - sensor.add_to_platform_abort() - - # Future should now be done - assert sensor.add_finished.done() -async def test_energy_power_sensor_add_to_platform_abort( - recorder_mock: Recorder, hass: HomeAssistant -) -> None: - """Test EnergyPowerSensor.add_to_platform_abort sets the future.""" - sensor = EnergyPowerSensor( +def _make_power_sensor() -> EnergyPowerSensor: + """Build an EnergyPowerSensor for the abort test.""" + return EnergyPowerSensor( source_type="battery", config={"stat_rate_inverted": "sensor.battery_power"}, unique_id="test_unique_id", entity_id="sensor.test_power", ) - # Future should not be done yet + +@pytest.mark.usefixtures("recorder_mock", "hass") +@pytest.mark.parametrize( + "make_sensor", + [ + pytest.param(_make_cost_sensor, id="cost"), + pytest.param(_make_power_sensor, id="power"), + ], +) +async def test_add_finished_resolved_on_abort( + make_sensor: Callable[[], EnergyCostSensor | EnergyPowerSensor], +) -> None: + """Test an aborted add resolves add_finished so SensorManager does not hang.""" + sensor = make_sensor() + assert not sensor.add_finished.done() - # Call abort sensor.add_to_platform_abort() - # Future should now be done assert sensor.add_finished.done()