Return multiple trains in Trafikverket Train (#106206)

* Add list of trains to coordinator

* Fix to work

* snapshot

* Fixes

* Fix

.coveragerc

@@ -1395,7 +1395,6 @@ omit =
     homeassistant/components/tradfri/switch.py
     homeassistant/components/trafikverket_train/__init__.py
     homeassistant/components/trafikverket_train/coordinator.py
-    homeassistant/components/trafikverket_train/sensor.py
     homeassistant/components/trafikverket_train/util.py
     homeassistant/components/trafikverket_weatherstation/__init__.py
     homeassistant/components/trafikverket_weatherstation/coordinator.py

homeassistant/components/trafikverket_train/coordinator.py

@@ -39,6 +39,8 @@ class TrainData:
     other_info: str | None
     deviation: str | None
     product_filter: str | None
+    departure_time_next: datetime | None
+    departure_time_next_next: datetime | None
 
 
 _LOGGER = logging.getLogger(__name__)
@@ -91,6 +93,7 @@ class TVDataUpdateCoordinator(DataUpdateCoordinator[TrainData]):
 
         when = dt_util.now()
         state: TrainStop | None = None
+        states: list[TrainStop] | None = None
         if self._time:
             departure_day = next_departuredate(self._weekdays)
             when = datetime.combine(
@@ -104,8 +107,12 @@ class TVDataUpdateCoordinator(DataUpdateCoordinator[TrainData]):
                     self.from_station, self.to_station, when, self._filter_product
                 )
             else:
-                state = await self._train_api.async_get_next_train_stop(
-                    self.from_station, self.to_station, when, self._filter_product
+                states = await self._train_api.async_get_next_train_stops(
+                    self.from_station,
+                    self.to_station,
+                    when,
+                    self._filter_product,
+                    number_of_stops=3,
                 )
         except InvalidAuthentication as error:
             raise ConfigEntryAuthFailed from error
@@ -117,6 +124,20 @@ class TVDataUpdateCoordinator(DataUpdateCoordinator[TrainData]):
                 f"Train departure {when} encountered a problem: {error}"
             ) from error
 
+        depart_next = None
+        depart_next_next = None
+        if not state and states:
+            state = states[0]
+            depart_next = (
+                states[1].advertised_time_at_location if len(states) > 1 else None
+            )
+            depart_next_next = (
+                states[2].advertised_time_at_location if len(states) > 2 else None
+            )
+
+        if not state:
+            raise UpdateFailed("Could not find any departures")
+
         departure_time = state.advertised_time_at_location
         if state.estimated_time_at_location:
             departure_time = state.estimated_time_at_location
@@ -125,7 +146,7 @@ class TVDataUpdateCoordinator(DataUpdateCoordinator[TrainData]):
 
         delay_time = state.get_delay_time()
 
-        states = TrainData(
+        return TrainData(
             departure_time=_get_as_utc(departure_time),
             departure_state=state.get_state().value,
             cancelled=state.canceled,
@@ -136,6 +157,6 @@ class TVDataUpdateCoordinator(DataUpdateCoordinator[TrainData]):
             other_info=_get_as_joined(state.other_information),
             deviation=_get_as_joined(state.deviations),
             product_filter=self._filter_product,
+            departure_time_next=_get_as_utc(depart_next),
+            departure_time_next_next=_get_as_utc(depart_next_next),
         )
-
-        return states

homeassistant/components/trafikverket_train/sensor.py

@@ -105,6 +105,20 @@ SENSOR_TYPES: tuple[TrafikverketSensorEntityDescription, ...] = (
         icon="mdi:alert",
         value_fn=lambda data: data.deviation,
     ),
+    TrafikverketSensorEntityDescription(
+        key="departure_time_next",
+        translation_key="departure_time_next",
+        icon="mdi:clock",
+        device_class=SensorDeviceClass.TIMESTAMP,
+        value_fn=lambda data: data.departure_time_next,
+    ),
+    TrafikverketSensorEntityDescription(
+        key="departure_time_next_next",
+        translation_key="departure_time_next_next",
+        icon="mdi:clock",
+        device_class=SensorDeviceClass.TIMESTAMP,
+        value_fn=lambda data: data.departure_time_next_next,
+    ),
 )
 
 

homeassistant/components/trafikverket_train/strings.json

@@ -69,6 +69,22 @@
           }
         }
       },
