Repository: msz/hammox
Pull request #93 (2021-11-08 15:53:00)

Extend `protect/2` to accept a list of behviour modules
My team loves Hammox since it was recommended to us by Jose - and we really
appreciate all of the work you have put into it!

We often write modules that implement multiple behaviours and have gotten into
the habit of doing something similar to this in many of our tests

```
  setup_all do
    Map.merge(
      Hammox.protect(My.Module, My.Module.Contract),
      Hammox.protect(My.Module, My.Module.OtherContract)
    )
  end
```

This PR extends the api to `Hammox.protect` to be able to take in an
implementation module and a list of behaviours so that we could replace the
above with

```
  setup_all do
    Hammox.protect(My.Module, [My.Module.Contract, My.Module.OtherContract]),
  end
```

Please let me know if you have ideas of alternative implementations for this,
or if this does not fit with your design, feel free to disregard!

Once again - thanks for all of your work on this library - it has been very
valuable to our application!



Patch:
diff --git a/lib/hammox.ex b/lib/hammox.ex
index 6751a65..75b81d0 100644
--- a/lib/hammox.ex
+++ b/lib/hammox.ex
@@ -108,8 +108,16 @@ defmodule Hammox do
   def protect(mfa, behaviour_name)
 
   @spec protect(module :: module(), funs :: [function_arity_pair()]) :: %{atom() => fun()}
-  def protect(module, funs) when is_atom(module) and is_list(funs),
-    do: protect(module, module, funs)
+  def protect(module, [{function, arity} | _] = funs)
+      when is_atom(module) and is_atom(function) and (is_integer(arity) or is_list(arity)),
+      do: protect(module, module, funs)
+
+  def protect(module, [behaviour | _] = behaviour_names)
+      when is_atom(module) and is_atom(behaviour) do
+    Enum.reduce(behaviour_names, %{}, fn behaviour_name, acc ->
+      Map.merge(acc, protect(module, behaviour_name))
+    end)
+  end
 
   @spec protect(mfa :: mfa(), behaviour_name :: module()) :: fun()
   def protect({module, function_name, arity} = mfa, behaviour_name)
@@ -198,6 +206,38 @@ defmodule Hammox do
   } = Hammox.protect(TestCalculator, Calculator, add: [2, 3], multiply: 2)
   ```
 
+  ## Batch usage for multiple behviours
+
+  You can decorate all functions defined by any number of behaviours by passing an
+  implementation module and a list of behaviour modules.
+
+  The returned map is useful as the return value for a test setup callback to
+  set test context for all tests to use.
+
+  Example:
+  ```elixir
+  defmodule Calculator do
+    @callback add(integer(), integer()) :: integer()
+    @callback multiply(integer(), integer()) :: integer()
+  end
+
+  defmodule AdditionalCalculator do
+    @callback subtract(integer(), integer()) :: integer()
+  end
+
+  defmodule TestCalculator do
+    def add(a, b), do: a + b
+    def multiply(a, b), do: a * b
+    def subtract(a, b), do: a - b
+  end
+
+  %{
+    add_2: add_2,
+    multiply_2: multiply_2
+    subtract_2: subtract_2
+  } = Hammox.protect(TestCalculator, [Calculator, AdditionalCalculator])
+  ```
+
   ## Behaviour-implementation shortcuts
 
   Often, there exists one "default" implementation for a behaviour. A common
diff --git a/lib/hammox/protect.ex b/lib/hammox/protect.ex
index 027605f..f1b48c0 100644
--- a/lib/hammox/protect.ex
+++ b/lib/hammox/protect.ex
@@ -22,18 +22,34 @@ defmodule Hammox.Protect do
   `Hammox.protect/3` in batch usage.
   - `:funs` — An optional explicit list of functions you'd like to protect.
   Equivalent to the third parameter of `Hammox.protect/3` in batch usage.
+
+  Additionally multiple `behaviour` and `funs` options can be provided for
+  modules that implement multiple behaviours
+  - note: the `funs` options are optional but specific to the `behaviour` that
+    precedes them
+
+  ```
+  use Hammox.Protect,
+    module: Hammox.Test.MultiBehaviourImplementation,
+    behaviour: Hammox.Test.SmallBehaviour,
+    # the `funs` opt below effects the funs protected from `SmallBehaviour`
+    funs: [foo: 0, other_foo: 1],
+    behaviour: Hammox.Test.AdditionalBehaviour
+    # with no `funs` pt provided after `AdditionalBehaviour`, all callbacks
+    # will be protected
+  ````
   """
   alias Hammox.Utils
 
   defmacro __using__(opts) do
     opts_block =
       quote do
-        {module, behaviour, funs} = Hammox.Protect.extract_opts!(unquote(opts))
+        mod_behaviour_funs = Hammox.Protect.extract_opts!(unquote(opts))
       end
 
     funs_block =
       quote unquote: false do
-        for {name, arity} <- funs do
+        for {module, behaviour, funs} <- mod_behaviour_funs, {name, arity} <- funs do
           def unquote(name)(
                 unquote_splicing(
                   Enum.map(
@@ -76,7 +92,6 @@ defmodule Hammox.Protect do
   @doc false
   def extract_opts!(opts) do
     module = Keyword.get(opts, :module)
-    behaviour = Keyword.get(opts, :behaviour)
 
     if is_nil(module) do
       raise ArgumentError,
@@ -89,17 +104,45 @@ defmodule Hammox.Protect do
         """
     end
 
