Extend protect/2 to accept a list of behviour modules

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

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/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

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

test/support/additional_behaviour.ex

@@ -0,0 +1,5 @@
+defmodule Hammox.Test.AdditionalBehaviour do
+  @moduledoc false
+
+  @callback additional_foo() :: number()
+end

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