feat: pull stop/trip information in with Vehicle stream

apps/vehicles/lib/parser.ex

@@ -3,13 +3,16 @@ defmodule Vehicles.Parser do
   alias Vehicles.Vehicle
 
   @spec parse(JsonApi.Item.t()) :: Vehicle.t()
-  def parse(%JsonApi.Item{id: id, attributes: attributes, relationships: relationships}) do
+  def parse(
+        %JsonApi.Item{id: id, attributes: attributes, relationships: relationships},
+        opts \\ []
+      ) do
     %Vehicle{
       id: id,
       route_id: optional_id(relationships["route"]),
       trip_id: optional_id(relationships["trip"]),
-      shape_id: shape(relationships["trip"]),
-      stop_id: stop_id(relationships["stop"]),
+      shape_id: shape(relationships["trip"], opts),
+      stop_id: stop_id(relationships["stop"], opts),
       direction_id: attributes["direction_id"],
       status: status(attributes["current_status"]),
       longitude: attributes["longitude"],
@@ -28,27 +31,23 @@ defmodule Vehicles.Parser do
   defp optional_id([]), do: nil
   defp optional_id([%JsonApi.Item{id: id}]), do: id
 
-  @spec stop_id([JsonApi.Item.t()]) :: Stops.Stop.id_t() | nil
-  defp stop_id([%JsonApi.Item{id: stop_id}]) do
-    case Stops.Repo.get_parent(stop_id) do
-      %Stops.Stop{id: id} -> id
-      _ -> nil
-    end
+  @spec stop_id([JsonApi.Item.t()], Keyword.t()) :: Stops.Stop.id_t() | nil
+  defp stop_id([%JsonApi.Item{id: stop_id}], opts) do
+    parent_stations = Keyword.get(opts, :parent_stations, %{})
+    Map.get(parent_stations, stop_id, stop_id)
   end
 
-  defp stop_id(_) do
+  defp stop_id(_, _) do
     nil
   end
 
-  @spec shape([JsonApi.Item.t()]) :: Routes.Shape.id_t() | nil
-  defp shape([%JsonApi.Item{id: trip_id}]) do
-    case Schedules.Repo.trip(trip_id) do
-      %Schedules.Trip{shape_id: id} -> id
-      _ -> nil
-    end
+  @spec shape([JsonApi.Item.t()], Keyword.t()) :: Routes.Shape.id_t() | nil
+  defp shape([%JsonApi.Item{id: trip_id}], opts) do
+    trip_shapes = Keyword.get(opts, :trip_shapes, %{})
+    Map.get(trip_shapes, trip_id, nil)
   end
 
-  defp shape(_) do
+  defp shape(_, _) do
     nil
   end
 

apps/vehicles/lib/repo.ex

@@ -38,8 +38,7 @@ defmodule Vehicles.Repo do
         # For simplicity's sake, our repo always returns a single vehicle for a trip.
         # We sort by vehicle id to ensure that all servers will return the same vehicle.
         # See https://github.com/mbta/dotcom/pull/2753 for further details.
-        [vehicle | _] = Enum.sort_by(vehicles, fn %{id: id} -> id end)
-        vehicle
+        Enum.min_by(vehicles, fn %{id: id} -> id end)
     end
   end
 

apps/vehicles/lib/stream.ex

@@ -23,11 +23,13 @@ defmodule Vehicles.Stream do
   def init(opts) do
     producer_consumer = Keyword.fetch!(opts, :subscribe_to)
     broadcast_fn = Keyword.get(opts, :broadcast_fn, &PubSub.broadcast/3)
-    {:consumer, %{broadcast_fn: broadcast_fn}, subscribe_to: [producer_consumer]}
+
+    {:consumer, %{broadcast_fn: broadcast_fn, trip_shapes: %{}, parent_stations: %{}},
+     subscribe_to: [producer_consumer]}
   end
 
   def handle_events(events, _from, state) do
-    :ok = Enum.each(events, &send_event(&1, state.broadcast_fn))
+    state = Enum.reduce(events, state, &send_event/2)
     {:noreply, [], state}
   end
 
@@ -36,33 +38,51 @@ defmodule Vehicles.Stream do
            event: :remove,
            data: %JsonApi{data: data}
          },
-         broadcast_fn
+         state
        ) do