-    module_with_callbacks = behaviour || module
+    mods_and_funs =
+      opts
+      |> Keyword.take([:behaviour, :funs])
+      |> case do
+        # just the module in opts
+        [] ->
+          [{module, get_funs!(module)}]
+
+        # module and funs in opts
+        [{:funs, funs}] ->
+          [{module, funs}]
+
+        # module multiple behaviours with or without funs
+        behaviours_and_maybe_funs ->
+          reduce_opts_to_behaviours_and_funs({behaviours_and_maybe_funs, []})
+      end
 
-    funs = Keyword.get_lazy(opts, :funs, fn -> get_funs!(module_with_callbacks) end)
+    mods_and_funs
+    |> Enum.map(fn {module_with_callbacks, funs} ->
+      if funs == [] do
+        raise ArgumentError,
+          message:
+            "The module #{inspect(module_with_callbacks)} does not contain any callbacks. Please use a behaviour with at least one callback."
+      end
 
-    if funs == [] do
-      raise ArgumentError,
-        message:
-          "The module #{inspect(module_with_callbacks)} does not contain any callbacks. Please use a behaviour with at least one callback."
-    end
+      {module, module_with_callbacks, funs}
+    end)
+  end
+
+  defp reduce_opts_to_behaviours_and_funs({[], acc}) do
+    acc
+  end
+
+  defp reduce_opts_to_behaviours_and_funs({[{:behaviour, behaviour}, {:funs, funs} | rest], acc}) do
+    reduce_opts_to_behaviours_and_funs({rest, [{behaviour, funs} | acc]})
+  end
 
-    {module, behaviour, funs}
+  defp reduce_opts_to_behaviours_and_funs({[{:behaviour, behaviour} | rest], acc}) do
+    reduce_opts_to_behaviours_and_funs({rest, [{behaviour, get_funs!(behaviour)} | acc]})
   end
 
   @doc false


Test patch:
diff --git a/test/hammox/protect_test.exs b/test/hammox/protect_test.exs
index 383e4ae..f65bdb1 100644
--- a/test/hammox/protect_test.exs
+++ b/test/hammox/protect_test.exs
@@ -57,4 +57,45 @@ defmodule Hammox.ProtectTest do
   test "using Protect creates protected versions of functions from given behaviour and implementation" do
     assert_raise Hammox.TypeMatchError, fn -> behaviour_wrong_typespec() end
   end
