Minor fixes Trafikverket Train (#72996)

* Minor fixes Trafikverket Train

* Remove ConfigEntryAuthFailed

homeassistant/components/trafikverket_train/__init__.py

@@ -1,15 +1,39 @@
 """The trafikverket_train component."""
 from __future__ import annotations
 
-from homeassistant.config_entries import ConfigEntry
+from pytrafikverket import TrafikverketTrain
-from homeassistant.core import HomeAssistant
 
-from .const import PLATFORMS
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_API_KEY
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
+
+from .const import CONF_FROM, CONF_TO, DOMAIN, PLATFORMS
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Trafikverket Train from a config entry."""
 
+    http_session = async_get_clientsession(hass)
+    train_api = TrafikverketTrain(http_session, entry.data[CONF_API_KEY])
+
+    try:
+        to_station = await train_api.async_get_train_station(entry.data[CONF_TO])
+        from_station = await train_api.async_get_train_station(entry.data[CONF_FROM])
+    except ValueError as error:
+        if "Invalid authentication" in error.args[0]:
+            raise ConfigEntryAuthFailed from error
+        raise ConfigEntryNotReady(
+            f"Problem when trying station {entry.data[CONF_FROM]} to {entry.data[CONF_TO]}. Error: {error} "
+        ) from error
+
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = {
+        CONF_TO: to_station,
+        CONF_FROM: from_station,
+        "train_api": train_api,
+    }
+
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
 
     return True

homeassistant/components/trafikverket_train/sensor.py

@@ -10,10 +10,8 @@ from pytrafikverket.trafikverket_train import StationInfo, TrainStop
 
 from homeassistant.components.sensor import SensorDeviceClass, SensorEntity
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_API_KEY, CONF_NAME, CONF_WEEKDAY, WEEKDAYS
+from homeassistant.const import CONF_NAME, CONF_WEEKDAY, WEEKDAYS
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
-from homeassistant.helpers.aiohttp_client import async_get_clientsession
 from homeassistant.helpers.device_registry import DeviceEntryType
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -42,24 +40,11 @@ async def async_setup_entry(
 ) -> None:
     """Set up the Trafikverket sensor entry."""
 
-    httpsession = async_get_clientsession(hass)
+    train_api = hass.data[DOMAIN][entry.entry_id]["train_api"]
-    train_api = TrafikverketTrain(httpsession, entry.data[CONF_API_KEY])
+    to_station = hass.data[DOMAIN][entry.entry_id][CONF_TO]
-
+    from_station = hass.data[DOMAIN][entry.entry_id][CONF_FROM]
-    try:
+    get_time: str | None = entry.data.get(CONF_TIME)
-        to_station = await train_api.async_get_train_station(entry.data[CONF_TO])
+    train_time = dt.parse_time(get_time) if get_time else None
-        from_station = await train_api.async_get_train_station(entry.data[CONF_FROM])
-    except ValueError as error:
-        if "Invalid authentication" in error.args[0]:
-            raise ConfigEntryAuthFailed from error
-        raise ConfigEntryNotReady(
-            f"Problem when trying station {entry.data[CONF_FROM]} to {entry.data[CONF_TO]}. Error: {error} "
-        ) from error
-
-    train_time = (
-        dt.parse_time(entry.data.get(CONF_TIME, ""))
-        if entry.data.get(CONF_TIME)
-        else None
-    )
 
     async_add_entities(
         [
@@ -157,8 +142,8 @@ class TrainSensor(SensorEntity):
                 _state = await self._train_api.async_get_next_train_stop(
                     self._from_station, self._to_station, when
                 )
-        except ValueError as output_error:
+        except ValueError as error:
-            _LOGGER.error("Departure %s encountered a problem: %s", when, output_error)
+            _LOGGER.error("Departure %s encountered a problem: %s", when, error)
 
         if not _state:
             self._attr_available = False

tests/components/trafikverket_train/test_config_flow.py

@@ -66,6 +66,51 @@ async def test_form(hass: HomeAssistant) -> None:
     )
 
 
+async def test_form_entry_already_exist(hass: HomeAssistant) -> None:
+    """Test flow aborts when entry already exist."""
+
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_API_KEY: "1234567890",
+            CONF_NAME: "Stockholm C to Uppsala C at 10:00",
+            CONF_FROM: "Stockholm C",
+            CONF_TO: "Uppsala C",
+            CONF_TIME: "10:00",
+            CONF_WEEKDAY: WEEKDAYS,
+        },
+        unique_id=f"stockholmc-uppsalac-10:00-{WEEKDAYS}",
+    )
+    entry.add_to_hass(hass)
+
+    result = await hass.config_entries.flow.async_init(
+        DOMAIN, context={"source": config_entries.SOURCE_USER}
+    )
+    assert result["type"] == RESULT_TYPE_FORM
+    assert result["errors"] == {}
+
+    with patch(
+        "homeassistant.components.trafikverket_train.config_flow.TrafikverketTrain.async_get_train_station",
+    ), patch(
+        "homeassistant.components.trafikverket_train.async_setup_entry",
+        return_value=True,
+    ):
+        result2 = 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 result2["type"] == RESULT_TYPE_ABORT
+    assert result2["reason"] == "already_configured"
+
+
 @pytest.mark.parametrize(
     "error_message,base_error",
     [