fix: only adjust headsign to last predicted stop on subway routes

apps/state/lib/state/trip/added.ex

@@ -10,7 +10,7 @@ defmodule State.Trip.Added do
     recordable: Model.Trip,
     hibernate: false
 
-  alias Model.{Prediction, Trip}
+  alias Model.{Prediction, Trip, Route}
 
   @impl GenServer
   def init(state) do
@@ -58,7 +58,7 @@ defmodule State.Trip.Added do
          [trip | _] <- State.Trip.by_id(rep_trip_id) do
       stop = parent_or_stop(prediction.stop_id)
 
-      headsign = if stop, do: stop.name, else: trip.headsign
+      headsign = if stop && subway?(trip.route_type), do: stop.name, else: trip.headsign
 
       [
         %{
@@ -80,12 +80,16 @@ defmodule State.Trip.Added do
 
   defp prediction_to_trip_via_shape(prediction) do
     stop = parent_or_stop(prediction.stop_id)
+    route = State.Route.by_id(prediction.route_id)
+
+    stop =
+      if subway?(route),
+        do: stop,
+        else: last_stop_for_prediction_pattern(prediction, stop) || stop
 
     if stop == nil do
       []
     else
-      route = State.Route.by_id(prediction.route_id)
-
       [
         %Trip{
           id: prediction.trip_id,
@@ -103,11 +107,42 @@ defmodule State.Trip.Added do
     end
   end
 
+  defp last_stop_id_on_shape(_, _, nil), do: nil
+
+  defp last_stop_id_on_shape(%{priority: p} = shape, prediction, stop) when p >= 0 do
+    shape_stops =
+      State.StopsOnRoute.by_route_id(
+        prediction.route_id,
+        direction_id: prediction.direction_id,
+        shape_ids: [shape.id]
+      )
+
+    if Enum.any?(shape_stops, &(&1 in [stop.id, stop.parent_station])) do
+      List.last(shape_stops)
+    end
+  end
+
+  defp last_stop_id_on_shape(_, _, _) do
+    nil
+  end
+
+  defp last_stop_for_prediction_pattern(prediction, stop) do
+    [prediction.route_id]
+    |> State.Shape.select_routes(prediction.direction_id)
+    |> Stream.filter(&(&1.route_id == prediction.route_id))
+    |> Enum.find_value(&last_stop_id_on_shape(&1, prediction, stop))
+    |> State.Stop.by_id()
+  end
+
   defp parent_or_stop(stop_id) do
     case State.Stop.by_id(stop_id) do
       %{parent_station: nil} = stop -> stop
       %{parent_station: id} -> State.Stop.by_id(id)
       _other -> nil
     end
   end
+
+  defp subway?(%Route{type: type}), do: subway?(type)
+  defp subway?(type) when type in [0, 1], do: true
+  defp subway?(_), do: false
 end

apps/state/test/state/trip/added_test.exs

@@ -3,6 +3,7 @@ defmodule State.Trip.AddedTest do
   use ExUnit.Case
   import State.Trip.Added
   @trip_id "added_trip"
+  @scheduled_trip_id "scheduled_trip"
   @route_id "route"
   @route_pattern_id "pattern"
   @route_type 3
@@ -15,7 +16,9 @@ defmodule State.Trip.AddedTest do
     schedule_relationship: :added
   }
 
-  setup_all do
+  setup do
+    new_state([])
+
     stops = [
       %Model.Stop{id: "other", name: "Other"},
       %Model.Stop{id: "child", parent_station: "parent", name: "Child"},
@@ -39,11 +42,6 @@ defmodule State.Trip.AddedTest do
     :ok
   end
 
-  setup do
-    new_state([])
-    :ok
-  end
-
   defp insert_predictions(predictions) do
     State.Prediction.new_state(predictions)
     # wait for update
@@ -152,7 +150,7 @@ defmodule State.Trip.AddedTest do
 
   describe "handle_event/4 with shapes" do
     setup do
-      trip_id = "scheduled_trip"
+      trip_id = @scheduled_trip_id
       shape_id = "shape_id"
 
       State.Shape.new_state([
@@ -181,13 +179,26 @@ defmodule State.Trip.AddedTest do
       {:ok, %{shape: shape, prediction: prediction}}
     end
 
-    test "if there's a matching shape for the route/direction, uses the last stop predicted stop",
+    test "if there's a matching shape for the route/direction, uses the last stop from that shape",
          %{prediction: prediction} do
       predictions = [
         %{@prediction | stop_sequence: 3, stop_id: "child"},
         %{@prediction | stop_sequence: 2, stop_id: "other"}
       ]
 
+      insert_predictions(predictions)
+      assert [%{headsign: "Last Stop on Shape"}] = by_id(@trip_id)
+    end
+
+    test "if there's a matching shape for the route/direction and it's a subway route, uses the last predicted stop",
+         %{prediction: prediction} do
+      State.Route.new_state([%Model.Route{id: @route_id, type: 0}])
+
+      predictions = [
+        %{@prediction | stop_sequence: 3, stop_id: "child"},
+        %{@prediction | stop_sequence: 2, stop_id: "other"}
+      ]
+
       insert_predictions(predictions)
       assert [%{headsign: "Parent"}] = by_id(@trip_id)
     end
@@ -210,7 +221,7 @@ defmodule State.Trip.AddedTest do
       prediction: prediction
     } do
       State.Schedule.new_state([
-        %Model.Schedule{trip_id: "scheduled_trip", stop_sequence: 2, stop_id: "shape"}
+        %Model.Schedule{trip_id: @scheduled_trip_id, stop_sequence: 2, stop_id: "shape"}
       ])
 
       State.StopsOnRoute.update!()
@@ -236,19 +247,19 @@ defmodule State.Trip.AddedTest do
     test "creates a trip with revenue value set to :REVENUE",
          %{prediction: prediction} do
       insert_predictions([%{prediction | revenue: :REVENUE}])
-      assert [%{headsign: "Parent", revenue: :REVENUE}] = by_id(@trip_id)
+      assert [%{headsign: "Last Stop on Shape", revenue: :REVENUE}] = by_id(@trip_id)
     end
 
     test "creates a trip with revenue value set to false",
          %{prediction: prediction} do
       insert_predictions([%{prediction | revenue: :NON_REVENUE}])
-      assert [%{headsign: "Parent", revenue: :NON_REVENUE}] = by_id(@trip_id)
+      assert [%{headsign: "Last Stop on Shape", revenue: :NON_REVENUE}] = by_id(@trip_id)
     end
   end
 
   describe "handle_event/4 with route patterns" do
     setup do
-      trip_id = "scheduled_trip"
+      trip_id = @scheduled_trip_id
 
       State.RoutePattern.new_state([
         %Model.RoutePattern{id: @route_pattern_id, representative_trip_id: trip_id}
@@ -272,19 +283,36 @@ defmodule State.Trip.AddedTest do
       end)
     end
 
-    test "if there's a matching route_pattern, use the last predicted stop" do
+    test "if there's a matching route_pattern, use the representative trip" do
+      insert_predictions([%{@prediction | stop_id: "child"}])
+      assert [%{headsign: "Headsign"}] = by_id(@trip_id)
+    end
+
+    test "if there's a matching route_pattern, and it's a subway route, use the last stop" do
+      State.Route.new_state([%Model.Route{id: @route_id, type: 0}])
+
+      State.Trip.new_state([
+        %Model.Trip{
+          id: @scheduled_trip_id,
+          route_id: @route_id,
+          direction_id: 0,
+          headsign: "Headsign",
+          route_type: 0
+        }
+      ])
+
       insert_predictions([%{@prediction | stop_id: "child"}])
       assert [%{headsign: "Parent"}] = by_id(@trip_id)
     end
 
     test "creates a trip with revenue value set to :REVENUE" do
       insert_predictions([%{@prediction | stop_id: "child", revenue: :REVENUE}])
-      assert [%{headsign: "Parent", revenue: :REVENUE}] = by_id(@trip_id)
+      assert [%{headsign: "Headsign", revenue: :REVENUE}] = by_id(@trip_id)
     end
 
     test "creates a trip with revenue value set to :NON_REVENUE" do
       insert_predictions([%{@prediction | stop_id: "child", revenue: :NON_REVENUE}])
-      assert [%{headsign: "Parent", revenue: :NON_REVENUE}] = by_id(@trip_id)
+      assert [%{headsign: "Headsign", revenue: :NON_REVENUE}] = by_id(@trip_id)
     end
   end
 end