From 6ba9b2f28e1010ff3dcd619f7a0eda749b97d8ce Mon Sep 17 00:00:00 2001
From: Kieran Eglin <kieran.eglin@gmail.com>
Date: Thu, 29 Feb 2024 15:23:47 -0800
Subject: [PATCH] Adds new maybe_limit method for queries

---
 lib/pinchflat/media.ex        |  9 +++++++--
 lib/pinchflat/repo.ex         | 11 +++++++++++
 test/pinchflat/media_test.exs | 17 +++++++++++++++++
 test/pinchflat/repo_test.exs  | 29 +++++++++++++++++++++++++++++
 4 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/lib/pinchflat/media.ex b/lib/pinchflat/media.ex
index e9fb633..6f93d1b 100644
--- a/lib/pinchflat/media.ex
+++ b/lib/pinchflat/media.ex
@@ -41,12 +41,14 @@ defmodule Pinchflat.Media do
 
   Returns [%MediaItem{}, ...].
   """
-  def list_pending_media_items_for(%Source{} = source) do
+  def list_pending_media_items_for(%Source{} = source, opts \\ []) do
+    limit = Keyword.get(opts, :limit, nil)
     media_profile = Repo.preload(source, :media_profile).media_profile
 
     MediaItem
     |> where([mi], mi.source_id == ^source.id and is_nil(mi.media_filepath))
     |> where(^build_format_clauses(media_profile))
+    |> Repo.maybe_limit(limit)
     |> Repo.all()
   end
 
@@ -55,9 +57,12 @@ defmodule Pinchflat.Media do
 
   Returns [%MediaItem{}, ...].
   """
-  def list_downloaded_media_items_for(%Source{} = source) do
+  def list_downloaded_media_items_for(%Source{} = source, opts \\ []) do
+    limit = Keyword.get(opts, :limit, nil)
+
     MediaItem
     |> where([mi], mi.source_id == ^source.id and not is_nil(mi.media_filepath))
+    |> Repo.maybe_limit(limit)
     |> Repo.all()
   end
 
diff --git a/lib/pinchflat/repo.ex b/lib/pinchflat/repo.ex
index 1ade75a..a4d3b85 100644
--- a/lib/pinchflat/repo.ex
+++ b/lib/pinchflat/repo.ex
@@ -3,6 +3,8 @@ defmodule Pinchflat.Repo do
     otp_app: :pinchflat,
     adapter: Ecto.Adapters.SQLite3
 
+  import Ecto.Query, warn: false
+
   @doc """
   It's not immediately obvious if an Oban job qualifies as unique, so this method
   attempts creating a job and checks for the `conflict?` field in the returned job.
@@ -16,4 +18,13 @@ defmodule Pinchflat.Repo do
       err -> err
     end
   end
+
+  @doc """
+  Applies a limit to a query if provided, otherwise returns the query as-is.
+
+  Returns %Ecto.Query{}.
+  """
+  def maybe_limit(query, limit) do
+    if limit, do: limit(query, ^limit), else: query
+  end
 end
diff --git a/test/pinchflat/media_test.exs b/test/pinchflat/media_test.exs
index fdd87d2..15ca76d 100644
--- a/test/pinchflat/media_test.exs
+++ b/test/pinchflat/media_test.exs
@@ -64,6 +64,14 @@ defmodule Pinchflat.MediaTest do
 
       assert Media.list_pending_media_items_for(source) == []
     end
+
+    test "optionally accepts a limit" do
+      source = source_fixture()
+      media_item = media_item_fixture(%{source_id: source.id, media_filepath: nil})
+
+      assert Media.list_pending_media_items_for(source, limit: 1) == [media_item]
+      assert Media.list_pending_media_items_for(source, limit: 0) == []
+    end
   end
 
   describe "list_pending_media_items_for/1 when testing shorts" do
@@ -196,6 +204,15 @@ defmodule Pinchflat.MediaTest do
 
       assert Media.list_downloaded_media_items_for(source) == [media_item]
     end
+
+    test "optionally accepts a limit" do
+      source = source_fixture()
+      _media_item = media_item_fixture(%{source_id: source.id, media_filepath: nil})
+      media_item = media_item_fixture(%{source_id: source.id, media_filepath: "/video/#{Faker.File.file_name(:video)}"})
+
+      assert Media.list_downloaded_media_items_for(source, limit: 1) == [media_item]
+      assert Media.list_downloaded_media_items_for(source, limit: 0) == []
+    end
   end
 
   describe "search/1" do
diff --git a/test/pinchflat/repo_test.exs b/test/pinchflat/repo_test.exs
index 0d7eea2..eb465f3 100644
--- a/test/pinchflat/repo_test.exs
+++ b/test/pinchflat/repo_test.exs
@@ -1,6 +1,9 @@
 defmodule Pinchflat.RepoTest do
   use Pinchflat.DataCase
+  import Pinchflat.ProfilesFixtures
 
+  alias Pinchflat.Repo
+  alias Pinchflat.Profiles.MediaProfile
   alias Pinchflat.JobFixtures.TestJobWorker
 
   describe "insert_unique_job/1" do
@@ -23,4 +26,30 @@ defmodule Pinchflat.RepoTest do
       assert {:error, _} = Pinchflat.Repo.insert_unique_job(%Ecto.Changeset{})
     end
   end
+
+  describe "maybe_limit/2" do
+    test "applies a limit if provided" do
+      media_profile_fixture()
+      media_profile_fixture()
+
+      result =
+        MediaProfile
+        |> Repo.maybe_limit(1)
+        |> Repo.aggregate(:count, :id)
+
+      assert result == 1
+    end
+
+    test "does not apply a limit if not provided" do
+      media_profile_fixture()
+      media_profile_fixture()
+
+      result =
+        MediaProfile
+        |> Repo.maybe_limit(nil)
+        |> Repo.aggregate(:count, :id)
+
+      assert result == 2
+    end
+  end
 end