Schedule a final poll refresh before streaming out the notifications

2 files changed, 12 insertions, 2 deletions

diff --git a/lib/pleroma/workers/poll_worker.ex b/lib/pleroma/workers/poll_worker.ex
index a61c5eac1..574daa9ba 100644
--- a/lib/pleroma/workers/poll_worker.ex
+++ b/lib/pleroma/workers/poll_worker.ex
@@ -22,6 +22,10 @@ defmodule Pleroma.Workers.PollWorker do
   def perform(%Job{args: %{"op" => "poll_end", "activity_id" => activity_id}}) do
     with {_, %Activity{} = activity} <- {:activity, Activity.get_by_id(activity_id)},
          {:ok, notifications} <- Notification.create_poll_notifications(activity) do
+      # Schedule a final refresh
+      __MODULE__.new(%{"op" => "refresh", "activity_id" => activity_id})
+      |> Oban.insert()
+
       Notification.stream(notifications)
     else
       {:activity, nil} -> {:cancel, :poll_activity_not_found}
@@ -32,8 +36,8 @@ defmodule Pleroma.Workers.PollWorker do
   def perform(%Job{args: %{"op" => "refresh", "activity_id" => activity_id}}) do
     with {_, %Activity{object: object}} <-
            {:activity, Activity.get_by_id_with_object(activity_id)},
-          {:ok, naive_closed} <- NaiveDateTime.from_iso8601(object.data["closed"]),
-          {_, :lt} <- {:closed_compare, NaiveDateTime.compare(object.updated_at, naive_closed)},
+         {:ok, naive_closed} <- NaiveDateTime.from_iso8601(object.data["closed"]),
+         {_, :lt} <- {:closed_compare, NaiveDateTime.compare(object.updated_at, naive_closed)},
          {_, {:ok, _object}} <- {:refetch, Fetcher.refetch_object(object)} do
       stream_update(activity_id)
 
diff --git a/test/pleroma/workers/poll_worker_test.exs b/test/pleroma/workers/poll_worker_test.exs
index 749df8aff..e1c67f057 100644
--- a/test/pleroma/workers/poll_worker_test.exs
+++ b/test/pleroma/workers/poll_worker_test.exs
@@ -44,6 +44,12 @@ defmodule Pleroma.Workers.PollWorkerTest do
       # Ensure notifications were streamed out when job executes
       assert called(Pleroma.Web.Streamer.stream(["user", "user:notification"], :_))
       assert called(Pleroma.Web.Push.send(:_))
+
+      # Ensure we scheduled a final refresh of the poll
+      assert_enqueued(
+        worker: PollWorker,
+        args: %{"op" => "refresh", "activity_id" => activity.id}
+      )
     end
   end
 end