Expand imports within describe, closes #11726

Author: josevalim

lib/ex_unit/lib/ex_unit/callbacks.ex

@@ -167,6 +167,7 @@ defmodule ExUnit.Callbacks do
       @ex_unit_describe nil
       @ex_unit_setup []
       @ex_unit_setup_all []
+      @ex_unit_used_describes %{}
 
       @before_compile unquote(__MODULE__)
       import unquote(__MODULE__)
@@ -175,7 +176,8 @@ defmodule ExUnit.Callbacks do
 
   @doc false
   defmacro __before_compile__(env) do
-    [compile_callbacks(env, :setup), compile_callbacks(env, :setup_all)]
+    used_describes = Module.get_attribute(env.module, :ex_unit_used_describes)
+    [compile_setup(env, :setup, used_describes), compile_setup(env, :setup_all, %{})]
   end
 
   @doc """
@@ -208,8 +210,7 @@ defmodule ExUnit.Callbacks do
       do_setup(quote(do: _), block)
     else
       quote do
-        @ex_unit_setup ExUnit.Callbacks.__callback__(unquote(block), @ex_unit_describe) ++
-                         @ex_unit_setup
+        ExUnit.Callbacks.__setup__(__MODULE__, unquote(block))
       end
     end
   end
@@ -236,12 +237,31 @@ defmodule ExUnit.Callbacks do
 
   defp do_setup(context, block) do
     quote bind_quoted: [context: escape(context), block: escape(block)] do
-      name = :"__ex_unit_setup_#{length(@ex_unit_setup)}"
+      name = ExUnit.Callbacks.__setup__(__MODULE__)
       defp unquote(name)(unquote(context)), unquote(block)
-      @ex_unit_setup [{name, @ex_unit_describe} | @ex_unit_setup]
     end
   end
 
+  @doc false
+  def __setup__(module, callbacks) do
+    setup = Module.get_attribute(module, :ex_unit_setup)
+    Module.put_attribute(module, :ex_unit_setup, Enum.reverse(callbacks(callbacks), setup))
+  end
+
+  @doc false
+  def __setup__(module) do
+    setup = Module.get_attribute(module, :ex_unit_setup)
+
+    name =
+      case Module.get_attribute(module, :ex_unit_describe) do
+        {_line, _message, counter} -> :"__ex_unit_setup_#{counter}_#{length(setup)}"
+        nil -> :"__ex_unit_setup_#{length(setup)}"
+      end
+
+    Module.put_attribute(module, :ex_unit_setup, [name | setup])
+    name
+  end
+
   @doc """
   Defines a callback to be run before all tests in a case.
 
@@ -283,12 +303,7 @@ defmodule ExUnit.Callbacks do
       do_setup_all(quote(do: _), block)
     else
       quote do
-        @ex_unit_describe &&
-          raise "cannot invoke setup_all/1 inside describe as setup_all/1 " <>
-                  "always applies to all tests in a module"
-
-        @ex_unit_setup_all ExUnit.Callbacks.__callback__(unquote(block), nil) ++
-                             @ex_unit_setup_all
+        ExUnit.Callbacks.__setup_all__(__MODULE__, unquote(block))
       end
     end
   end
@@ -312,10 +327,48 @@ defmodule ExUnit.Callbacks do
 
   defp do_setup_all(context, block) do
     quote bind_quoted: [context: escape(context), block: escape(block)] do
-      @ex_unit_describe && raise "cannot invoke setup_all/2 inside describe"
-      name = :"__ex_unit_setup_all_#{length(@ex_unit_setup_all)}"
+      name = ExUnit.Callbacks.__setup_all__(__MODULE__)
       defp unquote(name)(unquote(context)), unquote(block)
-      @ex_unit_setup_all [{name, nil} | @ex_unit_setup_all]
+    end
+  end
+
+  @doc false
+  def __setup_all__(module, callbacks) do
+    no_describe!(module)
+    setup_all = Module.get_attribute(module, :ex_unit_setup_all)
+
+    Module.put_attribute(
+      module,
+      :ex_unit_setup_all,
+      Enum.reverse(callbacks(callbacks), setup_all)
+    )
+  end
+
+  @doc false
+  def __setup_all__(module) do
+    no_describe!(module)
+    setup_all = Module.get_attribute(module, :ex_unit_setup_all)
+    name = :"__ex_unit_setup_all_#{length(setup_all)}"
+    Module.put_attribute(module, :ex_unit_setup_all, [name | setup_all])
+    name
+  end
+
+  defp no_describe!(module) do
+    if Module.get_attribute(module, :ex_unit_describe) do
+      raise "cannot invoke setup_all/1-2 inside describe as setup_all " <>
+              "always applies to all tests in a module"
+    end
+  end
+
+  defp callbacks(callbacks) do
+    for k <- List.wrap(callbacks) do
+      if not is_atom(k) do
+        raise ArgumentError,
+              "setup/setup_all expect a callback name as an atom or " <>
+                "a list of callback names, got: #{inspect(k)}"
+      end
+
+      k
     end
   end
 
