Adjust energy to not override Entity.add_to_platform_xxx methods (#183775)

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

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()

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()