asserting on query counts using phoenix+ecto instumentation

phoenix_instrument.ex

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