Add Reauth flow to Wallbox integration

homeassistant/components/wallbox/__init__.py

@@ -9,16 +9,10 @@
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PASSWORD, CONF_USERNAME
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import HomeAssistantError
+from homeassistant.exceptions import ConfigEntryAuthFailed, HomeAssistantError
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
 
-from .const import (
-    CONF_CONNECTIONS,
-    CONF_DATA_KEY,
-    CONF_MAX_CHARGING_CURRENT_KEY,
-    CONF_STATION,
-    DOMAIN,
-)
+from .const import CONF_DATA_KEY, CONF_MAX_CHARGING_CURRENT_KEY, CONF_STATION, DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -48,7 +42,7 @@ def _authenticate(self):
             return True
         except requests.exceptions.HTTPError as wallbox_connection_error:
             if wallbox_connection_error.response.status_code == HTTPStatus.FORBIDDEN:
-                raise InvalidAuth from wallbox_connection_error
+                raise ConfigEntryAuthFailed from wallbox_connection_error
             raise ConnectionError from wallbox_connection_error
 
     def _validate(self):
@@ -112,18 +106,17 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         hass,
     )
 
-    await wallbox_coordinator.async_validate_input()
+    try:
+        await wallbox_coordinator.async_validate_input()
 
-    await wallbox_coordinator.async_config_entry_first_refresh()
+    except InvalidAuth as ex:
+        raise ConfigEntryAuthFailed from ex
 
-    hass.data.setdefault(DOMAIN, {CONF_CONNECTIONS: {}})
-    hass.data[DOMAIN][CONF_CONNECTIONS][entry.entry_id] = wallbox_coordinator
+    await wallbox_coordinator.async_config_entry_first_refresh()
 
-    for platform in PLATFORMS:
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = wallbox_coordinator
 
-        hass.async_create_task(
-            hass.config_entries.async_forward_entry_setup(entry, platform)
-        )
+    hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
     return True
 
@@ -132,7 +125,7 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
     if unload_ok:
-        hass.data[DOMAIN][CONF_CONNECTIONS].pop(entry.entry_id)
+        hass.data[DOMAIN].pop(entry.entry_id)
 
     return unload_ok
 

homeassistant/components/wallbox/config_flow.py

@@ -36,6 +36,18 @@ async def validate_input(hass: core.HomeAssistant, data):
 class ConfigFlow(config_entries.ConfigFlow, domain=COMPONENT_DOMAIN):
     """Handle a config flow for Wallbox."""
 
+    def __init__(self):
+        """Start the Wallbox config flow."""
+        self._reauth_entry = None
+
+    async def async_step_reauth(self, user_input=None):
+        """Perform reauth upon an API authentication error."""
+        self._reauth_entry = self.hass.config_entries.async_get_entry(
+            self.context["entry_id"]
+        )
+
+        return await self.async_step_user()
+
     async def async_step_user(self, user_input=None):
         """Handle the initial step."""
         if user_input is None:
@@ -47,14 +59,27 @@ async def async_step_user(self, user_input=None):
         errors = {}
 
         try:
-            info = await validate_input(self.hass, user_input)
+            await self.async_set_unique_id(user_input["station"])
+            if not self._reauth_entry:
+                self._abort_if_unique_id_configured()
+                info = await validate_input(self.hass, user_input)
+                return self.async_create_entry(title=info["title"], data=user_input)
+            if user_input["station"] == self._reauth_entry.data[CONF_STATION]:
+                self.hass.config_entries.async_update_entry(
+                    self._reauth_entry, data=user_input, unique_id=user_input["station"]
+                )
+                self.hass.async_create_task(
+                    self.hass.config_entries.async_reload(self._reauth_entry.entry_id)
+                )
+                return self.async_abort(reason="reauth_successful")
+            errors["base"] = "reauth_invalid"
         except ConnectionError:
             errors["base"] = "cannot_connect"
         except InvalidAuth:
             errors["base"] = "invalid_auth"