+    state = Enum.reduce(data, state, &remove_items/2)
+
     data
+    |> Enum.filter(&(&1.type == "vehicle"))
     |> Enum.map(& &1.id)
-    |> broadcast(:remove, broadcast_fn)
+    |> broadcast(:remove, state.broadcast_fn)
+
+    state
   end
 
   defp send_event(
          %V3Api.Stream.Event{
            event: type,
            data: %JsonApi{data: data}
          },
-         broadcast_fn
+         state
        ) do
+    state = Enum.reduce(data, state, &add_items/2)
+
     data
-    |> Enum.map(&Parser.parse/1)
-    |> broadcast(type, broadcast_fn)
+    |> Enum.filter(&(&1.type == "vehicle"))
+    |> Enum.map(
+      &Parser.parse(&1, parent_stations: state.parent_stations, trip_shapes: state.trip_shapes)
+    )
+    |> broadcast(type, state.broadcast_fn)
+
+    state
   end
 
   @typep broadcast_fn :: (atom, String.t(), any -> :ok | {:error, any})
   @spec broadcast([Vehicles.Vehicle.t() | String.t()], event_type, broadcast_fn) :: :ok
-  defp broadcast(data, type, broadcast_fn) do
+  defp broadcast(data, type, broadcast_fn)
+
+  defp broadcast([_ | _] = data, type, broadcast_fn) do
     Vehicles.PubSub
     |> broadcast_fn.("vehicles", {type, data})
     |> do_broadcast()
   end
 
+  defp broadcast(_data, _type, _broadcast_fn) do
+    :ok
+  end
+
   @spec do_broadcast(:ok | {:error, any}) :: :ok
   defp do_broadcast(:ok) do
     :ok
@@ -71,4 +91,46 @@ defmodule Vehicles.Stream do
   defp do_broadcast({:error, error}) do
     Logger.error("module=#{__MODULE__} error=#{inspect(error)}")
   end
+
+  defp add_items(%{type: "stop", id: id} = item, state) do
+    parent_station_id =
+      case item do
+        %{relationships: %{"parent_station" => [%{id: parent_station_id} | _]}} ->
+          parent_station_id
+
+        _ ->
+          id
+      end
+
+    %{state | parent_stations: Map.put(state.parent_stations, id, parent_station_id)}
+  end
+
+  defp add_items(%{type: "trip", id: id} = item, state) do
+    shape_id =
+      case item do
+        %{relationships: %{"shape" => [%{id: shape_id} | _]}} ->
+          shape_id
+
+        _ ->
+          nil
+      end
+
+    %{state | trip_shapes: Map.put(state.trip_shapes, id, shape_id)}
+  end
+
+  defp add_items(%{type: "vehicle"}, state) do
+    state
+  end
+
+  defp remove_items(%{type: "stop", id: id}, state) do
+    %{state | parent_stations: Map.delete(state.parent_stations, id)}
+  end
+
+  defp remove_items(%{type: "trip", id: id}, state) do
+    %{state | trip_shapes: Map.delete(state.trip_shapes, id)}
+  end
+
+  defp remove_items(%{}, state) do
+    state
+  end
 end

apps/vehicles/lib/vehicles.ex

@@ -39,7 +39,7 @@ defmodule Vehicles do
       V3Api.Stream.build_options(
         name: Vehicles.Api.SSES,
         path:
-          "/vehicles?fields[vehicle]=direction_id,current_status,longitude,latitude,bearing,occupancy_status"
+          "/vehicles?fields[vehicle]=direction_id,current_status,longitude,latitude,bearing,occupancy_status&fields[stop]=&fields[trip]=&include=stop,trip"
       )
 
     [

apps/vehicles/test/parser_test.exs

@@ -1,8 +1,7 @@
 defmodule Vehicles.ParserTest do
-  use ExUnit.Case
+  use ExUnit.Case, async: true
   alias Vehicles.Vehicle
   import Vehicles.Parser
-  import Mock
 
   @item %JsonApi.Item{
     attributes: %{
@@ -20,6 +19,10 @@ defmodule Vehicles.ParserTest do
     },
     type: "vehicle"
   }
+  @opts [
+    parent_stations: %{"72" => "place-72"},
+    trip_shapes: %{"25" => "25-shape"}
+  ]
 
   describe "parse/1" do
     test "parses an API response into a Vehicle struct" do
@@ -37,6 +40,8 @@ defmodule Vehicles.ParserTest do
       }
 
       assert parse(@item) == expected
