From 6434dac1833ebaf4d6c9962fd7a80a74854a6579 Mon Sep 17 00:00:00 2001 From: Paul Swartz Date: Tue, 15 Aug 2023 20:06:07 -0400 Subject: [PATCH] feat: pull stop/trip information in with Vehicle stream This avoids making a bunch of extra API calls for data that the streaming API can already provide. --- apps/vehicles/lib/parser.ex | 33 ++++++------- apps/vehicles/lib/repo.ex | 3 +- apps/vehicles/lib/stream.ex | 78 +++++++++++++++++++++++++++--- apps/vehicles/lib/vehicles.ex | 2 +- apps/vehicles/test/parser_test.exs | 54 +++++++-------------- apps/vehicles/test/stream_test.exs | 45 ++++++++++++++++- 6 files changed, 148 insertions(+), 67 deletions(-) diff --git a/apps/vehicles/lib/parser.ex b/apps/vehicles/lib/parser.ex index 815fbfac40..446ec13317 100644 --- a/apps/vehicles/lib/parser.ex +++ b/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 diff --git a/apps/vehicles/lib/repo.ex b/apps/vehicles/lib/repo.ex index 1254158120..d684e12c7a 100644 --- a/apps/vehicles/lib/repo.ex +++ b/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 diff --git a/apps/vehicles/lib/stream.ex b/apps/vehicles/lib/stream.ex index df0152d66f..e00d647b2e 100644 --- a/apps/vehicles/lib/stream.ex +++ b/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,11 +38,16 @@ 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( @@ -48,21 +55,34 @@ defmodule Vehicles.Stream do 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 diff --git a/apps/vehicles/lib/vehicles.ex b/apps/vehicles/lib/vehicles.ex index 250ee0121f..7bb18a4eac 100644 --- a/apps/vehicles/lib/vehicles.ex +++ b/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" ) [ diff --git a/apps/vehicles/test/parser_test.exs b/apps/vehicles/test/parser_test.exs index 6abec62dfa..ba8b3c581c 100644 --- a/apps/vehicles/test/parser_test.exs +++ b/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 diff --git a/apps/vehicles/test/stream_test.exs b/apps/vehicles/test/stream_test.exs index d1439e279a..ab2c18e0eb 100644 --- a/apps/vehicles/test/stream_test.exs +++ b/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