+      "departure_time_next": {
+        "name": "Departure time next",
+        "state_attributes": {
+          "product_filter": {
+            "name": "[%key:component::trafikverket_train::entity::sensor::departure_time::state_attributes::product_filter::name%]"
+          }
+        }
+      },
+      "departure_time_next_next": {
+        "name": "Departure time next after",
+        "state_attributes": {
+          "product_filter": {
+            "name": "[%key:component::trafikverket_train::entity::sensor::departure_time::state_attributes::product_filter::name%]"
+          }
+        }
+      },
       "departure_state": {
         "name": "Departure state",
         "state": {

tests/components/trafikverket_train/__init__.py

@@ -1 +1,28 @@
 """Tests for the Trafikverket Train integration."""
+from __future__ import annotations
+
+from homeassistant.components.trafikverket_ferry.const import (
+    CONF_FROM,
+    CONF_TIME,
+    CONF_TO,
+)
+from homeassistant.components.trafikverket_train.const import CONF_FILTER_PRODUCT
+from homeassistant.const import CONF_API_KEY, CONF_NAME, CONF_WEEKDAY, WEEKDAYS
+
+ENTRY_CONFIG = {
+    CONF_API_KEY: "1234567890",
+    CONF_FROM: "Stockholm C",
+    CONF_TO: "Uppsala C",
+    CONF_TIME: None,
+    CONF_WEEKDAY: WEEKDAYS,
+    CONF_NAME: "Stockholm C to Uppsala C",
+}
+ENTRY_CONFIG2 = {
+    CONF_API_KEY: "1234567890",
+    CONF_FROM: "Stockholm C",
+    CONF_TO: "Uppsala C",
+    CONF_TIME: "11:00:00",
+    CONF_WEEKDAY: WEEKDAYS,
+    CONF_NAME: "Stockholm C to Uppsala C",
+}
+OPTIONS_CONFIG = {CONF_FILTER_PRODUCT: "Regionaltåg"}

tests/components/trafikverket_train/conftest.py

@@ -0,0 +1,160 @@
+"""Fixtures for Trafikverket Train integration tests."""
+from __future__ import annotations
+
+from datetime import datetime, timedelta
+from unittest.mock import patch
+
+import pytest
+from pytrafikverket.trafikverket_train import TrainStop
+
+from homeassistant.components.trafikverket_train.const import DOMAIN
+from homeassistant.config_entries import SOURCE_USER
+from homeassistant.core import HomeAssistant
+from homeassistant.util import dt as dt_util
+
+from . import ENTRY_CONFIG, ENTRY_CONFIG2, OPTIONS_CONFIG
+
+from tests.common import MockConfigEntry
+
+
+@pytest.fixture(name="load_int")
+async def load_integration_from_entry(
+    hass: HomeAssistant,
+    get_trains: list[TrainStop],
+    get_train_stop: TrainStop,
+) -> MockConfigEntry:
+    """Set up the Trafikverket Train integration in Home Assistant."""
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        source=SOURCE_USER,
+        data=ENTRY_CONFIG,
+        options=OPTIONS_CONFIG,
+        entry_id="1",
+        unique_id="stockholmc-uppsalac--['mon', 'tue', 'wed', 'thu', 'fri', 'sat', 'sun']",
+    )
+    config_entry.add_to_hass(hass)
+    config_entry2 = MockConfigEntry(
+        domain=DOMAIN,
+        source=SOURCE_USER,
+        data=ENTRY_CONFIG2,
+        entry_id="2",
+        unique_id="stockholmc-uppsalac-1100-['mon', 'tue', 'wed', 'thu', 'fri', 'sat', 'sun']",
+    )
+    config_entry2.add_to_hass(hass)
+
+    with patch(
+        "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_get_next_train_stops",
+        return_value=get_trains,
+    ), patch(
+        "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_get_train_stop",
+        return_value=get_train_stop,
+    ), patch(
+        "homeassistant.components.trafikverket_train.TrafikverketTrain.async_get_train_station",
+    ):
+        await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+    return config_entry
+
+
+@pytest.fixture(name="get_trains")
+def fixture_get_trains() -> list[TrainStop]:
+    """Construct TrainStop Mock."""
+
+    depart1 = TrainStop(
+        id=13,
+        canceled=False,
+        advertised_time_at_location=datetime(2023, 5, 1, 12, 0, tzinfo=dt_util.UTC),
+        estimated_time_at_location=datetime(2023, 5, 1, 12, 0, tzinfo=dt_util.UTC),
+        time_at_location=datetime(2023, 5, 1, 12, 0, tzinfo=dt_util.UTC),
+        other_information=["Some other info"],
+        deviations=None,
+        modified_time=datetime(2023, 5, 1, 12, 0, tzinfo=dt_util.UTC),
+        product_description=["Regionaltåg"],
+    )
+    depart2 = TrainStop(
+        id=14,
+        canceled=False,
+        advertised_time_at_location=datetime(2023, 5, 1, 12, 0, tzinfo=dt_util.UTC)
+        + timedelta(minutes=15),
+        estimated_time_at_location=None,
+        time_at_location=None,
+        other_information=["Some other info"],
+        deviations=None,
+        modified_time=datetime(2023, 5, 1, 12, 0, tzinfo=dt_util.UTC),
+        product_description=["Regionaltåg"],
+    )
+    depart3 = TrainStop(
+        id=15,
+        canceled=False,
+        advertised_time_at_location=datetime(2023, 5, 1, 12, 0, tzinfo=dt_util.UTC)
+        + timedelta(minutes=30),
+        estimated_time_at_location=None,
+        time_at_location=None,
+        other_information=["Some other info"],
+        deviations=None,
+        modified_time=datetime(2023, 5, 1, 12, 0, tzinfo=dt_util.UTC),
+        product_description=["Regionaltåg"],
+    )
+
+    return [depart1, depart2, depart3]
+
+
+@pytest.fixture(name="get_trains_next")
+def fixture_get_trains_next() -> list[TrainStop]:
+    """Construct TrainStop Mock."""
+
+    depart1 = TrainStop(
+        id=13,
+        canceled=False,
+        advertised_time_at_location=datetime(2023, 5, 1, 17, 0, tzinfo=dt_util.UTC),
+        estimated_time_at_location=datetime(2023, 5, 1, 17, 0, tzinfo=dt_util.UTC),
+        time_at_location=datetime(2023, 5, 1, 17, 0, tzinfo=dt_util.UTC),
+        other_information=None,
+        deviations=None,
+        modified_time=datetime(2023, 5, 1, 12, 0, tzinfo=dt_util.UTC),
+        product_description=["Regionaltåg"],
+    )
+    depart2 = TrainStop(
+        id=14,
+        canceled=False,
+        advertised_time_at_location=datetime(2023, 5, 1, 17, 0, tzinfo=dt_util.UTC)
+        + timedelta(minutes=15),
+        estimated_time_at_location=None,
+        time_at_location=None,
+        other_information=["Some other info"],
+        deviations=None,
+        modified_time=datetime(2023, 5, 1, 12, 0, tzinfo=dt_util.UTC),
+        product_description=["Regionaltåg"],
+    )
+    depart3 = TrainStop(
+        id=15,
+        canceled=False,
+        advertised_time_at_location=datetime(2023, 5, 1, 17, 0, tzinfo=dt_util.UTC)
+        + timedelta(minutes=30),
+        estimated_time_at_location=None,
+        time_at_location=None,
+        other_information=["Some other info"],
+        deviations=None,
+        modified_time=datetime(2023, 5, 1, 12, 0, tzinfo=dt_util.UTC),
+        product_description=["Regionaltåg"],
+    )
+
+    return [depart1, depart2, depart3]
+
+
+@pytest.fixture(name="get_train_stop")
+def fixture_get_train_stop() -> TrainStop:
+    """Construct TrainStop Mock."""
+
+    return TrainStop(
+        id=13,
+        canceled=False,
+        advertised_time_at_location=datetime(2023, 5, 1, 11, 0, tzinfo=dt_util.UTC),
+        estimated_time_at_location=datetime(2023, 5, 1, 11, 0, tzinfo=dt_util.UTC),
+        time_at_location=datetime(2023, 5, 1, 11, 0, tzinfo=dt_util.UTC),
+        other_information=None,
+        deviations=None,
+        modified_time=datetime(2023, 5, 1, 11, 0, tzinfo=dt_util.UTC),
+        product_description=["Regionaltåg"],
+    )

tests/components/trafikverket_train/snapshots/test_sensor.ambr

@@ -0,0 +1,216 @@
+# serializer version: 1
+# name: test_sensor_next
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'Data provided by Trafikverket',
+      'device_class': 'timestamp',
+      'friendly_name': 'Stockholm C to Uppsala C Departure time',
+      'icon': 'mdi:clock',
+      'product_filter': 'Regionaltåg',
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.stockholm_c_to_uppsala_c_departure_time',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '2023-05-01T12:00:00+00:00',
+  })
+# ---
+# name: test_sensor_next.1
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'Data provided by Trafikverket',
+      'device_class': 'enum',
+      'friendly_name': 'Stockholm C to Uppsala C Departure state',
+      'icon': 'mdi:clock',
+      'options': list([
+        'on_time',
+        'delayed',
+        'canceled',
+      ]),
+      'product_filter': 'Regionaltåg',
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.stockholm_c_to_uppsala_c_departure_state',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'on_time',
+  })
+# ---
+# name: test_sensor_next.10
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'Data provided by Trafikverket',
+      'device_class': 'timestamp',
+      'friendly_name': 'Stockholm C to Uppsala C Departure time next',
+      'icon': 'mdi:clock',
+      'product_filter': 'Regionaltåg',
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.stockholm_c_to_uppsala_c_departure_time_next',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '2023-05-01T17:15:00+00:00',
+  })
+# ---
+# name: test_sensor_next.11
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'Data provided by Trafikverket',
+      'device_class': 'timestamp',
+      'friendly_name': 'Stockholm C to Uppsala C Departure time next after',
+      'icon': 'mdi:clock',
+      'product_filter': 'Regionaltåg',
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.stockholm_c_to_uppsala_c_departure_time_next_after',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '2023-05-01T17:30:00+00:00',
+  })
+# ---
+# name: test_sensor_next.2
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'Data provided by Trafikverket',
+      'device_class': 'timestamp',
+      'friendly_name': 'Stockholm C to Uppsala C Actual time',
+      'icon': 'mdi:clock',
+      'product_filter': 'Regionaltåg',
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.stockholm_c_to_uppsala_c_actual_time',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '2023-05-01T12:00:00+00:00',
+  })
+# ---
+# name: test_sensor_next.3
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'Data provided by Trafikverket',
+      'friendly_name': 'Stockholm C to Uppsala C Other information',
+      'icon': 'mdi:information-variant',
+      'product_filter': 'Regionaltåg',
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.stockholm_c_to_uppsala_c_other_information',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'Some other info',
+  })
+# ---
+# name: test_sensor_next.4
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'Data provided by Trafikverket',
+      'device_class': 'timestamp',
+      'friendly_name': 'Stockholm C to Uppsala C Departure time next',
+      'icon': 'mdi:clock',
+      'product_filter': 'Regionaltåg',
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.stockholm_c_to_uppsala_c_departure_time_next',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '2023-05-01T12:15:00+00:00',
+  })
+# ---
+# name: test_sensor_next.5
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'Data provided by Trafikverket',
+      'device_class': 'timestamp',
+      'friendly_name': 'Stockholm C to Uppsala C Departure time next after',
+      'icon': 'mdi:clock',
+      'product_filter': 'Regionaltåg',
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.stockholm_c_to_uppsala_c_departure_time_next_after',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '2023-05-01T12:30:00+00:00',
+  })
+# ---
+# name: test_sensor_next.6
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'Data provided by Trafikverket',
+      'device_class': 'timestamp',
+      'friendly_name': 'Stockholm C to Uppsala C Departure time',
+      'icon': 'mdi:clock',
+      'product_filter': 'Regionaltåg',
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.stockholm_c_to_uppsala_c_departure_time',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '2023-05-01T17:00:00+00:00',
+  })
+# ---
+# name: test_sensor_next.7
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'Data provided by Trafikverket',
+      'device_class': 'enum',
+      'friendly_name': 'Stockholm C to Uppsala C Departure state',
+      'icon': 'mdi:clock',
+      'options': list([
+        'on_time',
+        'delayed',
+        'canceled',
+      ]),
+      'product_filter': 'Regionaltåg',
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.stockholm_c_to_uppsala_c_departure_state',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'on_time',
+  })
+# ---
+# name: test_sensor_next.8
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'Data provided by Trafikverket',
+      'device_class': 'timestamp',
+      'friendly_name': 'Stockholm C to Uppsala C Actual time',
+      'icon': 'mdi:clock',
+      'product_filter': 'Regionaltåg',
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.stockholm_c_to_uppsala_c_actual_time',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '2023-05-01T17:00:00+00:00',
+  })
+# ---
+# name: test_sensor_next.9
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'Data provided by Trafikverket',
+      'friendly_name': 'Stockholm C to Uppsala C Other information',
+      'icon': 'mdi:information-variant',
+      'product_filter': 'Regionaltåg',
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.stockholm_c_to_uppsala_c_other_information',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': 'unknown',
+  })
+# ---
+# name: test_sensor_single_stop
+  StateSnapshot({
+    'attributes': ReadOnlyDict({
+      'attribution': 'Data provided by Trafikverket',
+      'device_class': 'timestamp',
+      'friendly_name': 'Stockholm C to Uppsala C Departure time',
+      'icon': 'mdi:clock',
+    }),
+    'context': <ANY>,
+    'entity_id': 'sensor.stockholm_c_to_uppsala_c_departure_time_2',
+    'last_changed': <ANY>,
+    'last_updated': <ANY>,
+    'state': '2023-05-01T11:00:00+00:00',
+  })
+# ---

