From d84c6af55df12700aebfc18963e2b0640aa59728 Mon Sep 17 00:00:00 2001
From: G Johansson <goran.johansson@shiftit.se>
Date: Sat, 14 May 2022 02:42:11 +0200
Subject: [PATCH] Add multiple departures to Trafikverket Ferry (#71484)

---
 .../trafikverket_ferry/coordinator.py         | 18 +++++++-----
 .../components/trafikverket_ferry/sensor.py   | 29 ++++++++++++++++---
 2 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/homeassistant/components/trafikverket_ferry/coordinator.py b/homeassistant/components/trafikverket_ferry/coordinator.py
index 96a1e58c6c9..7c2c64d49f0 100644
--- a/homeassistant/components/trafikverket_ferry/coordinator.py
+++ b/homeassistant/components/trafikverket_ferry/coordinator.py
@@ -77,8 +77,10 @@ class TVDataUpdateCoordinator(DataUpdateCoordinator):
             when = current_time
 
         try:
-            routedata: FerryStop = await self._ferry_api.async_get_next_ferry_stop(
-                self._from, self._to, when
+            routedata: list[
+                FerryStop
+            ] = await self._ferry_api.async_get_next_ferry_stops(
+                self._from, self._to, when, 3
             )
         except ValueError as error:
             raise UpdateFailed(
@@ -86,11 +88,13 @@ class TVDataUpdateCoordinator(DataUpdateCoordinator):
             ) from error
 
         states = {
-            "departure_time": routedata.departure_time,
-            "departure_from": routedata.from_harbor_name,
-            "departure_to": routedata.to_harbor_name,
-            "departure_modified": routedata.modified_time,
-            "departure_information": routedata.other_information,
+            "departure_time": routedata[0].departure_time,
+            "departure_from": routedata[0].from_harbor_name,
+            "departure_to": routedata[0].to_harbor_name,
+            "departure_modified": routedata[0].modified_time,
+            "departure_information": routedata[0].other_information,
+            "departure_time_next": routedata[1].departure_time,
+            "departure_time_next_next": routedata[2].departure_time,
         }
         _LOGGER.debug("States: %s", states)
         return states
diff --git a/homeassistant/components/trafikverket_ferry/sensor.py b/homeassistant/components/trafikverket_ferry/sensor.py
index d93aecd38c1..bab73d72210 100644
--- a/homeassistant/components/trafikverket_ferry/sensor.py
+++ b/homeassistant/components/trafikverket_ferry/sensor.py
@@ -38,7 +38,7 @@ class TrafikverketRequiredKeysMixin:
     """Mixin for required keys."""
 
     value_fn: Callable[[dict[str, Any]], StateType | datetime]
-    info_fn: Callable[[dict[str, Any]], StateType | list]
+    info_fn: Callable[[dict[str, Any]], StateType | list] | None
 
 
 @dataclass
@@ -80,6 +80,24 @@ SENSOR_TYPES: tuple[TrafikverketSensorEntityDescription, ...] = (
         info_fn=lambda data: data["departure_information"],
         entity_registry_enabled_default=False,
     ),
+    TrafikverketSensorEntityDescription(
+        key="departure_time_next",
+        name="Departure Time Next",
+        icon="mdi:clock",
+        device_class=SensorDeviceClass.TIMESTAMP,
+        value_fn=lambda data: as_utc(data["departure_time_next"]),
+        info_fn=None,
+        entity_registry_enabled_default=False,
+    ),
+    TrafikverketSensorEntityDescription(
+        key="departure_time_next_next",
+        name="Departure Time Next After",
+        icon="mdi:clock",
+        device_class=SensorDeviceClass.TIMESTAMP,
+        value_fn=lambda data: as_utc(data["departure_time_next_next"]),
+        info_fn=None,
+        entity_registry_enabled_default=False,
+    ),
 )
 
 
@@ -132,9 +150,12 @@ class FerrySensor(CoordinatorEntity[TVDataUpdateCoordinator], SensorEntity):
             self.coordinator.data
         )
 
-        self._attr_extra_state_attributes = {
-            "other_information": self.entity_description.info_fn(self.coordinator.data),
-        }
+        if self.entity_description.info_fn:
+            self._attr_extra_state_attributes = {
+                "other_information": self.entity_description.info_fn(
+                    self.coordinator.data
+                ),
+            }
 
     @callback
     def _handle_coordinator_update(self) -> None: