Add full test coverage to Trafikverket Train (#106247)

* Add full test coverage for Trafikverket Train

* Coordinator full coverage

* Now full coverage

* Add util

.coveragerc

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

tests/components/trafikverket_train/conftest.py

@@ -151,7 +151,7 @@ def fixture_get_train_stop() -> 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),
+        estimated_time_at_location=None,
         time_at_location=datetime(2023, 5, 1, 11, 0, tzinfo=dt_util.UTC),
         other_information=None,
         deviations=None,

tests/components/trafikverket_train/snapshots/test_init.ambr

@@ -0,0 +1,16 @@
+# serializer version: 1
+# name: test_auth_failed
+  FlowResultSnapshot({
+    'context': dict({
+      'entry_id': '1',
+      'source': 'reauth',
+      'title_placeholders': dict({
+        'name': 'Mock Title',
+      }),
+      'unique_id': '321',
+    }),
+    'flow_id': <ANY>,
+    'handler': 'trafikverket_train',
+    'step_id': 'reauth_confirm',
+  })
+# ---

tests/components/trafikverket_train/snapshots/test_sensor.ambr

@@ -214,3 +214,18 @@
     'state': '2023-05-01T11:00:00+00:00',
   })
 # ---
+# name: test_sensor_update_auth_failure
+  FlowResultSnapshot({
+    'context': dict({
+      'entry_id': '1',
+      'source': 'reauth',
+      'title_placeholders': dict({
+        'name': 'Mock Title',
+      }),
+      'unique_id': "stockholmc-uppsalac--['mon', 'tue', 'wed', 'thu', 'fri', 'sat', 'sun']",
+    }),
+    'flow_id': <ANY>,
+    'handler': 'trafikverket_train',
+    'step_id': 'reauth_confirm',
+  })
+# ---

tests/components/trafikverket_train/test_config_flow.py

@@ -11,6 +11,7 @@ from pytrafikverket.exceptions import (
     NoTrainStationFound,
     UnknownError,
 )