@@ -558,8 +611,8 @@ defmodule ExUnit.Callbacks do
     raise_merge_failed!(mod, original_value)
   end
 
-  defp merge(mod, context, data, original_value) when is_list(data) do
-    merge(mod, context, Map.new(data), original_value)
+  defp merge(mod, context, data, _original_value) when is_list(data) do
+    context_merge(mod, context, Map.new(data))
   end
 
   defp merge(mod, context, data, _original_value) when is_map(data) do
@@ -594,44 +647,112 @@ defmodule ExUnit.Callbacks do
     Macro.escape(contents, unquote: true)
   end
 
-  defp compile_callbacks(env, kind) do
-    callbacks = Module.get_attribute(env.module, :"ex_unit_#{kind}") |> Enum.reverse()
+  @doc false
+  def __describe__(module, line, message, fun) do
+    if Module.get_attribute(module, :ex_unit_describe) do
+      raise "cannot call \"describe\" inside another \"describe\". See the documentation " <>
+              "for ExUnit.Case.describe/2 on named setups and how to handle hierarchies"
+    end
+
+    used_describes = Module.get_attribute(module, :ex_unit_used_describes)
+
+    cond do
+      not is_binary(message) ->
+        raise ArgumentError, "describe name must be a string, got: #{inspect(message)}"
 
-    acc =
-      case callbacks do
-        [] ->
-          quote(do: context)
+      is_map_key(used_describes, message) ->
+        raise ExUnit.DuplicateDescribeError,
+              "describe #{inspect(message)} is already defined in #{inspect(module)}"
 
-        [h | t] ->
-          Enum.reduce(t, compile_merge(h), fn callback_describe, acc ->
+      true ->
+        :ok
+    end
+
+    if Module.get_attribute(module, :describetag) != [] do
+      raise "@describetag must be set inside describe/2 blocks"
+    end
+
+    setup = Module.get_attribute(module, :ex_unit_setup)
+    Module.put_attribute(module, :ex_unit_setup, [])
+    Module.put_attribute(module, :ex_unit_describe, {line, message, map_size(used_describes)})
+
+    try do
+      fun.(message, used_describes)
+    after
+      Module.put_attribute(module, :ex_unit_describe, nil)
+      Module.put_attribute(module, :ex_unit_setup, setup)
+      Module.delete_attribute(module, :describetag)
+
+      for attribute <- Module.get_attribute(module, :ex_unit_registered_describe_attributes) do
+        Module.delete_attribute(module, attribute)
+      end
+    end
+  end
+
+  @doc false
+  def __describe__(module, message, used_describes) do
+    {name, body} =
+      case Module.get_attribute(module, :ex_unit_setup) do
+        [] -> {nil, nil}
+        callbacks -> {:"__ex_unit_describe_#{map_size(used_describes)}", compile_setup(callbacks)}
+      end
+
+    used_describes = Map.put(used_describes, message, name)
+    Module.put_attribute(module, :ex_unit_used_describes, used_describes)
+    {name, body}
+  end
+
+  defp compile_setup(env, kind, describes) do
+    calls =
+      env.module
+      |> Module.get_attribute(:"ex_unit_#{kind}")
+      |> compile_setup()
+
+    describe_clauses =
+      for {describe, callback} <- describes,
+          callback != nil,
+          clause <- quote(do: (unquote(describe) -> unquote(callback)(var!(context)))),
+          do: clause
+
+    body =
+      if describe_clauses == [] do
+        calls
+      else
+        describe_clauses =
+          describe_clauses ++
             quote do
-              context = unquote(acc)
-              unquote(compile_merge(callback_describe))
+              _ -> var!(context)
             end
-          end)
+
+        quote do
+          var!(context) = unquote(calls)
+          case Map.get(var!(context), :describe, nil), do: unquote(describe_clauses)
+        end
       end
 
     quote do
-      def __ex_unit__(unquote(kind), context) do
-        describe = Map.get(context, :describe, nil)
-        unquote(acc)
-      end
+      def __ex_unit__(unquote(kind), var!(context)), do: unquote(body)
     end
   end
 
-  defp compile_merge({callback, nil}) do
-    quote do
-      unquote(__MODULE__).__merge__(__MODULE__, context, unquote(callback)(context))
-    end
+  defp compile_setup([]) do
+    quote(do: var!(context))
   end
 
