Skip to content

Commit 06435f6

Browse files
Implement config flow for trafikverket_train (home-assistant#65182)
1 parent 425b825 commit 06435f6

File tree

13 files changed

+711
-41
lines changed

13 files changed

+711
-41
lines changed

.coveragerc

+1
Original file line numberDiff line numberDiff line change
@@ -1272,6 +1272,7 @@ omit =
12721272
homeassistant/components/tradfri/light.py
12731273
homeassistant/components/tradfri/sensor.py
12741274
homeassistant/components/tradfri/switch.py
1275+
homeassistant/components/trafikverket_train/__init__.py
12751276
homeassistant/components/trafikverket_train/sensor.py
12761277
homeassistant/components/trafikverket_weatherstation/__init__.py
12771278
homeassistant/components/trafikverket_weatherstation/coordinator.py

CODEOWNERS

+2-1
Original file line numberDiff line numberDiff line change
@@ -1048,6 +1048,7 @@ build.json @home-assistant/supervisor
10481048
/homeassistant/components/tractive/ @Danielhiversen @zhulik @bieniu
10491049
/tests/components/tractive/ @Danielhiversen @zhulik @bieniu
10501050
/homeassistant/components/trafikverket_train/ @endor-force @gjohansson-ST
1051+
/tests/components/trafikverket_train/ @endor-force @gjohansson-ST
10511052
/homeassistant/components/trafikverket_weatherstation/ @endor-force @gjohansson-ST
10521053
/tests/components/trafikverket_weatherstation/ @endor-force @gjohansson-ST
10531054
/homeassistant/components/transmission/ @engrbm87 @JPHutchins
@@ -1201,4 +1202,4 @@ build.json @home-assistant/supervisor
12011202
/homeassistant/components/demo/weather.py @fabaff
12021203

12031204
# Remove codeowners from files
1204-
/homeassistant/components/*/translations/
1205+
/homeassistant/components/*/translations/
Original file line numberDiff line numberDiff line change
@@ -1 +1,21 @@
11
"""The trafikverket_train component."""
2+
from __future__ import annotations
3+
4+
from homeassistant.config_entries import ConfigEntry
5+
from homeassistant.core import HomeAssistant
6+
7+
from .const import PLATFORMS
8+
9+
10+
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
11+
"""Set up Trafikverket Train from a config entry."""
12+
13+
hass.config_entries.async_setup_platforms(entry, PLATFORMS)
14+
15+
return True
16+
17+
18+
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
19+
"""Unload Trafikverket Weatherstation config entry."""
20+
21+
return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
1+
"""Adds config flow for Trafikverket Train integration."""
2+
from __future__ import annotations
3+
4+
from typing import Any
5+
6+
from pytrafikverket import TrafikverketTrain
7+
import voluptuous as vol
8+
9+
from homeassistant import config_entries
10+
from homeassistant.const import CONF_API_KEY, CONF_NAME, CONF_WEEKDAY, WEEKDAYS
11+
from homeassistant.data_entry_flow import FlowResult
12+
from homeassistant.helpers.aiohttp_client import async_get_clientsession
13+
import homeassistant.helpers.config_validation as cv
14+
import homeassistant.util.dt as dt_util
15+
16+
from .const import CONF_FROM, CONF_TIME, CONF_TO, DOMAIN
17+
from .util import create_unique_id
18+
19+
ERROR_INVALID_AUTH = "Source: Security, message: Invalid authentication"
20+
ERROR_INVALID_STATION = "Could not find a station with the specified name"
21+
ERROR_MULTIPLE_STATION = "Found multiple stations with the specified name"
22+
23+
DATA_SCHEMA = vol.Schema(
24+
{
25+
vol.Required(CONF_API_KEY): cv.string,
26+
vol.Required(CONF_FROM): cv.string,
27+
vol.Required(CONF_TO): cv.string,
28+
vol.Optional(CONF_TIME): cv.string,
29+
vol.Required(CONF_WEEKDAY, default=WEEKDAYS): cv.multi_select(
30+
{day: day for day in WEEKDAYS}
31+
),
32+
}
33+
)
34+
DATA_SCHEMA_REAUTH = vol.Schema(
35+
{
36+
vol.Required(CONF_API_KEY): cv.string,
37+
}
38+
)
39+
40+
41+
class TVTrainConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
42+
"""Handle a config flow for Trafikverket Train integration."""
43+
44+
VERSION = 1
45+
46+
entry: config_entries.ConfigEntry | None
47+
48+
async def validate_input(
49+
self, api_key: str, train_from: str, train_to: str
50+
) -> None:
51+
"""Validate input from user input."""
52+
web_session = async_get_clientsession(self.hass)
53+
train_api = TrafikverketTrain(web_session, api_key)
54+
await train_api.async_get_train_station(train_from)
55+
await train_api.async_get_train_station(train_to)
56+
57+
async def async_step_reauth(
58+
self, user_input: dict[str, Any] | None = None
59+
) -> FlowResult:
60+
"""Handle re-authentication with Trafikverket."""
61+
62+
self.entry = self.hass.config_entries.async_get_entry(self.context["entry_id"])
63+
return await self.async_step_reauth_confirm()
64+
65+
async def async_step_reauth_confirm(
66+
self, user_input: dict[str, Any] | None = None
67+
) -> FlowResult:
68+
"""Confirm re-authentication with Trafikverket."""
69+
errors: dict[str, str] = {}
70+
71+
if user_input:
72+
api_key = user_input[CONF_API_KEY]
73+
74+
assert self.entry is not None
75+
try:
76+
await self.validate_input(
77+
api_key, self.entry.data[CONF_FROM], self.entry.data[CONF_TO]
78+
)
79+
except ValueError as err:
80+
if str(err) == ERROR_INVALID_AUTH:
81+
errors["base"] = "invalid_auth"
82+
elif str(err) == ERROR_INVALID_STATION:
83+
errors["base"] = "invalid_station"
84+
elif str(err) == ERROR_MULTIPLE_STATION:
85+
errors["base"] = "more_stations"
86+
else:
87+
errors["base"] = "cannot_connect"
88+
else:
89+
self.hass.config_entries.async_update_entry(
90+
self.entry,
91+
data={
92+
**self.entry.data,
93+
CONF_API_KEY: api_key,
94+
},
95+
)
96+
await self.hass.config_entries.async_reload(self.entry.entry_id)
97+
return self.async_abort(reason="reauth_successful")
98+
99+
return self.async_show_form(
100+
step_id="reauth_confirm",
101+
data_schema=DATA_SCHEMA_REAUTH,
102+
errors=errors,
103+
)
104+
105+
async def async_step_import(self, config: dict[str, Any] | None) -> FlowResult:
106+
"""Import a configuration from config.yaml."""
107+
108+
return await self.async_step_user(user_input=config)
109+
110+
async def async_step_user(
111+
self, user_input: dict[str, Any] | None = None
112+
) -> FlowResult:
113+
"""Handle the user step."""
114+
errors: dict[str, str] = {}
115+
116+
if user_input is not None:
117+
api_key: str = user_input[CONF_API_KEY]
118+
train_from: str = user_input[CONF_FROM]
119+
train_to: str = user_input[CONF_TO]
120+
train_time: str | None = user_input.get(CONF_TIME)
121+
train_days: list = user_input[CONF_WEEKDAY]
122+
123+
name = f"{train_from} to {train_to}"
124+
if train_time:
125+
name = f"{train_from} to {train_to} at {train_time}"
126+
127+
try:
128+
await self.validate_input(api_key, train_from, train_to)
129+
except ValueError as err:
130+
if str(err) == ERROR_INVALID_AUTH:
131+
errors["base"] = "invalid_auth"
132+
elif str(err) == ERROR_INVALID_STATION:
133+
errors["base"] = "invalid_station"
134+
elif str(err) == ERROR_MULTIPLE_STATION:
135+
errors["base"] = "more_stations"
136+
else:
137+
errors["base"] = "cannot_connect"
138+
else:
139+
if train_time:
140+
if bool(dt_util.parse_time(train_time) is None):
141+
errors["base"] = "invalid_time"
142+
if not errors:
143+
unique_id = create_unique_id(
144+
train_from, train_to, train_time, train_days
145+
)
146+
await self.async_set_unique_id(unique_id)
147+
self._abort_if_unique_id_configured()
148+
return self.async_create_entry(
149+
title=name,
150+
data={
151+
CONF_API_KEY: api_key,
152+
CONF_NAME: name,
153+
CONF_FROM: train_from,
154+
CONF_TO: train_to,
155+
CONF_TIME: train_time,
156+
CONF_WEEKDAY: train_days,
157+
},
158+
)
159+
160+
return self.async_show_form(
161+
step_id="user",
162+
data_schema=DATA_SCHEMA,
163+
errors=errors,
164+
)
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
"""Adds constants for Trafikverket Train integration."""
2+
from homeassistant.const import Platform
3+
4+
DOMAIN = "trafikverket_train"
5+
PLATFORMS = [Platform.SENSOR]
6+
ATTRIBUTION = "Data provided by Trafikverket"
7+
8+
CONF_TRAINS = "trains"
9+
CONF_FROM = "from"
10+
CONF_TO = "to"
11+
CONF_TIME = "time"

homeassistant/components/trafikverket_train/manifest.json

+1
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
"documentation": "https://www.home-assistant.io/integrations/trafikverket_train",
55
"requirements": ["pytrafikverket==0.1.6.2"],
66
"codeowners": ["@endor-force", "@gjohansson-ST"],
7+
"config_flow": true,
78
"iot_class": "cloud_polling",
89
"loggers": ["pytrafikverket"]
910
}

