diff --git a/homeassistant/components/coronavirus/sensor.py b/homeassistant/components/coronavirus/sensor.py index d24c33a7752..7f0e0c230e6 100644 --- a/homeassistant/components/coronavirus/sensor.py +++ b/homeassistant/components/coronavirus/sensor.py @@ -1,6 +1,6 @@ """Sensor platform for the Corona virus.""" from homeassistant.const import ATTR_ATTRIBUTION -from homeassistant.helpers.entity import Entity +from homeassistant.helpers.update_coordinator import CoordinatorEntity from . import get_coordinator from .const import ATTRIBUTION, OPTION_WORLDWIDE @@ -23,7 +23,7 @@ async def async_setup_entry(hass, config_entry, async_add_entities): ) -class CoronavirusSensor(Entity): +class CoronavirusSensor(CoordinatorEntity): """Sensor representing corona virus data.""" name = None @@ -31,12 +31,12 @@ class CoronavirusSensor(Entity): def __init__(self, coordinator, country, info_type): """Initialize coronavirus sensor.""" + super().__init__(coordinator) if country == OPTION_WORLDWIDE: self.name = f"Worldwide Coronavirus {info_type}" else: self.name = f"{coordinator.data[country].country} Coronavirus {info_type}" self.unique_id = f"{country}-{info_type}" - self.coordinator = coordinator self.country = country self.info_type = info_type @@ -76,11 +76,3 @@ class CoronavirusSensor(Entity): def device_state_attributes(self): """Return device attributes.""" return {ATTR_ATTRIBUTION: ATTRIBUTION} - - async def async_added_to_hass(self): - """When entity is added to hass.""" - self.coordinator.async_add_listener(self.async_write_ha_state) - - async def async_will_remove_from_hass(self): - """When entity will be removed from hass.""" - self.coordinator.async_remove_listener(self.async_write_ha_state)