Repository: home-assistant/core
Pull request: Return multiple trains in Trafikverket Train
<!--
  You are amazing! Thanks for contributing to our project!
  Please, DO NOT DELETE ANY TEXT from this template! (unless instructed).
-->
## Proposed change
<!--
  Describe the big picture of your changes here to communicate to the
  maintainers why we should accept this pull request. If it fixes a bug
  or resolves a feature request, be sure to link to that issue in the
  additional information section.
-->
Extends the integration to return more departures if not specific timeslot has been set

## Type of change
<!--
  What type of change does your PR introduce to Home Assistant?
  NOTE: Please, check only 1! box!
  If your PR requires multiple boxes to be checked, you'll most likely need to
  split it into multiple PRs. This makes things easier and faster to code review.
-->

- [ ] Dependency upgrade
- [ ] Bugfix (non-breaking change which fixes an issue)
- [ ] New integration (thank you!)
- [x] New feature (which adds functionality to an existing integration)
- [ ] Deprecation (breaking change to happen in the future)
- [ ] Breaking change (fix/feature causing existing functionality to break)
- [ ] Code quality improvements to existing code or addition of tests

## Additional information
<!--
  Details are important, and help maintainers processing your PR.
  Please be sure to fill out additional details, if applicable.
-->

- This PR fixes or closes issue: fixes #93154
- This PR is related to issue: 
- Link to documentation pull request: https://github.com/home-assistant/home-assistant.io/pull/30459

## Checklist
<!--
  Put an `x` in the boxes that apply. You can also fill these out after
  creating the PR. If you're unsure about any of them, don't hesitate to ask.
  We're here to help! This is simply a reminder of what we are going to look
  for before merging your code.
-->

- [x] The code change is tested and works locally.
- [x] Local tests pass. **Your PR cannot be merged unless tests pass**
- [x] There is no commented out code in this PR.
- [x] I have followed the [development checklist][dev-checklist]
- [x] I have followed the [perfect PR recommendations][perfect-pr]
- [x] The code has been formatted using Ruff (`ruff format homeassistant tests`)
- [x] Tests have been added to verify that the new code works.

If user exposed functionality or configuration variables are added/changed:

- [x] Documentation added/updated for [www.home-assistant.io][docs-repository]

If the code communicates with devices, web services, or third-party tools:

- [ ] The [manifest file][manifest-docs] has all fields filled out correctly.  
      Updated and included derived files by running: `python3 -m script.hassfest`.
- [ ] New or updated dependencies have been added to `requirements_all.txt`.  
      Updated by running `python3 -m script.gen_requirements_all`.
- [ ] For the updated dependencies - a link to the changelog, or at minimum a diff between library versions is added to the PR description.
- [ ] Untested files have been added to `.coveragerc`.

<!--
  This project is very active and we have a high turnover of pull requests.

  Unfortunately, the number of incoming pull requests is higher than what our
  reviewers can review and merge so there is a long backlog of pull requests
  waiting for review. You can help here!
  
  By reviewing another pull request, you will help raise the code quality of
  that pull request and the final review will be faster. This way the general
  pace of pull request reviews will go up and your wait time will go down.
  
  When picking a pull request to review, try to choose one that hasn't yet
  been reviewed.

  Thanks for helping out!
-->

To help with the load of incoming pull requests:

- [x] I have reviewed two other [open pull requests][prs] in this repository.

[prs]: https://github.com/home-assistant/core/pulls?q=is%3Aopen+is%3Apr+-author%3A%40me+-draft%3Atrue+-label%3Awaiting-for-upstream+sort%3Acreated-desc+review%3Anone+-status%3Afailure

<!--
  Thank you for contributing <3

  Below, some useful links you could explore:
-->
[dev-checklist]: https://developers.home-assistant.io/docs/development_checklist/
[manifest-docs]: https://developers.home-assistant.io/docs/creating_integration_manifest/
[quality-scale]: https://developers.home-assistant.io/docs/integration_quality_scale_index/
[docs-repository]: https://github.com/home-assistant/home-assistant.io
[perfect-pr]: https://developers.home-assistant.io/docs/review-process/#creating-the-perfect-pr


Problem statement:
Trafikverket train - get multiple departures for same route
### The problem

Is it possible to retrieve multiple departures for the same route instead of just the next one? The trains over the Oresund bridge can be very "unpredictable" 

### What version of Home Assistant Core has the issue?

2023.5.3

### What was the last working version of Home Assistant Core?

_No response_

### What type of installation are you running?

Home Assistant OS

### Integration causing the issue