+
+  defmodule MultiProtect do
+    use Hammox.Protect,
+      module: Hammox.Test.MultiBehaviourImplementation,
+      behaviour: Hammox.Test.SmallBehaviour,
+      behaviour: Hammox.Test.AdditionalBehaviour
+  end
+
+  test "using Protect with multiple behaviour opts creates expected functions" do
+    # Hammox.Test.SmallBehaviour
+    assert_raise Hammox.TypeMatchError, fn -> MultiProtect.foo() end
+    assert 1 == MultiProtect.other_foo()
+    assert 1 == MultiProtect.other_foo(10)
+
+    # Hammox.Test.AdditionalBehaviour
+    assert 1 == MultiProtect.additional_foo()
+  end
+
+  defmodule MultiProtectWithFuns do
+    use Hammox.Protect,
+      module: Hammox.Test.MultiBehaviourImplementation,
+      behaviour: Hammox.Test.SmallBehaviour,
+      funs: [other_foo: 1],
+      behaviour: Hammox.Test.AdditionalBehaviour
+  end
+
+  test "using Protect with multiple behaviour / funs opts creates expected functions" do
+    # Hammox.Test.SmallBehaviour
+    assert_raise UndefinedFunctionError,
+                 ~r[MultiProtectWithFuns.foo/0 is undefined or private],
+                 fn -> apply(MultiProtectWithFuns, :foo, []) end
+
+    assert_raise UndefinedFunctionError,
+                 ~r[MultiProtectWithFuns.other_foo/0 is undefined or private],
+                 fn -> apply(MultiProtectWithFuns, :other_foo, []) end
+
+    assert 1 == MultiProtectWithFuns.other_foo(10)
+
+    # Hammox.Test.AdditionalBehaviour
+    assert 1 == MultiProtectWithFuns.additional_foo()
+  end
 end
diff --git a/test/hammox_test.exs b/test/hammox_test.exs
index 03e957c..08c6e74 100644
--- a/test/hammox_test.exs
+++ b/test/hammox_test.exs
@@ -71,6 +71,14 @@ defmodule HammoxTest do
       assert %{foo_0: _, other_foo_0: _, other_foo_1: _} =
                Hammox.protect(Hammox.Test.SmallImplementation, Hammox.Test.SmallBehaviour)
     end
+
+    test "decorate all functions from multiple behaviours" do
+      assert %{foo_0: _, other_foo_0: _, other_foo_1: _, additional_foo_0: _} =
+               Hammox.protect(Hammox.Test.MultiBehaviourImplementation, [
+                 Hammox.Test.SmallBehaviour,
+                 Hammox.Test.AdditionalBehaviour
+               ])
+    end
   end
 
   describe "protect/3" do
diff --git a/test/support/additional_behaviour.ex b/test/support/additional_behaviour.ex
new file mode 100644
index 0000000..b7b480b
--- /dev/null
+++ b/test/support/additional_behaviour.ex
@@ -0,0 +1,5 @@
+defmodule Hammox.Test.AdditionalBehaviour do
+  @moduledoc false
+
+  @callback additional_foo() :: number()
+end
diff --git a/test/support/multi_behaviour_implementation.ex b/test/support/multi_behaviour_implementation.ex
new file mode 100644
index 0000000..084900c
--- /dev/null
+++ b/test/support/multi_behaviour_implementation.ex
@@ -0,0 +1,20 @@
+defmodule Hammox.Test.MultiBehaviourImplementation do
+  @moduledoc false
+
+  @behaviour Hammox.Test.SmallBehaviour
+  @behaviour Hammox.Test.AdditionalBehaviour
+
+  @impl Hammox.Test.SmallBehaviour
+  def foo, do: :bar
+
+  @impl Hammox.Test.SmallBehaviour
+  def other_foo, do: 1
+
+  @impl Hammox.Test.SmallBehaviour
+  def other_foo(_), do: 1
+
+  @impl Hammox.Test.AdditionalBehaviour
+  def additional_foo(), do: 1
+
+  def nospec_fun, do: 1
+end