From a1b372e4ca59b9574908f8797607725c6c9328ac Mon Sep 17 00:00:00 2001 From: G Johansson Date: Sat, 4 Jun 2022 12:37:39 +0200 Subject: [PATCH] Minor fixes Trafikverket Train (#72996) * Minor fixes Trafikverket Train * Remove ConfigEntryAuthFailed --- .../components/trafikverket_train/__init__.py | 30 +++++++++++-- .../components/trafikverket_train/sensor.py | 31 ++++--------- .../trafikverket_train/test_config_flow.py | 45 +++++++++++++++++++ 3 files changed, 80 insertions(+), 26 deletions(-) diff --git a/homeassistant/components/trafikverket_train/__init__.py b/homeassistant/components/trafikverket_train/__init__.py index 4411ccab948d..ee026371b04e 100644 --- a/homeassistant/components/trafikverket_train/__init__.py +++ b/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 homeassistant.core import HomeAssistant +from pytrafikverket import TrafikverketTrain -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 diff --git a/homeassistant/components/trafikverket_train/sensor.py b/homeassistant/components/trafikverket_train/sensor.py index 4a419ff3b339..d9674e5373a8 100644 --- a/homeassistant/components/trafikverket_train/sensor.py +++ b/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 = TrafikverketTrain(httpsession, 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 - - train_time = ( - dt.parse_time(entry.data.get(CONF_TIME, "")) - if entry.data.get(CONF_TIME) - else None - ) + train_api = hass.data[DOMAIN][entry.entry_id]["train_api"] + to_station = hass.data[DOMAIN][entry.entry_id][CONF_TO] + from_station = hass.data[DOMAIN][entry.entry_id][CONF_FROM] + get_time: str | None = entry.data.get(CONF_TIME) + train_time = dt.parse_time(get_time) if get_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: - _LOGGER.error("Departure %s encountered a problem: %s", when, output_error) + except ValueError as error: + _LOGGER.error("Departure %s encountered a problem: %s", when, error) if not _state: self._attr_available = False diff --git a/tests/components/trafikverket_train/test_config_flow.py b/tests/components/trafikverket_train/test_config_flow.py index 09539584cbce..37788fc285b4 100644 --- a/tests/components/trafikverket_train/test_config_flow.py +++ b/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", [