_No response_

### Link to integration documentation on our website

_No response_

### Diagnostics information

_No response_

### Example YAML snippet

_No response_

### Anything in the logs that might be useful for us?

_No response_

### Additional information

_No response_

Hints from issue discussion:
Hey there @endor-force, @gjohansson-st, mind taking a look at this issue as it has been labeled with an integration (`trafikverket_train`) you are listed as a [code owner](https://github.com/home-assistant/core/blob/dev/CODEOWNERS#L1295) for? Thanks!

<details>
<summary>Code owner commands</summary>


Code owners of `trafikverket_train` can trigger bot actions by commenting:

- `@home-assistant close` Closes the issue.
- `@home-assistant rename Awesome new title` Renames the issue.
- `@home-assistant reopen` Reopen the issue.
- `@home-assistant unassign trafikverket_train` Removes the current integration label and assignees on the issue, add the integration domain after the command.

</details>


<sub><sup>(message by CodeOwnersMention)</sup></sub>

---

[trafikverket_train documentation](https://www.home-assistant.io/integrations/trafikverket_train)
[trafikverket_train source](https://github.com/home-assistant/core/tree/dev/homeassistant/components/trafikverket_train)
<sub><sup>(message by IssueLinks)</sup></sub>
This has been obviously resting for a long time but I'm prepping now to add it in so multiple sensors can be shown with different departure times.

Commit message: Fix to work

Patch to review:
diff --git a/.coveragerc b/.coveragerc
index adb0458d6f67bc..6cf95669feafdc 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -1386,7 +1386,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
diff --git a/homeassistant/components/trafikverket_train/coordinator.py b/homeassistant/components/trafikverket_train/coordinator.py
index 91a7e9f07b2c57..084970b93e3e10 100644
--- a/homeassistant/components/trafikverket_train/coordinator.py
+++ b/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 @@ async def _async_update_data(self) -> 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 @@ async def _async_update_data(self) -> 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,18 @@ async def _async_update_data(self) -> TrainData:
                 f"Train departure {when} encountered a problem: {error}"
             ) from error
 
+        depart_next = None
+        depart_next_next = None
+        if not state and states and len(states) > 0:
+            state = states[0]
+            depart_next = states[1].advertised_time_at_location if states[1] else None
+            depart_next_next = (
+                states[2].advertised_time_at_location if 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 +144,7 @@ async def _async_update_data(self) -> 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 +155,6 @@ async def _async_update_data(self) -> 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
diff --git a/homeassistant/components/trafikverket_train/sensor.py b/homeassistant/components/trafikverket_train/sensor.py
index 3aff376ab27206..68865a64cb5af3 100644
--- a/homeassistant/components/trafikverket_train/sensor.py
+++ b/homeassistant/components/trafikverket_train/sensor.py
@@ -105,6 +105,20 @@ class 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,
+    ),
 )
 
 
diff --git a/homeassistant/components/trafikverket_train/strings.json b/homeassistant/components/trafikverket_train/strings.json
index a2c286867b2e22..89542211a92477 100644
--- a/homeassistant/components/trafikverket_train/strings.json
+++ b/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": {
diff --git a/tests/components/trafikverket_train/__init__.py b/tests/components/trafikverket_train/__init__.py
index 060b6a344a17c3..9a02ebbf3b6369 100644
--- a/tests/components/trafikverket_train/__init__.py
+++ b/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"}
diff --git a/tests/components/trafikverket_train/conftest.py b/tests/components/trafikverket_train/conftest.py
new file mode 100644
index 00000000000000..dd9721a694efcd
--- /dev/null
+++ b/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"],
+    )
diff --git a/tests/components/trafikverket_train/test_config_flow.py b/tests/components/trafikverket_train/test_config_flow.py
index 1accd4b5a55107..c31d05bd038a8b 100644
--- a/tests/components/trafikverket_train/test_config_flow.py
+++ b/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",
diff --git a/tests/components/trafikverket_train/test_sensor.py b/tests/components/trafikverket_train/test_sensor.py
new file mode 100644
index 00000000000000..bd1a0b326f32e2
--- /dev/null
+++ b/tests/components/trafikverket_train/test_sensor.py
@@ -0,0 +1,72 @@
+"""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 homeassistant.config_entries import ConfigEntry
+from homeassistant.const import STATE_UNKNOWN
+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],
+) -> None:
+    """Test the Trafikverket Train sensor."""
+    state1 = hass.states.get("sensor.stockholm_c_to_uppsala_c_departure_time")
+    state2 = hass.states.get("sensor.stockholm_c_to_uppsala_c_departure_state")
+    state3 = hass.states.get("sensor.stockholm_c_to_uppsala_c_actual_time")
+    state4 = hass.states.get("sensor.stockholm_c_to_uppsala_c_other_information")
+    state5 = hass.states.get("sensor.stockholm_c_to_uppsala_c_departure_time_next")
+    state6 = hass.states.get(
+        "sensor.stockholm_c_to_uppsala_c_departure_time_next_after"
+    )
+
+    assert state1.state == "2023-05-01T12:00:00+00:00"
+    assert state2.state == "on_time"
+    assert state3.state == "2023-05-01T12:00:00+00:00"
+    assert state4.state == "Some other info"
+    assert state5.state == "2023-05-01T12:15:00+00:00"
+    assert state6.state == "2023-05-01T12:30:00+00:00"
+
+    assert state1.attributes["icon"] == "mdi:clock"
+    assert state1.attributes["product_filter"] == "Regionaltåg"
+    assert state2.attributes["icon"] == "mdi:clock"
+
+    with patch(
+        "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_get_next_train_stops",
+        return_value=get_trains_next,
+    ):
+        freezer.tick(timedelta(minutes=6))
+        async_fire_time_changed(hass)
+        await hass.async_block_till_done()
+
+    state1 = hass.states.get("sensor.stockholm_c_to_uppsala_c_departure_time")
+    assert state1.state == "2023-05-01T17:00:00+00:00"
+    state2 = hass.states.get("sensor.stockholm_c_to_uppsala_c_other_information")
+    assert state2.state == STATE_UNKNOWN
+
+
+async def test_sensor_single_stop(
+    hass: HomeAssistant,
+    freezer: FrozenDateTimeFactory,
+    entity_registry_enabled_by_default: None,
+    load_int: ConfigEntry,
+    get_trains_next: list[TrainStop],
+) -> None:
+    """Test the Trafikverket Train sensor."""
+    state1 = hass.states.get("sensor.stockholm_c_to_uppsala_c_departure_time_2")
+
+    assert state1.state == "2023-05-01T11:00:00+00:00"
+
+    assert state1.attributes["icon"] == "mdi:clock"
+    assert "product_filter" not in state1.attributes


Code review comments:
diff_hunk: @@ -117,6 +124,18 @@ async def _async_update_data(self) -> TrainData:
                 f"Train departure {when} encountered a problem: {error}"
             ) from error
 
+        depart_next = None
+        depart_next_next = None
+        if not state and states and len(states) > 0:
+            state = states[0]
+            depart_next = states[1].advertised_time_at_location if states[1] else None
path: homeassistant/components/trafikverket_train/coordinator.py
text: @user1:
Since when is `if states[1]` the same like `if key in states`?

@author:
states is a list containing `TrainStop` objects so we only need to check if index 1 of that list has an object and thereby can pull the required info. Or am I not understanding what you mean?

@user1:
If states only has 1 entry, doesn't `states[1]` raise an IndexError?

diff_hunk: @@ -0,0 +1,72 @@
+"""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 homeassistant.config_entries import ConfigEntry
+from homeassistant.const import STATE_UNKNOWN
+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],
+) -> None:
+    """Test the Trafikverket Train sensor."""
+    state1 = hass.states.get("sensor.stockholm_c_to_uppsala_c_departure_time")
+    state2 = hass.states.get("sensor.stockholm_c_to_uppsala_c_departure_state")
+    state3 = hass.states.get("sensor.stockholm_c_to_uppsala_c_actual_time")
+    state4 = hass.states.get("sensor.stockholm_c_to_uppsala_c_other_information")
+    state5 = hass.states.get("sensor.stockholm_c_to_uppsala_c_departure_time_next")
+    state6 = hass.states.get(
+        "sensor.stockholm_c_to_uppsala_c_departure_time_next_after"
+    )
+
+    assert state1.state == "2023-05-01T12:00:00+00:00"
+    assert state2.state == "on_time"
+    assert state3.state == "2023-05-01T12:00:00+00:00"
+    assert state4.state == "Some other info"
+    assert state5.state == "2023-05-01T12:15:00+00:00"
+    assert state6.state == "2023-05-01T12:30:00+00:00"
+
+    assert state1.attributes["icon"] == "mdi:clock"
+    assert state1.attributes["product_filter"] == "Regionaltåg"
+    assert state2.attributes["icon"] == "mdi:clock"
path: tests/components/trafikverket_train/test_sensor.py
text: @user1:
```suggestion
    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
```

@author:
Of course 👍