add helper to assert that a LiveView will receive a message

lib/phoenix_live_view/test/live_view_test.ex

@@ -1817,6 +1817,156 @@ defmodule Phoenix.LiveViewTest do
     end
   end
 
+  @doc """
+  Asserts the LiveView process will receive a message within `timeout`.
+  The default `timeout` is [ExUnit](https://hexdocs.pm/ex_unit/ExUnit.html#configure/1)'s
+  `assert_receive_timeout` (100 ms).
+
+  This requires Erlang/OTP 27 or later.
+
+  A zero-arity setup function may be passed as the third argument. It is
+  invoked after tracing is enabled and before waiting for the message. To
+  customize the timeout, pass it as the fourth argument.
+
+  ## Examples
+
+      assert_will_receive view, {:test_message, num}
+      assert num == 1
+
+      assert_will_receive view, {:test_message, num}, fn ->
+        send(view.pid, {:test_message, 1})
+      end
+
+      assert_will_receive view, {:test_message, num}, fn ->
+        send(view.pid, {:test_message, 1})
+      end, 1000
+  """
+  defmacro assert_will_receive(
+             view,
+             pattern,
+             fun \\ quote(do: fn -> nil end),
+             timeout \\ Application.fetch_env!(:ex_unit, :assert_receive_timeout)
+           ) do
+    {pattern, guard} = extract_guard(pattern)
+    ref = Macro.unique_var(:ref, __MODULE__)
+    tracer = Macro.unique_var(:tracer, __MODULE__)
+    received = Macro.unique_var(:received, __MODULE__)
+    message = Macro.unique_var(:message, __MODULE__)
+    generated_pattern = mark_vars_as_generated(pattern)
+    generated_guard = guard && mark_vars_as_generated(guard)
+
+    trace_pattern =
+      quote do
+        {^unquote(ref), unquote(generated_pattern)}
+      end
+      |> apply_guard(generated_guard)
+
+    quote do
+      {unquote(ref), unquote(tracer)} =
+        Phoenix.LiveViewTest.__start_assert_will_receive__(unquote(view))
+
+      unquote(received) =
+        try do
+          unquote(fun).()
+          assert_receive unquote(trace_pattern), unquote(timeout)
+        after
+          Phoenix.LiveViewTest.__stop_assert_will_receive__(unquote(ref), unquote(tracer))
+        end
+
+      {^unquote(ref), unquote(message)} = unquote(received)
+      unquote(pattern) = unquote(message)
+      unquote(message)
+    end
+  end
+
+  defp extract_guard({:when, _, [pattern, guard]}), do: {pattern, guard}
+  defp extract_guard(pattern), do: {pattern, nil}
+
+  defp apply_guard(pattern, nil), do: pattern
+  defp apply_guard(pattern, guard), do: {:when, [], [pattern, guard]}
+
+  defp mark_vars_as_generated(pattern) do
+    Macro.prewalk(pattern, fn
+      {name, meta, context} when is_atom(name) and is_atom(context) ->
+        {name, [generated: true] ++ meta, context}
+
+      other ->
+        other
+    end)
+  end
+
+  if Code.ensure_loaded?(:trace) and function_exported?(:trace, :session_create, 3) do
+    @doc false
+    def __start_assert_will_receive__(%View{pid: pid}) do
+      ref = make_ref()
+      parent = self()
+      tracer = spawn(fn -> assert_will_receive_loop(parent, ref) end)
+      session = :trace.session_create(:phoenix_live_view_test, tracer, [])
+
+      try do
+        :trace.process(session, pid, true, [:receive])
+        {ref, {tracer, session}}
+      rescue
+        exception ->
+          __stop_assert_will_receive__(ref, {tracer, session})
+          reraise exception, __STACKTRACE__
+      catch
+        kind, reason ->
+          __stop_assert_will_receive__(ref, {tracer, session})
+          :erlang.raise(kind, reason, __STACKTRACE__)
+      end
+    end
+
+    @doc false
+    def __stop_assert_will_receive__(ref, {tracer, session}) do
+      :trace.session_destroy(session)
+
+      monitor_ref = Process.monitor(tracer)
+      send(tracer, {ref, :stop})
+
+      receive do
+        {:DOWN, ^monitor_ref, :process, ^tracer, _reason} ->
+          :ok
+      after
+        5_000 ->
+          Process.exit(tracer, :kill)
+
+          receive do
+            {:DOWN, ^monitor_ref, :process, ^tracer, _reason} -> :ok
+          end
+      end
+
+      flush_assert_will_receive(ref)
+    end
+
+    defp assert_will_receive_loop(parent, ref) do
+      receive do
+        {:trace, _pid, :receive, message} ->
+          send(parent, {ref, message})
+          assert_will_receive_loop(parent, ref)
+
+        {^ref, :stop} ->
+          :ok
+      end
+    end
+
+    defp flush_assert_will_receive(ref) do
+      receive do
+        {^ref, _message} -> flush_assert_will_receive(ref)
+      after
+        0 -> :ok
+      end
+    end
+  else
+    @doc false
+    def __start_assert_will_receive__(%View{}) do
+      raise "assert_will_receive requires Erlang/OTP 27 or later"
+    end
+
+    @doc false
+    def __stop_assert_will_receive__(_ref, _tracer), do: :ok
+  end
+
   @doc """
   Follows the redirect from a `render_*` action or an `{:error, redirect}`
   tuple.

test/phoenix_live_view/integrations/assert_will_receive_test.exs

@@ -0,0 +1,92 @@
+if Code.ensure_loaded?(:trace) and function_exported?(:trace, :session_create, 3) do
+  defmodule Phoenix.LiveView.AssertWillReceiveTest do
+    use ExUnit.Case, async: true
+
+    import Phoenix.LiveViewTest
+    alias Phoenix.LiveViewTest.Support.Endpoint
+
+    @endpoint Endpoint
+
+    setup do
+      {:ok, conn: Plug.Test.init_test_session(Phoenix.ConnTest.build_conn(), %{})}
+    end
+
+    defmodule AssertWillReceiveLive do
+      use Phoenix.LiveView
+
+      def mount(_params, _session, socket) do
+        {:ok, socket}
+      end
+
+      def handle_info({:test_message, _num}, socket) do
+        {:noreply, socket}
+      end
+
+      def render(assigns) do
+        ~H""
+      end
+    end
+
+    describe "assert_will_receive" do
+      test "asserts the LiveView process will receive a message", %{conn: conn} do
+        {:ok, view, _} = live_isolated(conn, AssertWillReceiveLive)
+
+        Task.start(fn ->
+          Process.sleep(25)
+          send(view.pid, {:test_message, 1})
+        end)
+
+        assert_will_receive(view, {:test_message, num})
+        assert num == 1
+      end
+
+      test "runs setup function after tracing is set up", %{conn: conn} do
+        {:ok, view, _} = live_isolated(conn, AssertWillReceiveLive)
+
+        assert_will_receive(view, {:test_message, num}, fn ->
+          send(view.pid, {:test_message, 1})
+        end)
+
+        assert num == 1
+      end
+
+      test "supports guards", %{conn: conn} do
+        {:ok, view, _} = live_isolated(conn, AssertWillReceiveLive)
+
+        assert_will_receive(
+          view,
+          {:test_message, num} when num > 1,
+          fn ->
+            send(view.pid, {:test_message, 1})
+            send(view.pid, {:test_message, 2})
+          end,
+          1000
+        )
+
+        assert num == 2
+      end
+
+      test "stops tracing after the assertion", %{conn: conn} do
+        {:ok, view, _} = live_isolated(conn, AssertWillReceiveLive)
+
+        assert_will_receive(
+          view,
+          {:test_message, :during},
+          fn ->
+            send(view.pid, {:test_message, :during})
+          end,
+          1000
+        )
+
+        send(view.pid, {:test_message, :after})
+
+        receive do
+          {ref, {:test_message, :after}} when is_reference(ref) ->
+            flunk("expected assert_will_receive to stop tracing")
+        after
+          50 -> :ok
+        end
+      end
+    end
+  end
+end