Fix Aurora integration including externalizing API to PyPi and adding config_flow (#43045)

Co-authored-by: Pawel <pszafer@gmail.com>

.coveragerc

@@ -65,6 +65,9 @@ omit =
     homeassistant/components/asterisk_mbox/*
     homeassistant/components/aten_pe/*
     homeassistant/components/atome/*
+    homeassistant/components/aurora/__init__.py
+    homeassistant/components/aurora/binary_sensor.py
+    homeassistant/components/aurora/const.py
     homeassistant/components/aurora_abb_powerone/sensor.py
     homeassistant/components/avea/light.py
     homeassistant/components/avion/light.py

CODEOWNERS

@@ -48,6 +48,7 @@ homeassistant/components/atag/* @MatsNL
 homeassistant/components/aten_pe/* @mtdcr
 homeassistant/components/atome/* @baqs
 homeassistant/components/august/* @bdraco
+homeassistant/components/aurora/* @djtimca
 homeassistant/components/aurora_abb_powerone/* @davet2001
 homeassistant/components/auth/* @home-assistant/core
 homeassistant/components/automation/* @home-assistant/core

homeassistant/components/aurora/__init__.py

@@ -1 +1,130 @@
 """The aurora component."""
+
+import asyncio
+from datetime import timedelta
+import logging
+
+from auroranoaa import AuroraForecast
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import 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 .const import (
+    AURORA_API,
+    CONF_THRESHOLD,
+    COORDINATOR,
+    DEFAULT_POLLING_INTERVAL,
+    DEFAULT_THRESHOLD,
+    DOMAIN,
+)
+
+_LOGGER = logging.getLogger(__name__)
+
+PLATFORMS = ["binary_sensor"]
+
+
+async def async_setup(hass: HomeAssistant, config: dict):
+    """Set up the Aurora component."""
+    hass.data.setdefault(DOMAIN, {})
+
+    return True
+
+
+async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry):
+    """Set up Aurora from a config entry."""
+
+    conf = entry.data
+    options = entry.options
+
+    session = aiohttp_client.async_get_clientsession(hass)
+    api = AuroraForecast(session)
+
+    longitude = conf[CONF_LONGITUDE]
+    latitude = conf[CONF_LATITUDE]
+    polling_interval = DEFAULT_POLLING_INTERVAL
+    threshold = options.get(CONF_THRESHOLD, DEFAULT_THRESHOLD)
+    name = conf[CONF_NAME]
+
+    coordinator = AuroraDataUpdateCoordinator(
+        hass=hass,
+        name=name,
+        polling_interval=polling_interval,
+        api=api,
+        latitude=latitude,
+        longitude=longitude,
+        threshold=threshold,
+    )
+
+    await coordinator.async_refresh()
+
+    if not coordinator.last_update_success:
+        raise ConfigEntryNotReady
+
+    hass.data[DOMAIN][entry.entry_id] = {
+        COORDINATOR: coordinator,
+        AURORA_API: api,
+    }
+
+    for component in PLATFORMS:
+        hass.async_create_task(
+            hass.config_entries.async_forward_entry_setup(entry, component)
+        )
+
+    return True
+
+
+async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry):
+    """Unload a config entry."""
+    unload_ok = all(
+        await asyncio.gather(
+            *[
+                hass.config_entries.async_forward_entry_unload(entry, component)
+                for component in PLATFORMS
+            ]
+        )
+    )
+    if unload_ok:
+        hass.data[DOMAIN].pop(entry.entry_id)
+
+    return unload_ok
+
+
+class AuroraDataUpdateCoordinator(DataUpdateCoordinator):
+    """Class to manage fetching data from the NOAA Aurora API."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        name: str,
+        polling_interval: int,
+        api: str,
+        latitude: float,
+        longitude: float,
+        threshold: float,
+    ):
+        """Initialize the data updater."""
+
+        super().__init__(
+            hass=hass,
+            logger=_LOGGER,
+            name=name,
+            update_interval=timedelta(minutes=polling_interval),
+        )
+
+        self.api = api
+        self.name = name
+        self.latitude = int(latitude)
+        self.longitude = int(longitude)
+        self.threshold = int(threshold)
+
+    async def _async_update_data(self):
+        """Fetch the data from the NOAA Aurora Forecast."""
+
+        try:
+            return await self.api.get_forecast_data(self.longitude, self.latitude)
+        except ConnectionError as error:
+            raise UpdateFailed(f"Error updating from NOAA: {error}") from error

homeassistant/components/aurora/binary_sensor.py

@@ -1,146 +1,75 @@
 """Support for aurora forecast data sensor."""
-from datetime import timedelta
 import logging
-from math import floor
 
-from aiohttp.hdrs import USER_AGENT
+from homeassistant.components.binary_sensor import BinarySensorEntity
-import requests
+from homeassistant.const import ATTR_NAME
-import voluptuous as vol
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
-from homeassistant.components.binary_sensor import PLATFORM_SCHEMA, BinarySensorEntity
+from . import AuroraDataUpdateCoordinator
-from homeassistant.const import ATTR_ATTRIBUTION, CONF_NAME
+from .const import (
-import homeassistant.helpers.config_validation as cv
+    ATTR_IDENTIFIERS,
-from homeassistant.util import Throttle
+    ATTR_MANUFACTURER,
+    ATTR_MODEL,
+    ATTRIBUTION,
+    COORDINATOR,
+    DOMAIN,
+)
 
 _LOGGER = logging.getLogger(__name__)
 
-ATTRIBUTION = "Data provided by the National Oceanic and Atmospheric Administration"
-CONF_THRESHOLD = "forecast_threshold"
 
-DEFAULT_DEVICE_CLASS = "visible"
+async def async_setup_entry(hass, entry, async_add_entries):
-DEFAULT_NAME = "Aurora Visibility"
+    """Set up the binary_sensor platform."""
-DEFAULT_THRESHOLD = 75
+    coordinator = hass.data[DOMAIN][entry.entry_id][COORDINATOR]
+    name = coordinator.name
 
-HA_USER_AGENT = "Home Assistant Aurora Tracker v.0.1.0"
+    entity = AuroraSensor(coordinator, name)
 
-MIN_TIME_BETWEEN_UPDATES = timedelta(minutes=5)
+    async_add_entries([entity])
-
-URL = "http://services.swpc.noaa.gov/text/aurora-nowcast-map.txt"
-
-PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
-    {
-        vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,
-        vol.Optional(CONF_THRESHOLD, default=DEFAULT_THRESHOLD): cv.positive_int,
-    }
-)
 
 
-def setup_platform(hass, config, add_entities, discovery_info=None):
+class AuroraSensor(CoordinatorEntity, BinarySensorEntity):
-    """Set up the aurora sensor."""
-    if None in (hass.config.latitude, hass.config.longitude):
-        _LOGGER.error("Lat. or long. not set in Home Assistant config")
-        return False
-
-    name = config[CONF_NAME]
-    threshold = config[CONF_THRESHOLD]
-
-    try:
-        aurora_data = AuroraData(hass.config.latitude, hass.config.longitude, threshold)
-        aurora_data.update()
-    except requests.exceptions.HTTPError as error:
-        _LOGGER.error("Connection to aurora forecast service failed: %s", error)
-        return False
-
-    add_entities([AuroraSensor(aurora_data, name)], True)
-
-
-class AuroraSensor(BinarySensorEntity):
     """Implementation of an aurora sensor."""
 
-    def __init__(self, aurora_data, name):
+    def __init__(self, coordinator: AuroraDataUpdateCoordinator, name):
-        """Initialize the sensor."""
+        """Define the binary sensor for the Aurora integration."""
-        self.aurora_data = aurora_data
+        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 f"{self._name}"
+        return self._name
 
     @property
     def is_on(self):
         """Return true if aurora is visible."""
-        return self.aurora_data.is_visible if self.aurora_data else False
+        return self.coordinator.data > self.coordinator.threshold
-
-    @property
-    def device_class(self):
-        """Return the class of this device."""
-        return DEFAULT_DEVICE_CLASS
 
     @property
     def device_state_attributes(self):
         """Return the state attributes."""
-        attrs = {}
+        return {"attribution": ATTRIBUTION}
 
-        if self.aurora_data:
+    @property
-            attrs["visibility_level"] = self.aurora_data.visibility_level
+    def icon(self):
-            attrs["message"] = self.aurora_data.is_visible_text
+        """Return the icon for the sensor."""
-            attrs[ATTR_ATTRIBUTION] = ATTRIBUTION
+        return "mdi:hazard-lights"
-        return attrs
 
-    def update(self):
+    @property
-        """Get the latest data from Aurora API and updates the states."""
+    def device_info(self):
-        self.aurora_data.update()
+        """Define the device based on name."""
-
+        return {
-
+            ATTR_IDENTIFIERS: {(DOMAIN, self._unique_id)},
-class AuroraData:
+            ATTR_NAME: self.coordinator.name,
-    """Get aurora forecast."""
+            ATTR_MANUFACTURER: "NOAA",
-
+            ATTR_MODEL: "Aurora Visibility Sensor",
-    def __init__(self, latitude, longitude, threshold):
+        }
-        """Initialize the data object."""
-        self.latitude = latitude
-        self.longitude = longitude
-        self.headers = {USER_AGENT: HA_USER_AGENT}
-        self.threshold = int(threshold)
-        self.is_visible = None
-        self.is_visible_text = None
-        self.visibility_level = None
-
-    @Throttle(MIN_TIME_BETWEEN_UPDATES)
-    def update(self):
-        """Get the latest data from the Aurora service."""
-        try:
-            self.visibility_level = self.get_aurora_forecast()
-            if int(self.visibility_level) > self.threshold:
-                self.is_visible = True
-                self.is_visible_text = "visible!"
-            else:
-                self.is_visible = False
-                self.is_visible_text = "nothing's out"
-
-        except requests.exceptions.HTTPError as error:
-            _LOGGER.error("Connection to aurora forecast service failed: %s", error)
-            return False
-
-    def get_aurora_forecast(self):
-        """Get forecast data and parse for given long/lat."""
-        raw_data = requests.get(URL, headers=self.headers, timeout=5).text
-        # We discard comment rows (#)
-        # We split the raw text by line (\n)
-        # For each line we trim leading spaces and split by spaces
-        forecast_table = [
-            row.strip().split()
-            for row in raw_data.split("\n")
-            if not row.startswith("#")
-        ]
-
-        # Convert lat and long for data points in table
-        # Assumes self.latitude belongs to [-90;90[ (South to North)
-        # Assumes self.longitude belongs to [-180;180[ (West to East)
-        # No assumptions made regarding the number of rows and columns
-        converted_latitude = floor((self.latitude + 90) * len(forecast_table) / 180)
-        converted_longitude = floor(
-            (self.longitude + 180) * len(forecast_table[converted_latitude]) / 360
-        )
-
-        return forecast_table[converted_latitude][converted_longitude]

homeassistant/components/aurora/config_flow.py

@@ -0,0 +1,110 @@
+"""Config flow for SpaceX Launches and Starman."""
+import logging
+
+from auroranoaa import AuroraForecast
+import voluptuous as vol
+
+from homeassistant import config_entries
+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
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
+    """Handle a config flow for NOAA Aurora Integration."""
+
+    VERSION = 1
+    CONNECTION_CLASS = config_entries.CONN_CLASS_CLOUD_POLL
+
+    @staticmethod
+    @callback
+    def async_get_options_flow(config_entry):
+        """Get the options flow for this handler."""
+        return OptionsFlowHandler(config_entry)
+
+    async def async_step_user(self, user_input=None):
+        """Handle the initial step."""
+        errors = {}
+
+        if user_input is not None:
+            name = user_input[CONF_NAME]
+            longitude = user_input[CONF_LONGITUDE]
+            latitude = user_input[CONF_LATITUDE]
+
+            session = aiohttp_client.async_get_clientsession(self.hass)
+            api = AuroraForecast(session=session)
+
+            try:
+                await api.get_forecast_data(longitude, latitude)
+            except ConnectionError:
+                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]}"
+                )
+                self._abort_if_unique_id_configured()
+                return self.async_create_entry(
+                    title=f"Aurora - {name}", data=user_input
+                )
+
+        return self.async_show_form(
+            step_id="user",
+            data_schema=vol.Schema(
+                {
+                    vol.Required(CONF_NAME, default=DEFAULT_NAME): str,
+                    vol.Required(
+                        CONF_LONGITUDE,
+                        default=self.hass.config.longitude,
+                    ): vol.All(
+                        vol.Coerce(float),
+                        vol.Range(min=-180, max=180),
+                    ),
+                    vol.Required(
+                        CONF_LATITUDE,
+                        default=self.hass.config.latitude,
+                    ): vol.All(
+                        vol.Coerce(float),
+                        vol.Range(min=-90, max=90),
+                    ),
+                }
+            ),
+            errors=errors,
+        )
+
+
+class OptionsFlowHandler(config_entries.OptionsFlow):
+    """Handle options flow changes."""
+
+    def __init__(self, config_entry):
+        """Initialize options flow."""
+        self.config_entry = config_entry
+
+    async def async_step_init(self, user_input=None):
+        """Manage options."""
+
+        if user_input is not None:
+            return self.async_create_entry(title="", data=user_input)
+
+        return self.async_show_form(
+            step_id="init",
+            data_schema=vol.Schema(
+                {
+                    vol.Required(
+                        CONF_THRESHOLD,
+                        default=self.config_entry.options.get(
+                            CONF_THRESHOLD, DEFAULT_THRESHOLD
+                        ),
+                    ): vol.All(
+                        vol.Coerce(int),
+                        vol.Range(min=0, max=100),
+                    ),
+                }
+            ),
+        )

homeassistant/components/aurora/const.py

@@ -0,0 +1,13 @@
+"""Constants for the Aurora integration."""
+
+DOMAIN = "aurora"
+COORDINATOR = "coordinator"
+AURORA_API = "aurora_api"
+ATTR_IDENTIFIERS = "identifiers"
+ATTR_MANUFACTURER = "manufacturer"
+ATTR_MODEL = "model"
+DEFAULT_POLLING_INTERVAL = 5
+CONF_THRESHOLD = "forecast_threshold"
+DEFAULT_THRESHOLD = 75
+ATTRIBUTION = "Data provided by the National Oceanic and Atmospheric Administration"
+DEFAULT_NAME = "Aurora Visibility"

homeassistant/components/aurora/manifest.json

@@ -2,5 +2,7 @@
   "domain": "aurora",
   "name": "Aurora",
   "documentation": "https://www.home-assistant.io/integrations/aurora",
-  "codeowners": []
+  "config_flow": true,
+  "codeowners": ["@djtimca"],
+  "requirements": ["auroranoaa==0.0.1"]
 }

homeassistant/components/aurora/strings.json

@@ -0,0 +1,26 @@
+{
+    "title": "NOAA Aurora Sensor",
+    "config": {
+      "step": {
+        "user": {
+          "data": {
+            "name": "[%key:common::config_flow::data::name%]",
+            "longitude": "[%key:common::config_flow::data::longitude%]",
+            "latitude": "[%key:common::config_flow::data::latitude%]"
+          }
+        }
+      },
+      "error": {
+        "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]"
+      }
+    },
+    "options": {
+      "step": {
+        "init": {
+          "data": {
+            "threshold": "Threshold (%)"
+          }
+        }
+      }
+    }
+  }

homeassistant/components/aurora/translations/en.json

@@ -0,0 +1,26 @@
+{
+    "config": {
+        "error": {
+            "cannot_connect": "Failed to connect"
+        },
+        "step": {
+            "user": {
+                "data": {
+                    "latitude": "Latitude",
+                    "longitude": "Longitude",
+                    "name": "Name"
+                }
+            }
+        }
+    },
+    "options": {
+        "step": {
+            "init": {
+                "data": {
+                    "threshold": "Threshold (%)"
+                }
+            }
+        }
+    },
+    "title": "NOAA Aurora Sensor"
+}

homeassistant/generated/config_flows.py

@@ -21,6 +21,7 @@ FLOWS = [
     "arcam_fmj",
     "atag",
     "august",
+    "aurora",
     "avri",
     "awair",
     "axis",

requirements_all.txt

@@ -296,6 +296,9 @@ asyncpysupla==0.0.5
 # homeassistant.components.aten_pe
 atenpdu==0.3.0
 
+# homeassistant.components.aurora
+auroranoaa==0.0.1
+
 # homeassistant.components.aurora_abb_powerone
 aurorapy==0.2.6
 

requirements_test_all.txt

@@ -173,6 +173,9 @@ arcam-fmj==0.5.3
 # homeassistant.components.upnp
 async-upnp-client==0.14.13
 
+# homeassistant.components.aurora
+auroranoaa==0.0.1
+
 # homeassistant.components.stream
 av==8.0.2
 

tests/components/aurora/test_binary_sensor.py

@@ -1,60 +0,0 @@
-"""The tests for the Aurora sensor platform."""
-import re
-
-from homeassistant.components.aurora import binary_sensor as aurora
-
-from tests.common import load_fixture
-
-
-def test_setup_and_initial_state(hass, requests_mock):
-    """Test that the component is created and initialized as expected."""
-    uri = re.compile(r"http://services\.swpc\.noaa\.gov/text/aurora-nowcast-map\.txt")
-    requests_mock.get(uri, text=load_fixture("aurora.txt"))
-
-    entities = []
-
-    def mock_add_entities(new_entities, update_before_add=False):
-        """Mock add entities."""
-        if update_before_add:
-            for entity in new_entities:
-                entity.update()
-
-        for entity in new_entities:
-            entities.append(entity)
-
-    config = {"name": "Test", "forecast_threshold": 75}
-    aurora.setup_platform(hass, config, mock_add_entities)
-
-    aurora_component = entities[0]
-    assert len(entities) == 1
-    assert aurora_component.name == "Test"
-    assert aurora_component.device_state_attributes["visibility_level"] == "0"
-    assert aurora_component.device_state_attributes["message"] == "nothing's out"
-    assert not aurora_component.is_on
-
-
-def test_custom_threshold_works(hass, requests_mock):
-    """Test that the config can take a custom forecast threshold."""
-    uri = re.compile(r"http://services\.swpc\.noaa\.gov/text/aurora-nowcast-map\.txt")
-    requests_mock.get(uri, text=load_fixture("aurora.txt"))
-
-    entities = []
-
-    def mock_add_entities(new_entities, update_before_add=False):
-        """Mock add entities."""
-        if update_before_add:
-            for entity in new_entities:
-                entity.update()
-
-        for entity in new_entities:
-            entities.append(entity)
-
-    config = {"name": "Test", "forecast_threshold": 1}
-    hass.config.longitude = 18.987
-    hass.config.latitude = 69.648
-
-    aurora.setup_platform(hass, config, mock_add_entities)
-
-    aurora_component = entities[0]
-    assert aurora_component.aurora_data.visibility_level == "16"
-    assert aurora_component.is_on

tests/components/aurora/test_config_flow.py

@@ -0,0 +1,113 @@
+"""Test the Aurora config flow."""
+
+from homeassistant import config_entries, data_entry_flow, setup
+from homeassistant.components.aurora.const import DOMAIN
+
+from tests.async_mock import patch
+from tests.common import MockConfigEntry
+
+DATA = {
+    "name": "Home",
+    "latitude": -10,
+    "longitude": 10.2,
+}
+
+
+async def test_form(hass):
+    """Test we get the form."""
+    await setup.async_setup_component(hass, "persistent_notification", {})
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    assert result["type"] == "form"
+    assert result["errors"] == {}
+
+    with patch(
+        "homeassistant.components.aurora.config_flow.AuroraForecast.get_forecast_data",
+        return_value=True,
+    ), patch(
+        "homeassistant.components.aurora.async_setup", return_value=True
+    ) as mock_setup, patch(
+        "homeassistant.components.aurora.async_setup_entry",
+        return_value=True,
+    ) as mock_setup_entry:
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            DATA,
+        )
+
+    assert result2["type"] == "create_entry"
+    assert result2["title"] == "Aurora - Home"
+    assert result2["data"] == DATA
+    await hass.async_block_till_done()
+    assert len(mock_setup.mock_calls) == 1
+    assert len(mock_setup_entry.mock_calls) == 1
+
+
+async def test_form_cannot_connect(hass):
+    """Test if invalid response or no connection returned from the API."""
+
+    await setup.async_setup_component(hass, "persistent_notification", {})
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+
+    with patch(
+        "homeassistant.components.aurora.AuroraForecast.get_forecast_data",
+        side_effect=ConnectionError,
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            DATA,
+        )
+
+    assert result["type"] == "form"
+    assert result["step_id"] == "user"
+    assert result["errors"] == {"base": "cannot_connect"}
+
+
+async def test_with_unknown_error(hass):
+    """Test with unknown error response from the API."""
+    await setup.async_setup_component(hass, "persistent_notification", {})
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+
+    with patch(
+        "homeassistant.components.aurora.AuroraForecast.get_forecast_data",
+        side_effect=Exception,
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            DATA,
+        )
+
+    assert result["type"] == "form"
+    assert result["step_id"] == "user"
+    assert result["errors"] == {"base": "unknown"}
+
+
+async def test_option_flow(hass):
+    """Test option flow."""
+    entry = MockConfigEntry(domain=DOMAIN, data=DATA)
+    entry.add_to_hass(hass)
+
+    assert not entry.options
+
+    with patch("homeassistant.components.aurora.async_setup_entry", return_value=True):
+        result = await hass.config_entries.options.async_init(
+            entry.entry_id,
+            data=None,
+        )
+
+    assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
+    assert result["step_id"] == "init"
+
+    result = await hass.config_entries.options.async_configure(
+        result["flow_id"],
+        user_input={"forecast_threshold": 65},
+    )
+
+    assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
+    assert result["title"] == ""
+    assert result["data"]["forecast_threshold"] == 65