-        else:
-            return self.async_create_entry(title=info["title"], data=user_input)
 
         return self.async_show_form(
-            step_id="user", data_schema=STEP_USER_DATA_SCHEMA, errors=errors
+            step_id="user",
+            data_schema=STEP_USER_DATA_SCHEMA,
+            errors=errors,
         )

homeassistant/components/wallbox/number.py

@@ -8,12 +8,7 @@
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import InvalidAuth
-from .const import (
-    CONF_CONNECTIONS,
-    CONF_MAX_AVAILABLE_POWER_KEY,
-    CONF_MAX_CHARGING_CURRENT_KEY,
-    DOMAIN,
-)
+from .const import CONF_MAX_AVAILABLE_POWER_KEY, CONF_MAX_CHARGING_CURRENT_KEY, DOMAIN
 
 
 @dataclass
@@ -35,8 +30,7 @@ class WallboxNumberEntityDescription(NumberEntityDescription):
 
 async def async_setup_entry(hass, config, async_add_entities):
     """Create wallbox sensor entities in HASS."""
-    coordinator = hass.data[DOMAIN][CONF_CONNECTIONS][config.entry_id]
-
+    coordinator = hass.data[DOMAIN][config.entry_id]
     # Check if the user is authorized to change current, if so, add number component:
     try:
         await coordinator.async_set_charging_current(

homeassistant/components/wallbox/sensor.py

@@ -23,7 +23,6 @@
     CONF_ADDED_RANGE_KEY,
     CONF_CHARGING_POWER_KEY,
     CONF_CHARGING_SPEED_KEY,
-    CONF_CONNECTIONS,
     CONF_COST_KEY,
     CONF_CURRENT_MODE_KEY,
     CONF_DEPOT_PRICE_KEY,
@@ -120,7 +119,7 @@ class WallboxSensorEntityDescription(SensorEntityDescription):
 
 async def async_setup_entry(hass, config, async_add_entities):
     """Create wallbox sensor entities in HASS."""
-    coordinator = hass.data[DOMAIN][CONF_CONNECTIONS][config.entry_id]
+    coordinator = hass.data[DOMAIN][config.entry_id]
 
     async_add_entities(
         [

homeassistant/components/wallbox/strings.json

@@ -7,15 +7,23 @@
           "username": "[%key:common::config_flow::data::username%]",
           "password": "[%key:common::config_flow::data::password%]"
         }
+      },
+      "reauth_confirm": {
+        "data": {
+          "username": "[%key:common::config_flow::data::username%]",
+          "password": "[%key:common::config_flow::data::password%]"
+        }
       }
     },
     "error": {
       "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]",
       "invalid_auth": "[%key:common::config_flow::error::invalid_auth%]",
-      "unknown": "[%key:common::config_flow::error::unknown%]"
+      "unknown": "[%key:common::config_flow::error::unknown%]",
+      "reauth_invalid": "Re-authentication failed; Serial Number does not match original"
     },
     "abort": {
-      "already_configured": "[%key:common::config_flow::abort::already_configured_device%]"
+      "already_configured": "[%key:common::config_flow::abort::already_configured_device%]",
+      "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]"
     }
   }
 }

homeassistant/components/wallbox/translations/en.json

@@ -1,14 +1,22 @@
 {
     "config": {
         "abort": {
-            "already_configured": "Device is already configured"
+            "already_configured": "Device is already configured",
+            "reauth_successful": "Re-authentication was successful"
         },
         "error": {
             "cannot_connect": "Failed to connect",
             "invalid_auth": "Invalid authentication",
+            "reauth_invalid": "Re-authentication failed; Serial Number does not match original",
             "unknown": "Unexpected error"
         },
         "step": {
+            "reauth_confirm": {
+                "data": {
+                    "password": "Password",
+                    "username": "Username"
+                }
+            },
             "user": {
                 "data": {
                     "password": "Password",
@@ -17,6 +25,5 @@
                 }
             }
         }
-    },
-    "title": "Wallbox"
+    }
 }

homeassistant/components/wallbox/translations/nl.json

