Add reconfigure flow to Trafikverket Train (#136000)

homeassistant/components/trafikverket_train/config_flow.py

@@ -16,6 +16,7 @@ from pytrafikverket import (
 import voluptuous as vol
 
 from homeassistant.config_entries import (
+    SOURCE_RECONFIGURE,
     ConfigEntry,
     ConfigFlow,
     ConfigFlowResult,
@@ -146,6 +147,18 @@ class TVTrainConfigFlow(ConfigFlow, domain=DOMAIN):
 
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
+    ) -> ConfigFlowResult:
+        """Handle the user step."""
+        return await self.async_step_initial(user_input)
+
+    async def async_step_reconfigure(
+        self, user_input: dict[str, Any] | None = None
+    ) -> ConfigFlowResult:
+        """Handle the user step."""
+        return await self.async_step_initial(user_input)
+
+    async def async_step_initial(
+        self, user_input: dict[str, Any] | None = None
     ) -> ConfigFlowResult:
         """Handle the user step."""
         errors: dict[str, str] = {}
@@ -185,6 +198,22 @@ class TVTrainConfigFlow(ConfigFlow, domain=DOMAIN):
                             CONF_FILTER_PRODUCT: filter_product,
                         }
                     )
+
+                    if self.source == SOURCE_RECONFIGURE:
+                        reconfigure_entry = self._get_reconfigure_entry()
+                        return self.async_update_reload_and_abort(
+                            reconfigure_entry,
+                            title=name,
+                            data={
+                                CONF_API_KEY: api_key,
+                                CONF_NAME: name,
+                                CONF_FROM: self._from_stations[0].signature,
+                                CONF_TO: self._to_stations[0].signature,
+                                CONF_TIME: train_time,
+                                CONF_WEEKDAY: train_days,
+                            },
+                            options={CONF_FILTER_PRODUCT: filter_product},
+                        )
                     return self.async_create_entry(
                         title=name,
                         data={
@@ -201,7 +230,7 @@ class TVTrainConfigFlow(ConfigFlow, domain=DOMAIN):
                 return await self.async_step_select_stations()
 
         return self.async_show_form(
-            step_id="user",
+            step_id="initial",
             data_schema=self.add_suggested_values_to_schema(
                 DATA_SCHEMA, user_input or {}
             ),
@@ -238,6 +267,21 @@ class TVTrainConfigFlow(ConfigFlow, domain=DOMAIN):
                     CONF_FILTER_PRODUCT: filter_product,
                 }
             )
+            if self.source == SOURCE_RECONFIGURE:
+                reconfigure_entry = self._get_reconfigure_entry()
+                return self.async_update_reload_and_abort(
+                    reconfigure_entry,
+                    title=name,
+                    data={
+                        CONF_API_KEY: api_key,
+                        CONF_NAME: name,
+                        CONF_FROM: train_from,
+                        CONF_TO: train_to,
+                        CONF_TIME: train_time,
+                        CONF_WEEKDAY: train_days,
+                    },
+                    options={CONF_FILTER_PRODUCT: filter_product},
+                )
             return self.async_create_entry(
                 title=name,
                 data={

homeassistant/components/trafikverket_train/strings.json

@@ -2,7 +2,8 @@
   "config": {
     "abort": {
       "already_configured": "[%key:common::config_flow::abort::already_configured_account%]",
-      "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]"
+      "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]",
+      "reconfigure_successful": "[%key:common::config_flow::abort::reconfigure_successful%]"
     },
     "error": {
       "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]",
@@ -13,7 +14,7 @@
       "incorrect_api_key": "Invalid API key for selected account"
     },
     "step": {
-      "user": {
+      "initial": {
         "data": {
           "api_key": "[%key:common::config_flow::data::api_key%]",
           "to": "To station",
@@ -45,10 +46,10 @@
     "step": {
       "init": {
         "data": {
-          "filter_product": "[%key:component::trafikverket_train::config::step::user::data::filter_product%]"
+          "filter_product": "[%key:component::trafikverket_train::config::step::initial::data::filter_product%]"
         },
         "data_description": {
-          "filter_product": "[%key:component::trafikverket_train::config::step::user::data_description::filter_product%]"
+          "filter_product": "[%key:component::trafikverket_train::config::step::initial::data_description::filter_product%]"
         }
       }
     }

tests/components/trafikverket_train/test_config_flow.py

@@ -25,6 +25,8 @@ from homeassistant.const import CONF_API_KEY, CONF_NAME, CONF_WEEKDAY, WEEKDAYS
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 
+from . import ENTRY_CONFIG, OPTIONS_CONFIG
+
 from tests.common import MockConfigEntry
 
 
@@ -37,6 +39,7 @@ async def test_form(
         DOMAIN, context={"source": config_entries.SOURCE_USER}
     )
     assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "initial"
     assert result["errors"] == {}
 
     with (
@@ -218,7 +221,7 @@ async def test_flow_fails(
     )
 
     assert result["type"] is FlowResultType.FORM
-    assert result["step_id"] == config_entries.SOURCE_USER
+    assert result["step_id"] == "initial"
 
     with (
         patch(
@@ -260,7 +263,7 @@ async def test_flow_fails_departures(
     )
 
     assert result["type"] is FlowResultType.FORM
-    assert result["step_id"] == config_entries.SOURCE_USER
+    assert result["step_id"] == "initial"
 
     with (
         patch(
@@ -574,3 +577,328 @@ async def test_options_flow(
 
         assert result["type"] is FlowResultType.CREATE_ENTRY
         assert result["data"] == {"filter_product": None}
+
+
+async def test_reconfigure_flow(
+    hass: HomeAssistant, get_train_stations: list[StationInfoModel]
+) -> None:
+    """Test reconfigure flow."""
+
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data=ENTRY_CONFIG,
+        options=OPTIONS_CONFIG,
+        entry_id="1",
+        version=2,
+        minor_version=1,
+    )
+    config_entry.add_to_hass(hass)
+    result = await config_entry.start_reconfigure_flow(hass)
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "initial"
+    assert result["errors"] == {}
+
+    with (
+        patch(
+            "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_search_train_stations",
+            side_effect=get_train_stations,
+        ),
+        patch(
+            "homeassistant.components.trafikverket_train.async_setup_entry",
+            return_value=True,
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                CONF_API_KEY: "1234567890",
+                CONF_FROM: "Stockholm C",
+                CONF_TO: "Uppsala C",
+                CONF_TIME: "10:00",
+                CONF_WEEKDAY: ["mon", "fri"],
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "reconfigure_successful"
+
+
+async def test_reconfigure_multiple_stations(
+    hass: HomeAssistant, get_multiple_train_stations: list[StationInfoModel]
+) -> None:
+    """Test we can reconfigure with multiple stations."""
+
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data=ENTRY_CONFIG,
+        options=OPTIONS_CONFIG,
+        entry_id="1",
+        version=2,
+        minor_version=1,
+    )
+    config_entry.add_to_hass(hass)
+    result = await config_entry.start_reconfigure_flow(hass)
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "initial"
+    assert result["errors"] == {}
+
+    with (
+        patch(
+            "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_search_train_stations",
+            side_effect=get_multiple_train_stations,
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                CONF_API_KEY: "1234567890",
+                CONF_FROM: "Stockholm C",
+                CONF_TO: "Uppsala C",
+                CONF_TIME: "10:00",
+                CONF_WEEKDAY: ["mon", "fri"],
+            },
+        )
+        await hass.async_block_till_done()
+
+    with (
+        patch(
+            "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_search_train_stations",
+            side_effect=get_multiple_train_stations,
+        ),
+        patch(
+            "homeassistant.components.trafikverket_train.async_setup_entry",
+            return_value=True,
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                CONF_FROM: "Csu",
+                CONF_TO: "Ups",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "reconfigure_successful"
+
+
+async def test_reconfigure_entry_already_exist(
+    hass: HomeAssistant, get_train_stations: list[StationInfoModel]
+) -> None:
+    """Test flow aborts when entry already exist in a reconfigure flow."""
+
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_API_KEY: "1234567890",
+            CONF_NAME: "Stockholm C to Uppsala C at 10:00",
+            CONF_FROM: "Cst",
+            CONF_TO: "U",
+            CONF_TIME: "10:00",
+            CONF_WEEKDAY: WEEKDAYS,
+            CONF_FILTER_PRODUCT: None,
+        },
+        version=2,
+        minor_version=1,
+    )
+    config_entry.add_to_hass(hass)
+
+    config_entry2 = MockConfigEntry(
+        domain=DOMAIN,
+        data=ENTRY_CONFIG,
+        options=OPTIONS_CONFIG,
+        entry_id="1",
+        version=2,
+        minor_version=1,
+    )
+    config_entry2.add_to_hass(hass)
+    result = await config_entry2.start_reconfigure_flow(hass)
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["errors"] == {}
+
+    with (
+        patch(
+            "homeassistant.components.trafikverket_train.config_flow.TrafikverketTrain.async_get_train_stop",
+        ),
+        patch(
+            "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_search_train_stations",
+            side_effect=get_train_stations,
+        ),
+        patch(
+            "homeassistant.components.trafikverket_train.async_setup_entry",
+            return_value=True,
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                CONF_API_KEY: "1234567890",
+                CONF_FROM: "Stockholm C",
+                CONF_TO: "Uppsala C",
+                CONF_TIME: "10:00",
+                CONF_WEEKDAY: WEEKDAYS,
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "already_configured"
+
+
+@pytest.mark.parametrize(
+    ("side_effect", "p_error"),
+    [
+        (
+            InvalidAuthentication,
+            {"base": "invalid_auth"},
+        ),
+        (
+            NoTrainStationFound,
+            {"from": "invalid_station", "to": "invalid_station"},
+        ),
+        (
+            Exception,
+            {"base": "cannot_connect"},
+        ),
+    ],
+)
+async def test_reconfigure_flow_fails(
+    hass: HomeAssistant,
+    side_effect: Exception,
+    p_error: dict[str, str],
+    get_train_stations: list[StationInfoModel],
+) -> None:
+    """Test config flow errors."""
+    config_entry2 = MockConfigEntry(
+        domain=DOMAIN,
+        data=ENTRY_CONFIG,
+        options=OPTIONS_CONFIG,
+        entry_id="1",
+        version=2,
+        minor_version=1,
+    )
+    config_entry2.add_to_hass(hass)
+    result = await config_entry2.start_reconfigure_flow(hass)
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "initial"
+
+    with (
+        patch(
+            "homeassistant.components.trafikverket_train.config_flow.TrafikverketTrain.async_search_train_stations",
+            side_effect=side_effect(),
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            user_input={
+                CONF_API_KEY: "1234567890",
+                CONF_FROM: "Stockholm C",
+                CONF_TO: "Uppsala C",
+            },
+        )
+
+    assert result["errors"] == p_error
+
+    with (
+        patch(
+            "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_search_train_stations",
+            side_effect=get_train_stations,
+        ),
+        patch(
+            "homeassistant.components.trafikverket_train.async_setup_entry",
+            return_value=True,
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                CONF_API_KEY: "1234567890",
+                CONF_FROM: "Stockholm C",
+                CONF_TO: "Uppsala C",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "reconfigure_successful"
+
+
+@pytest.mark.parametrize(
+    ("side_effect", "p_error"),
+    [
+        (
+            NoTrainStationFound,
+            {"from": "invalid_station", "to": "invalid_station"},
+        ),
+        (
+            UnknownError,
+            {"base": "cannot_connect"},
+        ),
+    ],
+)
+async def test_reconfigure_flow_fails_departures(
+    hass: HomeAssistant,
+    side_effect: Exception,
+    p_error: dict[str, str],
+    get_train_stations: list[StationInfoModel],
+) -> None:
+    """Test config flow errors."""
+    config_entry2 = MockConfigEntry(
+        domain=DOMAIN,
+        data=ENTRY_CONFIG,
+        options=OPTIONS_CONFIG,
+        entry_id="1",
+        version=2,
+        minor_version=1,
+    )
+    config_entry2.add_to_hass(hass)
+    result = await config_entry2.start_reconfigure_flow(hass)
+
+    assert result["type"] is FlowResultType.FORM
+    assert result["step_id"] == "initial"
+
+    with (
+        patch(
+            "homeassistant.components.trafikverket_train.config_flow.TrafikverketTrain.async_search_train_stations",
+            side_effect=side_effect(),
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            user_input={
+                CONF_API_KEY: "1234567890",
+                CONF_FROM: "Stockholm C",
+                CONF_TO: "Uppsala C",
+            },
+        )
+
+    assert result["errors"] == p_error
+
+    with (
+        patch(
+            "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_search_train_stations",
+            side_effect=get_train_stations,
+        ),
+        patch(
+            "homeassistant.components.trafikverket_train.async_setup_entry",
+            return_value=True,
+        ),
+    ):
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            {
+                CONF_API_KEY: "1234567890",
+                CONF_FROM: "Stockholm C",
+                CONF_TO: "Uppsala C",
+            },
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] is FlowResultType.ABORT
+    assert result["reason"] == "reconfigure_successful"