+      # does not crash
+      parse(@item, @opts)
     end
 
     test "can handle a missing trip" do
@@ -56,6 +61,8 @@ defmodule Vehicles.ParserTest do
       }
 
       assert parse(item) == expected
+      # does not crash
+      parse(item, @opts)
     end
 
     test "can handle a missing stop" do
@@ -75,6 +82,8 @@ defmodule Vehicles.ParserTest do
       }
 
       assert parse(item) == expected
+      # does not crash
+      parse(item, @opts)
     end
 
     test "can handle a missing route" do
@@ -96,47 +105,18 @@ defmodule Vehicles.ParserTest do
       }
 
       assert parse(item) == expected
+      # does not crash
+      parse(item, @opts)
     end
 
     test "fetches parent stop if present" do
-      with_mock(Stops.Repo, [], get_parent: fn "72" -> %Stops.Stop{id: "place-72"} end) do
-        expected = %Vehicle{
-          id: "y1799",
-          route_id: "1",
-          stop_id: "place-72",
-          trip_id: "25",
-          shape_id: nil,
-          direction_id: 1,
-          status: :stopped,
-          latitude: 2.2,
-          longitude: 1.1,
-          bearing: 140
-        }
-
-        %Vehicle{} = parsed_vehicle = parse(@item)
-        assert parsed_vehicle == expected, "parsed vehicle is #{inspect(parsed_vehicle)}"
-      end
+      assert %Vehicle{stop_id: "place-72"} = parse(@item, @opts)
     end
 
     test "fetches shape if trip is present" do
-      with_mock(Schedules.Repo,
-        trip: fn "25" -> %Schedules.Trip{shape_id: "25-shape"} end
-      ) do
-        expected = %Vehicle{
-          id: "y1799",
-          route_id: "1",
-          stop_id: "72",
-          trip_id: "25",
-          shape_id: "25-shape",
-          direction_id: 1,
-          status: :stopped,
-          latitude: 2.2,
-          longitude: 1.1,
-          bearing: 140
-        }
-
-        assert parse(@item) == expected
-      end
+      assert %Vehicle{
+               shape_id: "25-shape"
+             } = parse(@item, @opts)
     end
 
     test "can handle occupancy status" do

apps/vehicles/test/stream_test.exs

@@ -2,6 +2,38 @@ defmodule Vehicles.StreamTest do
   use ExUnit.Case
   import ExUnit.CaptureLog
 
+  @stops %JsonApi{
+    data: [
+      %JsonApi.Item{
+        type: "stop",
+        id: "stop",
+        relationships: %{
+          "parent_station" => [
+            %JsonApi.Item{
+              type: "stop",
+              id: "place-stop"
+            }
+          ]
+        }
+      }
+    ]
+  }
+  @trips %JsonApi{
+    data: [
+      %JsonApi.Item{
+        type: "trip",
+        id: "trip",
+        relationships: %{
+          "shape" => [
+            %JsonApi.Item{
+              type: "shape",
+              id: "trip-shape"
+            }
+          ]
+        }
+      }
+    ]
+  }
   @vehicles %JsonApi{
     data: [
       %JsonApi.Item{
@@ -23,7 +55,9 @@ defmodule Vehicles.StreamTest do
   setup tags do
     {:ok, mock_api} =
       GenStage.from_enumerable([
-        %V3Api.Stream.Event{event: :reset, data: @vehicles}
+        %V3Api.Stream.Event{event: :add, data: @stops},
+        %V3Api.Stream.Event{event: :add, data: @trips},
+        %V3Api.Stream.Event{event: :add, data: @vehicles}
       ])
 
     name = :"stream_test_#{tags.line}"
@@ -47,7 +81,14 @@ defmodule Vehicles.StreamTest do
                  subscribe_to: mock_api
                )
 
-      assert_receive {:reset, [%Vehicles.Vehicle{id: "vehicle1"}]}
+      assert_receive {:add,
+                      [
+                        %Vehicles.Vehicle{
+                          id: "vehicle1",
+                          stop_id: "place-stop",
+                          shape_id: "trip-shape"
+                        }
+                      ]}
     end
 
     test "publishes :remove events as a list of IDs", %{name: name} do