LostKobrakai · August 9, 2018 19:34 · Aug 9, 2018
diff --git a/phoenix_instrument.ex b/phoenix_instrument.ex
@@ -0,0 +1,172 @@
+defmodule ConnectWeb.PhoenixInstrument do
+  @moduledoc """
+  Phoenix instrumentation to assert on the query count of controllers/views in
+  Phoenix.ConnCase tests.
+
+  ## Installing
+
+  Before using the module there are some things to setup:
+
+  * Add `{:ok, _} = ConnectWeb.PhoenixInstrument.start_link()` to `test/test_helper.exs`.
+  * Edit your `config/test.exs` to add the module as instrumentation for phoenix and ecto as follows:
+
+  ### Setup phoenix instrumentation
+
+      config :my_app, MyAppWeb.Endpoint,
+        instrumenters: [ConnectWeb.PhoenixInstrument]
+
+  ### Setup ecto instrumentation
+
+      config :my_app, MyApp.Repo,
+        […],
+        loggers: [{Ecto.LogEntry, :log, []}, {ConnectWeb.PhoenixInstrument, :ecto_log, []}]
+
+  ## Usage
+
+      use MyAppWeb.ConnCase
+      import ConnectWeb.PhoenixInstrument
+
+      test "some test", %{conn: conn} do
+        conn = get conn, Routes.page_path(conn, :index)
+        assert_query_limit (controller: 10, view: 0)
+      end
+  """
+  alias ConnectWeb.PhoenixInstrument.Registry, as: QueryCounter
+
+  @doc """
+  Does start the process needed for aggregate ecto queries.
+
+  ## Example
+
+      {:ok, pid} = start_link()
+
+  """
+  def start_link(_opts \\ []) do
+    Registry.start_link(keys: :unique, name: QueryCounter)
+  end
+
+  @doc """
+  Assert on the query count on previous controller hits of the current process.
+
+  Does list the queries and if they originated in the controller or view in the
+  limit is exhausted.
+
+  ## Example
+
+      conn = get conn, Routes.page_path(conn, :index)
+      assert_query_limit (limit: 5)
+
+      conn = get conn, Routes.page_path(conn, :index)
+      assert_query_limit (controller: 10, view: 0)
+
+  """
+  defmacro assert_query_limit(opts \\ [limit: 5])
+
+  defmacro assert_query_limit(limit: limit) when limit >= 0 do
+    quote do
+      {controller_queries, view_queries} = unquote(__MODULE__).receive_queries()
+      queries = controller_queries ++ view_queries
+      unquote(__MODULE__).list_length_pattern(unquote(limit), queries)
+    end
+  end
+
+  defmacro assert_query_limit(controller: controller_limit, view: view_limit)
+           when controller_limit >= 0 and view_limit >= 0 do
+    quote do
+      {controller_queries, view_queries} = unquote(__MODULE__).receive_queries()
+      unquote(__MODULE__).list_length_pattern(unquote(controller_limit), controller_queries)
+      unquote(__MODULE__).list_length_pattern(unquote(view_limit), view_queries)
+    end
+  end
+
+  defmacro assert_query_limit(opts) do
+    quote do
+      assert_query_limit(
+        callback,
+        controller: Keyword.fetch!(unquote(opts), :controller),
+        view: Keyword.fetch!(unquote(opts), :view)
+      )
+    end
+  end
+
+  @doc false
+  defmacro list_length_pattern(limit, list) do
+    pattern = if limit == 0, do: [], else: for(_ <- 1..limit, do: quote(do: _))
+
+    quote do
+      list = unquote(list)
+
+      unless length(list) <= unquote(limit) do
+        assert unquote(pattern) = list
+      end
+    end
+  end
+
+  @doc false
+  def receive_queries() do
+    controller_queries =
+      receive do
+        {:instrumented_queries_controller, queries} -> queries
+      after
+        0 -> []
+      end
+      |> Enum.map(&{:controller, &1})
+
+    view_queries =
+      receive do
+        {:instrumented_queries_view, queries} -> queries
+      after
+        0 -> []
+      end
+      |> Enum.map(&{:view, &1})
+
+    {controller_queries, view_queries}
+  end
+
+  ##############################################################################
+  # Instrumentation callbacks
+  #
+
+  @doc false
+  def ecto_log(log) do
+    with :error <- update_registry({log.caller_pid, :view}, log.query) do
+      update_registry({log.caller_pid, :controller}, log.query)
+    end
+
+    log
+  end
+
+  defp update_registry(key, query) do
+    Registry.update_value(QueryCounter, key, &[query | &1])
+  end
+
+  @doc false
+  def phoenix_controller_call(:start, _, _) do
+    {:ok, _} = Registry.register(QueryCounter, {self(), :controller}, [])
+    :ok
+  end
+
+  @doc false
+  def phoenix_controller_call(:stop, _, _) do
+    pid = self()
+    [{^pid, queries}] = Registry.lookup(QueryCounter, {self(), :controller})
+    send(self(), {:instrumented_queries_controller, queries})
+    :ok = Registry.unregister(QueryCounter, {self(), :controller})
+    :ok
+  end
+
+  @doc false
+  def phoenix_controller_render(:start, _, _) do
+    {:ok, _} = Registry.register(QueryCounter, {self(), :view}, [])
+    :ok
+  end
+
+  @doc false
+  def phoenix_controller_render(:stop, _, _) do
+    pid = self()
+    [{^pid, queries}] = Registry.lookup(QueryCounter, {self(), :view})
+    send(self(), {:instrumented_queries_view, queries})
+    :ok = Registry.unregister(QueryCounter, {self(), :view})
+    :ok
+  end
+end