+from pytrafikverket.trafikverket_train import TrainStop
 
 from homeassistant import config_entries
 from homeassistant.components.trafikverket_train.const import (
@@ -442,7 +443,11 @@ async def test_reauth_flow_error_departures(
     }
 
 
-async def test_options_flow(hass: HomeAssistant) -> None:
+async def test_options_flow(
+    hass: HomeAssistant,
+    get_trains: list[TrainStop],
+    get_train_stop: TrainStop,
+) -> None:
     """Test a reauthentication flow."""
     entry = MockConfigEntry(
         domain=DOMAIN,
@@ -459,36 +464,41 @@ async def test_options_flow(hass: HomeAssistant) -> None:
     entry.add_to_hass(hass)
 
     with patch(
-        "homeassistant.components.trafikverket_train.async_setup_entry",
+        "homeassistant.components.trafikverket_train.TrafikverketTrain.async_get_train_station",
-        return_value=True,
+    ), 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,
     ):
         assert await hass.config_entries.async_setup(entry.entry_id)
         await hass.async_block_till_done()
 
-    result = await hass.config_entries.options.async_init(entry.entry_id)
+        result = await hass.config_entries.options.async_init(entry.entry_id)
 
-    assert result["type"] == FlowResultType.FORM
+        assert result["type"] == FlowResultType.FORM
-    assert result["step_id"] == "init"
+        assert result["step_id"] == "init"
 
-    result = await hass.config_entries.options.async_configure(
+        result = await hass.config_entries.options.async_configure(
-        result["flow_id"],
+            result["flow_id"],
-        user_input={"filter_product": "SJ Regionaltåg"},
+            user_input={"filter_product": "SJ Regionaltåg"},
-    )
+        )
-    await hass.async_block_till_done()
+        await hass.async_block_till_done()
 
-    assert result["type"] == FlowResultType.CREATE_ENTRY
+        assert result["type"] == FlowResultType.CREATE_ENTRY
-    assert result["data"] == {"filter_product": "SJ Regionaltåg"}
+        assert result["data"] == {"filter_product": "SJ Regionaltåg"}
 
-    result = await hass.config_entries.options.async_init(entry.entry_id)
+        result = await hass.config_entries.options.async_init(entry.entry_id)
 
-    assert result["type"] == FlowResultType.FORM
+        assert result["type"] == FlowResultType.FORM
-    assert result["step_id"] == "init"
+        assert result["step_id"] == "init"
 
-    result = await hass.config_entries.options.async_configure(
+        result = await hass.config_entries.options.async_configure(
-        result["flow_id"],
+            result["flow_id"],
-        user_input={"filter_product": ""},
+            user_input={"filter_product": ""},
-    )
+        )
-    await hass.async_block_till_done()
+        await hass.async_block_till_done()
 
-    assert result["type"] == FlowResultType.CREATE_ENTRY
+        assert result["type"] == FlowResultType.CREATE_ENTRY
-    assert result["data"] == {"filter_product": None}
+        assert result["data"] == {"filter_product": None}

tests/components/trafikverket_train/test_init.py

@@ -0,0 +1,143 @@
+"""Test for Trafikverket Train component Init."""
+from __future__ import annotations
+
+from unittest.mock import patch
+
+from pytrafikverket.exceptions import InvalidAuthentication, NoTrainStationFound
+from pytrafikverket.trafikverket_train import TrainStop
+from syrupy.assertion import SnapshotAssertion
+
+from homeassistant import config_entries
+from homeassistant.components.trafikverket_train.const import DOMAIN
+from homeassistant.config_entries import SOURCE_REAUTH, SOURCE_USER
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.entity_registry import EntityRegistry
+
+from . import ENTRY_CONFIG, OPTIONS_CONFIG
+
+from tests.common import MockConfigEntry
+
+
+async def test_unload_entry(hass: HomeAssistant, get_trains: list[TrainStop]) -> None:
+    """Test unload an entry."""
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        source=SOURCE_USER,
+        data=ENTRY_CONFIG,
+        options=OPTIONS_CONFIG,
+        entry_id="1",
+        unique_id="321",
+    )
+    entry.add_to_hass(hass)
+
+    with patch(
+        "homeassistant.components.trafikverket_train.TrafikverketTrain.async_get_train_station",
+    ), patch(
+        "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_get_next_train_stops",
+        return_value=get_trains,
+    ) as mock_tv_train:
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert entry.state is config_entries.ConfigEntryState.LOADED
+    assert len(mock_tv_train.mock_calls) == 1
+
+    assert await hass.config_entries.async_unload(entry.entry_id)
+    await hass.async_block_till_done()
+    assert entry.state is config_entries.ConfigEntryState.NOT_LOADED
+
+
+async def test_auth_failed(
+    hass: HomeAssistant,
+    get_trains: list[TrainStop],
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test authentication failed."""
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        source=SOURCE_USER,
+        data=ENTRY_CONFIG,
+        options=OPTIONS_CONFIG,
+        entry_id="1",
+        unique_id="321",
+    )
+    entry.add_to_hass(hass)
+
+    with patch(
+        "homeassistant.components.trafikverket_train.TrafikverketTrain.async_get_train_station",
+        side_effect=InvalidAuthentication,
+    ):
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert entry.state is config_entries.ConfigEntryState.SETUP_ERROR
+
+    active_flows = entry.async_get_active_flows(hass, (SOURCE_REAUTH))
+    for flow in active_flows:
+        assert flow == snapshot
+
+
+async def test_no_stations(
+    hass: HomeAssistant,
+    get_trains: list[TrainStop],
+    snapshot: SnapshotAssertion,
+) -> None:
+    """Test stations are missing."""
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        source=SOURCE_USER,
+        data=ENTRY_CONFIG,
+        options=OPTIONS_CONFIG,
+        entry_id="1",
+        unique_id="321",
+    )
+    entry.add_to_hass(hass)
+
+    with patch(
+        "homeassistant.components.trafikverket_train.TrafikverketTrain.async_get_train_station",
+        side_effect=NoTrainStationFound,
+    ):
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert entry.state is config_entries.ConfigEntryState.SETUP_RETRY
+
+
+async def test_migrate_entity_unique_id(
+    hass: HomeAssistant,
+    get_trains: list[TrainStop],
+    snapshot: SnapshotAssertion,
+    entity_registry: EntityRegistry,
+) -> None:
+    """Test migration of entity unique id in old format."""
+    entry = MockConfigEntry(
+        domain=DOMAIN,
+        source=SOURCE_USER,
+        data=ENTRY_CONFIG,
+        options=OPTIONS_CONFIG,
+        entry_id="1",
+        unique_id="321",
+    )
+    entry.add_to_hass(hass)
+
+    entity = entity_registry.async_get_or_create(
+        DOMAIN,
+        "sensor",
+        "incorrect_unique_id",
+        config_entry=entry,
+        original_name="Stockholm C to Uppsala C",
+    )
+
+    with patch(
+        "homeassistant.components.trafikverket_train.TrafikverketTrain.async_get_train_station",
+    ), patch(
+        "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_get_next_train_stops",
+        return_value=get_trains,
+    ):
+        await hass.config_entries.async_setup(entry.entry_id)
+        await hass.async_block_till_done()
+
+    assert entry.state is config_entries.ConfigEntryState.LOADED
+
+    entity = entity_registry.async_get(entity.entity_id)
+    assert entity.unique_id == f"{entry.entry_id}-departure_time"

tests/components/trafikverket_train/test_sensor.py

@@ -5,10 +5,12 @@ from datetime import timedelta
 from unittest.mock import patch
 
 from freezegun.api import FrozenDateTimeFactory
+from pytrafikverket.exceptions import InvalidAuthentication, NoTrainAnnouncementFound
 from pytrafikverket.trafikverket_train import TrainStop
 from syrupy.assertion import SnapshotAssertion
 
-from homeassistant.config_entries import ConfigEntry
+from homeassistant.config_entries import SOURCE_REAUTH, ConfigEntry
+from homeassistant.const import STATE_UNAVAILABLE
 from homeassistant.core import HomeAssistant
 
 from tests.common import async_fire_time_changed
@@ -72,3 +74,84 @@ async def test_sensor_single_stop(
     assert state.state == "2023-05-01T11:00:00+00:00"
 
     assert state == snapshot
+
+
+async def test_sensor_update_auth_failure(
+    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 with authentication update failure."""
+    state = hass.states.get("sensor.stockholm_c_to_uppsala_c_departure_time_2")
+    assert state.state == "2023-05-01T11:00:00+00:00"
+
+    with patch(
+        "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_get_next_train_stops",
+        side_effect=InvalidAuthentication,
+    ), patch(
+        "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_get_train_stop",
+        side_effect=InvalidAuthentication,
+    ):
+        freezer.tick(timedelta(minutes=6))
+        async_fire_time_changed(hass)
+        await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.stockholm_c_to_uppsala_c_departure_time_2")
+    assert state.state == STATE_UNAVAILABLE
+    active_flows = load_int.async_get_active_flows(hass, (SOURCE_REAUTH))
+    for flow in active_flows:
+        assert flow == snapshot
+
+
+async def test_sensor_update_failure(
+    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 with update failure."""
+    state = hass.states.get("sensor.stockholm_c_to_uppsala_c_departure_time_2")
+    assert state.state == "2023-05-01T11:00:00+00:00"
+
+    with patch(
+        "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_get_next_train_stops",
+        side_effect=NoTrainAnnouncementFound,
+    ), patch(
+        "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_get_train_stop",
+        side_effect=NoTrainAnnouncementFound,
+    ):
+        freezer.tick(timedelta(minutes=6))
+        async_fire_time_changed(hass)
+        await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.stockholm_c_to_uppsala_c_departure_time_2")
+    assert state.state == STATE_UNAVAILABLE
+
+
+async def test_sensor_update_failure_no_state(
+    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 with update failure from empty state."""
+    state = hass.states.get("sensor.stockholm_c_to_uppsala_c_departure_time_2")
+    assert state.state == "2023-05-01T11:00:00+00:00"
+
+    with patch(
+        "homeassistant.components.trafikverket_train.coordinator.TrafikverketTrain.async_get_train_stop",
+        return_value=None,
+    ):
+        freezer.tick(timedelta(minutes=6))
+        async_fire_time_changed(hass)
+        await hass.async_block_till_done()
+
+    state = hass.states.get("sensor.stockholm_c_to_uppsala_c_departure_time_2")
+    assert state.state == STATE_UNAVAILABLE

tests/components/trafikverket_train/test_util.py

@@ -0,0 +1,25 @@
+"""The test for the Trafikverket train utils."""
+from __future__ import annotations
+
+from datetime import datetime
+
+from freezegun.api import FrozenDateTimeFactory
+
+from homeassistant.components.trafikverket_train.util import next_departuredate
+from homeassistant.const import WEEKDAYS
+from homeassistant.util import dt as dt_util
+
+
+async def test_sensor_next(
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test the Trafikverket Train utils."""
+
+    assert next_departuredate(WEEKDAYS) == dt_util.now().date()
+    freezer.move_to(datetime(2023, 12, 22))  # Friday
+    assert (
+        next_departuredate(["mon", "tue", "wed", "thu"])
+        == datetime(2023, 12, 25).date()
+    )
+    freezer.move_to(datetime(2023, 12, 25))  # Monday
+    assert next_departuredate(["fri", "sat", "sun"]) == datetime(2023, 12, 29).date()