Add sensor platform for Aurora integration (#43148)

* Removed pylint disable on unused after updating CI files that were out of date.
* Pylint still fails due to bug on DOMAIN import. Added disable check.
* Addressed PR comments
* Added import for ClientError to test_config_flow.py

Co-authored-by: Paulus Schoutsen <balloob@gmail.com>

.coveragerc

@@ -72,6 +72,7 @@ omit =
     homeassistant/components/aurora/__init__.py
     homeassistant/components/aurora/binary_sensor.py
     homeassistant/components/aurora/const.py
+    homeassistant/components/aurora/sensor.py
     homeassistant/components/aurora_abb_powerone/sensor.py
     homeassistant/components/avea/light.py
     homeassistant/components/avion/light.py

homeassistant/components/aurora/__init__.py

@@ -4,16 +4,26 @@ import asyncio
 from datetime import timedelta
 import logging
 
+from aiohttp import ClientError
 from auroranoaa import AuroraForecast
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_LATITUDE, CONF_LONGITUDE, CONF_NAME
+from homeassistant.const import ATTR_NAME, CONF_LATITUDE, CONF_LONGITUDE, CONF_NAME
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import aiohttp_client
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+from homeassistant.helpers.update_coordinator import (
+    CoordinatorEntity,
+    DataUpdateCoordinator,
+    UpdateFailed,
+)
 
 from .const import (
+    ATTR_ENTRY_TYPE,
+    ATTR_IDENTIFIERS,
+    ATTR_MANUFACTURER,
+    ATTR_MODEL,
+    ATTRIBUTION,
     AURORA_API,
     CONF_THRESHOLD,
     COORDINATOR,
@@ -24,7 +34,7 @@ from .const import (
 
 _LOGGER = logging.getLogger(__name__)
 
-PLATFORMS = ["binary_sensor"]
+PLATFORMS = ["binary_sensor", "sensor"]
 
 
 async def async_setup(hass: HomeAssistant, config: dict):
@@ -126,5 +136,54 @@ class AuroraDataUpdateCoordinator(DataUpdateCoordinator):
 
         try:
             return await self.api.get_forecast_data(self.longitude, self.latitude)
-        except ConnectionError as error:
+        except ClientError as error:
             raise UpdateFailed(f"Error updating from NOAA: {error}") from error
+
+
+class AuroraEntity(CoordinatorEntity):
+    """Implementation of the base Aurora Entity."""
+
+    def __init__(
+        self,
+        coordinator: AuroraDataUpdateCoordinator,
+        name: str,
+        icon: str,
+    ):
+        """Initialize the Aurora Entity."""
+
+        super().__init__(coordinator=coordinator)
+
+        self._name = name
+        self._unique_id = f"{self.coordinator.latitude}_{self.coordinator.longitude}"
+        self._icon = icon
+
+    @property
+    def unique_id(self):
+        """Define the unique id based on the latitude and longitude."""
+        return self._unique_id
+
+    @property
+    def name(self):
+        """Return the name of the sensor."""
+        return self._name
+
+    @property
+    def device_state_attributes(self):
+        """Return the state attributes."""
+        return {"attribution": ATTRIBUTION}
+
+    @property
+    def icon(self):
+        """Return the icon for the sensor."""
+        return self._icon
+
+    @property
+    def device_info(self):
+        """Define the device based on name."""
+        return {
+            ATTR_IDENTIFIERS: {(DOMAIN, self._unique_id)},
+            ATTR_NAME: self.coordinator.name,
+            ATTR_MANUFACTURER: "NOAA",
+            ATTR_MODEL: "Aurora Visibility Sensor",
+            ATTR_ENTRY_TYPE: "service",
+        }

homeassistant/components/aurora/binary_sensor.py

@@ -1,19 +1,10 @@
-"""Support for aurora forecast data sensor."""
+"""Support for Aurora Forecast binary sensor."""
 import logging
 
 from homeassistant.components.binary_sensor import BinarySensorEntity
-from homeassistant.const import ATTR_NAME
-from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from . import AuroraDataUpdateCoordinator
+from . import AuroraEntity
-from .const import (
+from .const import COORDINATOR, DOMAIN
-    ATTR_IDENTIFIERS,
-    ATTR_MANUFACTURER,
-    ATTR_MODEL,
-    ATTRIBUTION,
-    COORDINATOR,
-    DOMAIN,
-)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -21,55 +12,17 @@ _LOGGER = logging.getLogger(__name__)
 async def async_setup_entry(hass, entry, async_add_entries):
     """Set up the binary_sensor platform."""
     coordinator = hass.data[DOMAIN][entry.entry_id][COORDINATOR]
-    name = coordinator.name
+    name = f"{coordinator.name} Aurora Visibility Alert"
 
-    entity = AuroraSensor(coordinator, name)
+    entity = AuroraSensor(coordinator=coordinator, name=name, icon="mdi:hazard-lights")
 
     async_add_entries([entity])
 
 
-class AuroraSensor(CoordinatorEntity, BinarySensorEntity):
+class AuroraSensor(AuroraEntity, BinarySensorEntity):
     """Implementation of an aurora sensor."""
 
-    def __init__(self, coordinator: AuroraDataUpdateCoordinator, name):
-        """Define the binary sensor for the Aurora integration."""
-        super().__init__(coordinator=coordinator)
-
-        self._name = name
-        self.coordinator = coordinator
-        self._unique_id = f"{self.coordinator.latitude}_{self.coordinator.longitude}"
-
-    @property
-    def unique_id(self):
-        """Define the unique id based on the latitude and longitude."""
-        return self._unique_id
-
-    @property
-    def name(self):
-        """Return the name of the sensor."""
-        return self._name
-
     @property
     def is_on(self):
         """Return true if aurora is visible."""
         return self.coordinator.data > self.coordinator.threshold
-
-    @property
-    def device_state_attributes(self):
-        """Return the state attributes."""
-        return {"attribution": ATTRIBUTION}
-
-    @property
-    def icon(self):
-        """Return the icon for the sensor."""
-        return "mdi:hazard-lights"
-
-    @property
-    def device_info(self):
-        """Define the device based on name."""
-        return {
-            ATTR_IDENTIFIERS: {(DOMAIN, self._unique_id)},
-            ATTR_NAME: self.coordinator.name,
-            ATTR_MANUFACTURER: "NOAA",
-            ATTR_MODEL: "Aurora Visibility Sensor",
-        }

homeassistant/components/aurora/config_flow.py

@@ -1,6 +1,7 @@
 """Config flow for SpaceX Launches and Starman."""
 import logging
 
+from aiohttp import ClientError
 from auroranoaa import AuroraForecast
 import voluptuous as vol
 
@@ -9,7 +10,12 @@ from homeassistant.const import CONF_LATITUDE, CONF_LONGITUDE, CONF_NAME
 from homeassistant.core import callback
 from homeassistant.helpers import aiohttp_client
 
-from .const import CONF_THRESHOLD, DEFAULT_NAME, DEFAULT_THRESHOLD, DOMAIN
+from .const import (  # pylint: disable=unused-import
+    CONF_THRESHOLD,
+    DEFAULT_NAME,
+    DEFAULT_THRESHOLD,
+    DOMAIN,
+)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -40,14 +46,14 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
 
             try:
                 await api.get_forecast_data(longitude, latitude)
-            except ConnectionError:
+            except ClientError:
                 errors["base"] = "cannot_connect"
             except Exception:  # pylint: disable=broad-except
                 _LOGGER.exception("Unexpected exception")
                 errors["base"] = "unknown"
             else:
                 await self.async_set_unique_id(
-                    f"{DOMAIN}_{user_input[CONF_LONGITUDE]}_{user_input[CONF_LATITUDE]}"
+                    f"{user_input[CONF_LONGITUDE]}_{user_input[CONF_LATITUDE]}"
                 )
                 self._abort_if_unique_id_configured()
                 return self.async_create_entry(

homeassistant/components/aurora/const.py

@@ -6,6 +6,7 @@ AURORA_API = "aurora_api"
 ATTR_IDENTIFIERS = "identifiers"
 ATTR_MANUFACTURER = "manufacturer"
 ATTR_MODEL = "model"
+ATTR_ENTRY_TYPE = "entry_type"
 DEFAULT_POLLING_INTERVAL = 5
 CONF_THRESHOLD = "forecast_threshold"
 DEFAULT_THRESHOLD = 75

homeassistant/components/aurora/sensor.py

@@ -0,0 +1,36 @@
+"""Support for Aurora Forecast sensor."""
+import logging
+
+from homeassistant.const import PERCENTAGE
+
+from . import AuroraEntity
+from .const import COORDINATOR, DOMAIN
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def async_setup_entry(hass, entry, async_add_entries):
+    """Set up the sensor platform."""
+    coordinator = hass.data[DOMAIN][entry.entry_id][COORDINATOR]
+
+    entity = AuroraSensor(
+        coordinator=coordinator,
+        name=f"{coordinator.name} Aurora Visibility %",
+        icon="mdi:gauge",
+    )
+
+    async_add_entries([entity])
+
+
+class AuroraSensor(AuroraEntity):
+    """Implementation of an aurora sensor."""
+
+    @property
+    def state(self):
+        """Return % chance the aurora is visible."""
+        return self.coordinator.data
+
+    @property
+    def unit_of_measurement(self):
+        """Return the unit of measure."""
+        return PERCENTAGE

tests/components/aurora/test_config_flow.py

@@ -2,6 +2,8 @@
 
 from unittest.mock import patch
 
+from aiohttp import ClientError
+
 from homeassistant import config_entries, data_entry_flow, setup
 from homeassistant.components.aurora.const import DOMAIN
 
@@ -55,7 +57,7 @@ async def test_form_cannot_connect(hass):
 
     with patch(
         "homeassistant.components.aurora.AuroraForecast.get_forecast_data",
-        side_effect=ConnectionError,
+        side_effect=ClientError,
     ):
         result = await hass.config_entries.flow.async_configure(
             result["flow_id"],