@@ -1,7 +1,8 @@
 {
     "config": {
         "abort": {
-            "already_configured": "Apparaat is al geconfigureerd"
+            "already_configured": "Apparaat is al geconfigureerd",
+            "reauth_successful": "Herauthenticatie was succesvol"
         },
         "error": {
             "cannot_connect": "Kan geen verbinding maken",

tests/components/wallbox/test_config_flow.py

@@ -18,6 +18,7 @@
 )
 from homeassistant.core import HomeAssistant
 
+from tests.components.wallbox import entry, setup_integration
 from tests.components.wallbox.const import (
     CONF_ERROR,
     CONF_JWT,
@@ -162,3 +163,83 @@ async def test_form_validate_input(hass):
 
     assert result2["title"] == "Wallbox Portal"
     assert result2["data"]["station"] == "12345"
+
+
+async def test_form_reauth(hass):
+    """Test we handle reauth flow."""
+    await setup_integration(hass)
+    assert entry.state == config_entries.ConfigEntryState.LOADED
+
+    with requests_mock.Mocker() as mock_request:
+        mock_request.get(
+            "https://api.wall-box.com/auth/token/user",
+            text='{"jwt":"fakekeyhere","user_id":12345,"ttl":145656758,"error":false,"status":200}',
+            status_code=200,
+        )
+        mock_request.get(
+            "https://api.wall-box.com/chargers/status/12345",
+            json=test_response,
+            status_code=200,
+        )
+
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={
+                "source": config_entries.SOURCE_REAUTH,
+                "entry_id": entry.entry_id,
+            },
+        )
+
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                "station": "12345",
+                "username": "test-username",
+                "password": "test-password",
+            },
+        )
+
+    assert result2["type"] == "abort"
+    assert result2["reason"] == "reauth_successful"
+
+    await hass.config_entries.async_unload(entry.entry_id)
+
+
+async def test_form_reauth_invalid(hass):
+    """Test we handle reauth invalid flow."""
+    await setup_integration(hass)
+    assert entry.state == config_entries.ConfigEntryState.LOADED
+
+    with requests_mock.Mocker() as mock_request:
+        mock_request.get(
+            "https://api.wall-box.com/auth/token/user",
+            text='{"jwt":"fakekeyhere","user_id":12345,"ttl":145656758,"error":false,"status":200}',
+            status_code=200,
+        )
+        mock_request.get(
+            "https://api.wall-box.com/chargers/status/12345",
+            json=test_response,
+            status_code=200,
+        )
+
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN,
+            context={
+                "source": config_entries.SOURCE_REAUTH,
+                "entry_id": entry.entry_id,
+            },
+        )
+
+        result2 = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                "station": "12345678",
+                "username": "test-username",
+                "password": "test-password",
+            },
+        )
+
+    assert result2["type"] == "form"
+    assert result2["errors"] == {"base": "reauth_invalid"}
+
+    await hass.config_entries.async_unload(entry.entry_id)

tests/components/wallbox/test_init.py

@@ -3,10 +3,7 @@
 
 import requests_mock
 
-from homeassistant.components.wallbox import (
-    CONF_CONNECTIONS,
-    CONF_MAX_CHARGING_CURRENT_KEY,
-)
+from homeassistant.components.wallbox import CONF_MAX_CHARGING_CURRENT_KEY
 from homeassistant.config_entries import ConfigEntryState
 from homeassistant.core import HomeAssistant
 
@@ -78,7 +75,7 @@ async def test_wallbox_refresh_failed_invalid_auth(hass: HomeAssistant):
             status_code=403,
         )
 
-        wallbox = hass.data[DOMAIN][CONF_CONNECTIONS][entry.entry_id]
+        wallbox = hass.data[DOMAIN][entry.entry_id]
 
         await wallbox.async_refresh()
 
@@ -104,7 +101,7 @@ async def test_wallbox_refresh_failed_connection_error(hass: HomeAssistant):
             status_code=403,
         )
 
-        wallbox = hass.data[DOMAIN][CONF_CONNECTIONS][entry.entry_id]
+        wallbox = hass.data[DOMAIN][entry.entry_id]
 
         await wallbox.async_refresh()