diff --git a/custom_components/home_assistant_covid19_augsburg/__init__.py b/custom_components/home_assistant_covid19_augsburg/__init__.py index 6263d32..bb9fe06 100644 --- a/custom_components/home_assistant_covid19_augsburg/__init__.py +++ b/custom_components/home_assistant_covid19_augsburg/__init__.py @@ -67,9 +67,12 @@ async def get_coordinator(hass: HomeAssistant): if DOMAIN in hass.data: return hass.data[DOMAIN] - async def async_get_data() -> IncidenceData: + async def async_get_data() -> dict: crawler = CovidCrawler(hass) - return await crawler.crawl_incidence() + return { + "incidence": await crawler.crawl_incidence(), + "vaccination": await crawler.crawl_vaccination(), + } hass.data[DOMAIN] = DataUpdateCoordinator( hass, diff --git a/custom_components/home_assistant_covid19_augsburg/sensor.py b/custom_components/home_assistant_covid19_augsburg/sensor.py index 516b50c..558022b 100644 --- a/custom_components/home_assistant_covid19_augsburg/sensor.py +++ b/custom_components/home_assistant_covid19_augsburg/sensor.py @@ -1,3 +1,5 @@ +from dataclasses import asdict + from homeassistant.helpers.entity import Entity from . import get_coordinator @@ -41,18 +43,59 @@ class CoronaAugsburgSensor(Entity): @property def state(self): - return self.coordinator.data.incidence + return self.coordinator.data["incidence"].incidence @property def device_state_attributes(self): - return { - "date": self.coordinator.data.date, - "incidence": self.coordinator.data.incidence, - "total_cases": self.coordinator.data.total_cases, - "num_dead": self.coordinator.data.num_dead, - "num_recovered": self.coordinator.data.num_recovered, - "num_infected": self.coordinator.data.num_infected, - } + data = self.coordinator.data["incidence"] + return asdict(data) + + 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) + + +class CoronaAugsburgVaccinationSensor(Entity): + """Representation of vaccination data for the city of Augsburg""" + + def __init__(self, coordinator): + """Initialize sensor.""" + self.coordinator = coordinator + self._name = "COVID-19 Vaccinations Augsburg" + self._state = None + + @property + def available(self): + return self.coordinator.last_update_success and self.coordinator.data + + @property + def name(self): + return self._name + + @property + def unique_id(self): + return self._name + + @property + def icon(self): + return "mdi:biohazard" + + @property + def unit_of_measurement(self): + return "" + + @property + def state(self): + return self.coordinator.data["vaccination"].total_vaccinations + + @property + def device_state_attributes(self): + data = self.coordinator.data["vaccination"] + return asdict(data) async def async_added_to_hass(self): """When entity is added to hass."""