Delivery site selection in config flow

custom_components/frank_energie/__init__.py

@@ -1,4 +1,4 @@
-"""The Frank Energie component."""
+"""The Frank Energie integration."""
 from __future__ import annotations
 
 from homeassistant.config_entries import ConfigEntry
@@ -14,36 +14,12 @@
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
-    """Set up the Frank Energie component from a config entry."""
+    """Set up the Frank Energie integration from a config entry."""
 
     # For backwards compatibility, set unique ID
     if entry.unique_id is None or entry.unique_id == "frank_energie_component":
         hass.config_entries.async_update_entry(entry, unique_id=str("frank_energie"))
 
-    # Select site-reference, or find first one that has status 'IN_DELIVERY' if not set
-    if entry.data.get("site_reference") is None and entry.data.get(CONF_ACCESS_TOKEN) is not None:
-        api = FrankEnergie(
-            clientsession=async_get_clientsession(hass),
-            auth_token=entry.data.get(CONF_ACCESS_TOKEN, None),
-            refresh_token=entry.data.get(CONF_TOKEN, None),
-        )
-        me = await api.me()
-
-        # filter out all sites that are not in delivery
-        me.deliverySites = [site for site in me.deliverySites if site.status == "IN_DELIVERY"]
-
-        if len(me.deliverySites) == 0:
-            raise Exception("No suitable sites found for this account")
-
-        site = me.deliverySites[0]
-        hass.config_entries.async_update_entry(entry, data={**entry.data, "site_reference": site.reference})
-
-        # Update title
-        title = f"{site.address_street} {site.address_houseNumber}"
-        if site.address_houseNumberAddition is not None:
-            title += f" {site.address_houseNumberAddition}"
-        hass.config_entries.async_update_entry(entry, title=title)
-
     # Initialise the coordinator and save it as domain-data
     api = FrankEnergie(
         clientsession=async_get_clientsession(hass),

custom_components/frank_energie/config_flow.py

@@ -1,4 +1,4 @@
-"""Config flow for Picnic integration."""
+"""Config flow for Frank Energie integration."""
 from __future__ import annotations
 
 import logging
@@ -15,10 +15,11 @@
     CONF_USERNAME,
 )
 from homeassistant.data_entry_flow import FlowResult
+from homeassistant.helpers.selector import SelectSelector, SelectSelectorConfig, SelectSelectorMode
 from python_frank_energie import FrankEnergie
 from python_frank_energie.exceptions import AuthException
 
-from .const import DOMAIN
+from .const import DOMAIN, CONF_SITE
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -27,6 +28,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
     """Handle the config flow for Frank Energie."""
 
     VERSION = 1
+    sign_in_data = {}
 
     def __init__(self) -> None:
         """Initialize the config flow."""
@@ -61,16 +63,16 @@ async def async_step_login(self, user_input=None, errors=None) -> FlowResult:
                 _LOGGER.exception("Error during login", exc_info=ex)
                 return await self.async_step_login(errors={"base": "invalid_auth"})
 
-        data = {
+        self.sign_in_data = {
             CONF_USERNAME: user_input[CONF_USERNAME],
             CONF_ACCESS_TOKEN: auth.authToken,
-            CONF_TOKEN: auth.refreshToken,
+            CONF_TOKEN: auth.refreshToken
         }
 
         if self._reauth_entry:
             self.hass.config_entries.async_update_entry(
                 self._reauth_entry,
-                data=data,
+                data=self.sign_in_data,
             )
 
             self.hass.async_create_task(
@@ -79,10 +81,38 @@ async def async_step_login(self, user_input=None, errors=None) -> FlowResult:
 
             return self.async_abort(reason="reauth_successful")
 
-        await self.async_set_unique_id(user_input[CONF_USERNAME])
-        self._abort_if_unique_id_configured()
+        return await self.async_step_site(self.sign_in_data)
 
-        return await self._async_create_entry(data)
+    async def async_step_site(self, user_input=None, errors=None) -> FlowResult:
+        """Handle possible multi site accounts."""
+        if user_input.get(CONF_SITE) is None:
+            api = FrankEnergie(
+                auth_token=self.sign_in_data.get(CONF_ACCESS_TOKEN, None),
+                refresh_token=self.sign_in_data.get(CONF_TOKEN, None),
+            )
+            me = await api.me()
+
+            # filter out all sites that are not in delivery
+            me.deliverySites = [site for site in me.deliverySites if site.status == "IN_DELIVERY"]
+
+            if len(me.deliverySites) == 0:
+                raise Exception("No suitable sites found for this account")
+
+            site_options = [{"value": site.reference, "label": self.create_title(site)} for site in me.deliverySites]
+            default_site = me.deliverySites[0].reference
+
+            options = {vol.Required(CONF_SITE, default=default_site): SelectSelector(SelectSelectorConfig(
+                options=site_options,
+                mode=SelectSelectorMode.LIST,
+            ))}
+
+            return self.async_show_form(
+                step_id="site", data_schema=vol.Schema(options), errors=errors
+            )
+
+        self.sign_in_data[CONF_SITE] = user_input[CONF_SITE]
+
+        return await self._async_create_entry(self.sign_in_data)
 
     async def async_step_user(self, user_input=None) -> FlowResult:
         """Handle a flow initiated by the user."""
@@ -110,9 +140,19 @@ async def async_step_reauth(self, entry_data: Mapping[str, Any]) -> FlowResult:
         return await self.async_step_login()
 
     async def _async_create_entry(self, data):
-        await self.async_set_unique_id(data.get(CONF_USERNAME, "frank_energie"))
+        unique_id = data[CONF_SITE] + data[CONF_USERNAME]
+        await self.async_set_unique_id(unique_id)
+        # await self.async_set_unique_id(data.get(CONF_USERNAME, "frank_energie"))
         self._abort_if_unique_id_configured()
 
         return self.async_create_entry(
             title=data.get(CONF_USERNAME, "Frank Energie"), data=data
         )
+
+    @staticmethod
+    def create_title(site) -> str:
+        title = f"{site.address_street} {site.address_houseNumber}"
+        if site.address_houseNumberAddition is not None:
+            title += f" {site.address_houseNumberAddition}"
+
+        return title

custom_components/frank_energie/const.py

@@ -8,6 +8,7 @@
 COMPONENT_TITLE = "Frank Energie"
 
 CONF_COORDINATOR = "coordinator"
+CONF_SITE = "site_reference"
 ATTR_TIME = "from_time"
 
 DATA_ELECTRICITY = "electricity"

requirements.txt

@@ -1,4 +1,4 @@
-pytest-homeassistant-custom-component~=0.13.7
+pytest-homeassistant-custom-component~=0.13.134
 flake8~=5.0.4
-pytest~=7.2.1
-python-frank-energie~=4.1.0
+pytest~=8.2.0
+python-frank-energie~=6.0.0