tests/components/trafikverket_train/test_config_flow.py

@@ -196,7 +196,7 @@ async def test_flow_fails_departures(
     with patch(
         "homeassistant.components.trafikverket_train.config_flow.TrafikverketTrain.async_get_train_station",
     ), patch(
-        "homeassistant.components.trafikverket_train.config_flow.TrafikverketTrain.async_get_next_train_stop",
+        "homeassistant.components.trafikverket_train.config_flow.TrafikverketTrain.async_get_next_train_stops",
         side_effect=side_effect(),
     ), patch(
         "homeassistant.components.trafikverket_train.config_flow.TrafikverketTrain.async_get_train_stop",

tests/components/trafikverket_train/test_sensor.py

@@ -0,0 +1,74 @@
+"""The test for the Trafikverket train sensor platform."""
+from __future__ import annotations
+
+from datetime import timedelta
+from unittest.mock import patch
+
+from freezegun.api import FrozenDateTimeFactory
+from pytrafikverket.trafikverket_train import TrainStop
+from syrupy.assertion import SnapshotAssertion
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+
+from tests.common import async_fire_time_changed
+
+
+async def test_sensor_next(
+    hass: HomeAssistant,
+    freezer: FrozenDateTimeFactory,
+    entity_registry_enabled_by_default: None,
+    load_int: ConfigEntry,
+    get_trains_next: list[TrainStop],
+    get_train_stop: TrainStop,
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test the Trafikverket Train sensor."""
+    for entity in (
+        "sensor.stockholm_c_to_uppsala_c_departure_time",
+        "sensor.stockholm_c_to_uppsala_c_departure_state",
+        "sensor.stockholm_c_to_uppsala_c_actual_time",
+        "sensor.stockholm_c_to_uppsala_c_other_information",
+        "sensor.stockholm_c_to_uppsala_c_departure_time_next",
+        "sensor.stockholm_c_to_uppsala_c_departure_time_next_after",
+    ):
+        state = hass.states.get(entity)
+        assert state == snapshot
+
+    with patch(
+        "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_get_next_train_stops",
+        return_value=get_trains_next,
+    ), patch(
+        "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_get_train_stop",
+        return_value=get_train_stop,
+    ):
+        freezer.tick(timedelta(minutes=6))
+        async_fire_time_changed(hass)
+        await hass.async_block_till_done()
+
+    for entity in (
+        "sensor.stockholm_c_to_uppsala_c_departure_time",
+        "sensor.stockholm_c_to_uppsala_c_departure_state",
+        "sensor.stockholm_c_to_uppsala_c_actual_time",
+        "sensor.stockholm_c_to_uppsala_c_other_information",
+        "sensor.stockholm_c_to_uppsala_c_departure_time_next",
+        "sensor.stockholm_c_to_uppsala_c_departure_time_next_after",
+    ):
+        state = hass.states.get(entity)
+        assert state == snapshot
+
+
+async def test_sensor_single_stop(
+    hass: HomeAssistant,
+    freezer: FrozenDateTimeFactory,
+    entity_registry_enabled_by_default: None,
+    load_int: ConfigEntry,
+    get_trains_next: list[TrainStop],
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test the Trafikverket Train sensor."""
+    state = hass.states.get("sensor.stockholm_c_to_uppsala_c_departure_time_2")
+
+    assert state.state == "2023-05-01T11:00:00+00:00"
+
+    assert state == snapshot