homeassistant/components/trafikverket_train/sensor.py

+77-40
Original file line numberDiff line numberDiff line change
@@ -14,20 +14,22 @@
1414
SensorDeviceClass,
1515
SensorEntity,
1616
)
17+
from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry
1718
from homeassistant.const import CONF_API_KEY, CONF_NAME, CONF_WEEKDAY, WEEKDAYS
1819
from homeassistant.core import HomeAssistant
20+
from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
1921
from homeassistant.helpers.aiohttp_client import async_get_clientsession
2022
import homeassistant.helpers.config_validation as cv
23+
from homeassistant.helpers.device_registry import DeviceEntryType
24+
from homeassistant.helpers.entity import DeviceInfo
2125
from homeassistant.helpers.entity_platform import AddEntitiesCallback
2226
from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
23-
from homeassistant.util.dt import as_utc, get_time_zone
27+
from homeassistant.util.dt import as_utc, get_time_zone, parse_time
2428

25-
_LOGGER = logging.getLogger(__name__)
29+
from .const import CONF_FROM, CONF_TIME, CONF_TO, CONF_TRAINS, DOMAIN
30+
from .util import create_unique_id
2631

27-
CONF_TRAINS = "trains"
28-
CONF_FROM = "from"
29-
CONF_TO = "to"
30-
CONF_TIME = "time"
32+
_LOGGER = logging.getLogger(__name__)
3133