-  defp compile_merge({callback, {_line, describe}}) do
-    quote do
-      if unquote(describe) == describe do
-        unquote(compile_merge({callback, nil}))
-      else
-        context
+  defp compile_setup(callbacks) do
+    [h | t] = Enum.reverse(callbacks)
+
+    Enum.reduce(t, compile_setup_call(h), fn callback, acc ->
+      quote do
+        var!(context) = unquote(acc)
+        unquote(compile_setup_call(callback))
       end
+    end)
+  end
+
+  defp compile_setup_call(callback) do
+    quote do
+      unquote(__MODULE__).__merge__(__MODULE__, var!(context), unquote(callback)(var!(context)))
     end
   end
 end

lib/ex_unit/lib/ex_unit/case.ex

@@ -290,17 +290,15 @@ defmodule ExUnit.Case do
         :moduletag,
         :ex_unit_registered_test_attributes,
         :ex_unit_registered_describe_attributes,
-        :ex_unit_registered_module_attributes,
-        :ex_unit_used_describes
+        :ex_unit_registered_module_attributes
       ]
 
       Enum.each(attributes, &Module.register_attribute(module, &1, accumulate: true))
 
       attributes = [
         before_compile: ExUnit.Case,
         after_compile: ExUnit.Case,
-        ex_unit_async: false,
-        ex_unit_describe: nil
+        ex_unit_async: false
       ]
 
       Enum.each(attributes, fn {k, v} -> Module.put_attribute(module, k, v) end)
@@ -454,48 +452,23 @@ defmodule ExUnit.Case do
   setup steps involved.
   """
   defmacro describe(message, do: block) do
-    quote do
-      ExUnit.Case.__describe__(__MODULE__, __ENV__.line, unquote(message), fn ->
-        unquote(block)
-      end)
-    end
-  end
-
-  @doc false
-  def __describe__(module, line, message, fun) do
-    if Module.get_attribute(module, :ex_unit_describe) do
-      raise "cannot call \"describe\" inside another \"describe\". See the documentation " <>
-              "for ExUnit.Case.describe/2 on named setups and how to handle hierarchies"
-    end
-
-    cond do
-      not is_binary(message) ->
-        raise ArgumentError, "describe name must be a string, got: #{inspect(message)}"
-
-      message in Module.get_attribute(module, :ex_unit_used_describes) ->
-        raise ExUnit.DuplicateDescribeError,
-              "describe #{inspect(message)} is already defined in #{inspect(module)}"
-
-      true ->
-        :ok
-    end
-
-    if Module.get_attribute(module, :describetag) != [] do
-      raise "@describetag must be set inside describe/2 blocks"
-    end
+    definition =
+      quote unquote: false do
+        defp unquote(name)(var!(context)), do: unquote(body)
+      end
 
-    Module.put_attribute(module, :ex_unit_describe, {line, message})
-    Module.put_attribute(module, :ex_unit_used_describes, message)
+    quote do
+      ExUnit.Callbacks.__describe__(__MODULE__, __ENV__.line, unquote(message), fn
+        message, describes ->
+          res = unquote(block)
 
-    try do
-      fun.()
-    after
-      Module.put_attribute(module, :ex_unit_describe, nil)
-      Module.delete_attribute(module, :describetag)
+          case ExUnit.Callbacks.__describe__(__MODULE__, message, describes) do
+            {nil, nil} -> :ok
+            {name, body} -> unquote(definition)
+          end
 
-      for attribute <- Module.get_attribute(module, :ex_unit_registered_describe_attributes) do
-        Module.delete_attribute(module, attribute)
-      end
+          res
+      end)
     end
   end
 
@@ -554,11 +527,11 @@ defmodule ExUnit.Case do
 
     {name, describe, describe_line, describetag} =
       case Module.get_attribute(mod, :ex_unit_describe) do
-        {line, describe} ->
+        {line, describe, _counter} ->
           description = :"#{test_type} #{describe} #{name}"
           {description, describe, line, Module.get_attribute(mod, :describetag)}
 
-        _ ->
+        nil ->
           {:"#{test_type} #{name}", nil, nil, []}
       end
 

lib/ex_unit/test/ex_unit/describe_test.exs

@@ -27,14 +27,23 @@ defmodule ExUnit.DescribeTest do
       [setup_tag: :from_describe]
     end
 
-    test "from describe has higher precedence", context do
+    test "from describe runs later", context do
       assert context.setup_tag == :from_describe
     end
   end
 
+  describe "setup from import" do
+    import Map
+    setup :to_list
+
+    test "is expanded within describe block", context do
+      assert context.setup_tag == :from_module
+    end
+  end
+
   describe "failures" do
     test "when using setup_all inside describe" do
-      assert_raise RuntimeError, ~r"cannot invoke setup_all/2 inside describe", fn ->
+      assert_raise RuntimeError, ~r"cannot invoke setup_all/1-2 inside describe", fn ->
         defmodule Sample do
           use ExUnit.Case