From 4a64af7d8ed1fbf74ce0e05b94cec91ec9bd19f3 Mon Sep 17 00:00:00 2001
From: kotva006 <ryan.kotval@gmail.com>
Date: Wed, 8 May 2024 13:39:37 -0500
Subject: [PATCH] Updated the prediction schedule and corresponding test

---
 lib/predicted_schedule.ex        | 7 ++++++-
 test/predicted_schedule_test.exs | 8 ++++----
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/lib/predicted_schedule.ex b/lib/predicted_schedule.ex
index cc5f11866a..7ce4305a7b 100644
--- a/lib/predicted_schedule.ex
+++ b/lib/predicted_schedule.ex
@@ -19,9 +19,14 @@ defmodule PredictedSchedule do
           prediction: Prediction.t() | nil
         }
 
+  @predictions_repo Application.compile_env!(:dotcom, :repo_modules)[:predictions]
+
   def get(route_id, stop_id, opts \\ []) do
     schedules_fn = Keyword.get(opts, :schedules_fn, &Schedules.Repo.by_route_ids/2)
-    predictions_fn = Keyword.get(opts, :predictions_fn, &Predictions.Repo.all/1)
+
+    predictions_fn =
+      Keyword.get(opts, :predictions_fn, Function.capture(@predictions_repo, :all, 1))
+
     now = Keyword.get(opts, :now, Util.now())
     direction_id = Keyword.get(opts, :direction_id)
     sort_fn = Keyword.get(opts, :sort_fn, &sort_predicted_schedules/1)
diff --git a/test/predicted_schedule_test.exs b/test/predicted_schedule_test.exs
index 74acadc742..ae87d51c43 100644
--- a/test/predicted_schedule_test.exs
+++ b/test/predicted_schedule_test.exs
@@ -5,6 +5,7 @@ defmodule PredictedScheduleTest do
   alias Predictions.Prediction
   import PredictedSchedule
   import Mock
+  import Mox
 
   # set to the end of a month to uncover issues with sorting times as
   # structs, rather than as integers
@@ -174,9 +175,9 @@ defmodule PredictedScheduleTest do
     end
 
     test "attempts to return a list of predicted schedules for tomorrow, after no valid ones are left for today" do
-      predictions_fn = fn _opts ->
+      expect(Predictions.Repo.Mock, :all, fn _ ->
         @trip_predictions
-      end
+      end)
 
       schedules_fn = fn ["Teal"], _opts ->
         @trip_schedules
@@ -186,8 +187,7 @@ defmodule PredictedScheduleTest do
         group: fn _predictions, _schedules, _opts -> [] end do
         get("Teal", "stop1",
           now: Timex.shift(@base_time, minutes: 30),
-          schedules_fn: schedules_fn,
-          predictions_fn: predictions_fn
+          schedules_fn: schedules_fn
         )
 
         assert :meck.num_calls(PredictedSchedule, :group, :_) == 2