3234
ATTR_DEPARTURE_STATE = "departure_state"
3335
ATTR_CANCELED = "canceled"
@@ -66,43 +68,66 @@ async def async_setup_platform(
6668
async_add_entities: AddEntitiesCallback,
6769
discovery_info: DiscoveryInfoType | None = None,
6870
) -> None:
69-
"""Set up the departure sensor."""
70-
httpsession = async_get_clientsession(hass)
71-
train_api = TrafikverketTrain(httpsession, config[CONF_API_KEY])
72-
sensors = []
73-
station_cache = {}
71+
"""Import Trafikverket Train configuration from YAML."""
72+
_LOGGER.warning(
73+
# Config flow added in Home Assistant Core 2022.3, remove import flow in 2022.7
74+
"Loading Trafikverket Train via platform setup is deprecated; Please remove it from your configuration"
75+
)
76+
7477
for train in config[CONF_TRAINS]:
75-
try:
76-
trainstops = [train[CONF_FROM], train[CONF_TO]]
77-
for station in trainstops:
78-
if station not in station_cache:
79-
station_cache[station] = await train_api.async_get_train_station(
80-
station
81-
)
82-
83-
except ValueError as station_error:
84-
if "Invalid authentication" in station_error.args[0]:
85-
_LOGGER.error("Unable to set up up component: %s", station_error)
86-
return
87-
_LOGGER.error(
88-
"Problem when trying station %s to %s. Error: %s ",
89-
train[CONF_FROM],
90-
train[CONF_TO],
91-
station_error,
78+
79+
new_config = {
80+
CONF_API_KEY: config[CONF_API_KEY],
81+
CONF_FROM: train[CONF_FROM],
82+
CONF_TO: train[CONF_TO],
83+
CONF_TIME: str(train.get(CONF_TIME)),
84+
CONF_WEEKDAY: train.get(CONF_WEEKDAY, WEEKDAYS),
85+
}
86+
hass.async_create_task(
87+
hass.config_entries.flow.async_init(
88+
DOMAIN,
89+
context={"source": SOURCE_IMPORT},
90+
data=new_config,
9291
)
93-
continue
94-
95-
sensor = TrainSensor(
96-
train_api,
97-
train[CONF_NAME],
98-
station_cache[train[CONF_FROM]],
99-
station_cache[train[CONF_TO]],
100-
train[CONF_WEEKDAY],
101-
train.get(CONF_TIME),
10292
)
103-
sensors.append(sensor)
10493

105-
async_add_entities(sensors, update_before_add=True)
94+
95+
async def async_setup_entry(
96+
hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
97+
) -> None:
98+
"""Set up the Trafikverket sensor entry."""
99+
100+
httpsession = async_get_clientsession(hass)
101+
train_api = TrafikverketTrain(httpsession, entry.data[CONF_API_KEY])
102+
103+
try:
104+
to_station = await train_api.async_get_train_station(entry.data[CONF_TO])
105+
from_station = await train_api.async_get_train_station(entry.data[CONF_FROM])
106+
except ValueError as error:
107+
if "Invalid authentication" in error.args[0]:
108+
raise ConfigEntryAuthFailed from error
109+
raise ConfigEntryNotReady(
110+
f"Problem when trying station {entry.data[CONF_FROM]} to {entry.data[CONF_TO]}. Error: {error} "
111+
) from error
112+
113+
train_time = (
114+
parse_time(entry.data.get(CONF_TIME, "")) if entry.data.get(CONF_TIME) else None
115+
)
116+
117+
async_add_entities(
118+
[
119+
TrainSensor(
120+
train_api,
121+
entry.data[CONF_NAME],
122+
from_station,
123+
to_station,
124+
entry.data[CONF_WEEKDAY],
125+
train_time,
126+
entry.entry_id,
127+
)
128+
],
129+
True,
130+
)
106131

107132

108133
def next_weekday(fromdate: date, weekday: int) -> date:
@@ -144,7 +169,8 @@ def __init__(
144169
from_station: str,
145170
to_station: str,
146171
weekday: list,
147-
departuretime: time,
172+
departuretime: time | None,
173+
entry_id: str,
148174
) -> None:
149175
"""Initialize the sensor."""
150176
self._train_api = train_api
@@ -153,6 +179,17 @@ def __init__(
153179
self._to_station = to_station
154180
self._weekday = weekday
155181
self._time = departuretime
182+
self._attr_device_info = DeviceInfo(
183+
entry_type=DeviceEntryType.SERVICE,
184+
identifiers={(DOMAIN, entry_id)},
185+
manufacturer="Trafikverket",
186+
model="v1.2",
187+
name=name,
188+
configuration_url="https://api.trafikinfo.trafikverket.se/",
189+
)
190+
self._attr_unique_id = create_unique_id(
191+
from_station, to_station, departuretime, weekday
192+
)
156193

157194
async def async_update(self) -> None:
158195
"""Retrieve latest state."""

0 commit comments

Comments
 (0)