diff --git a/lib/hammox.ex b/lib/hammox.ex index 852dc6d..d4b66ac 100644 --- a/lib/hammox.ex +++ b/lib/hammox.ex @@ -109,10 +109,25 @@ defmodule Hammox do add_2.(1.5, 2.5) # throws Hammox.TypeMatchError ``` + + If pass module and list instead of MFA and module works like `protect/3` + ```elixir + # calling this + Hammox.protect(SomeModule, foo: 1) + # works same as this + Hammox.protect(SomeModule, SomeModule, foo: 1) + ``` """ - @spec protect(mfa :: mfa(), behaviour_name :: module()) :: fun() + @spec protect( + mfa() | module(), + module() | [{atom(), arity() | [arity()]}] + ) :: + fun() | map() def protect(mfa, behaviour_name) + def protect(module_name, funs) when is_atom(module_name and is_list(funs)), + do: protect(module_name, module_name, funs) + def protect({module_name, function_name, arity}, behaviour_name) when is_atom(module_name) and is_atom(function_name) and is_integer(arity) and is_atom(behaviour_name) do diff --git a/test/hammox_test.exs b/test/hammox_test.exs index dde14bf..de8bed6 100644 --- a/test/hammox_test.exs +++ b/test/hammox_test.exs @@ -19,6 +19,14 @@ defmodule HammoxTest do ) end) end + + test "decorate multiple functions" do + assert %{foo_0: _, other_foo_1: _} = + Hammox.protect(Hammox.Test.SmallBehaviour, + foo: 0, + other_foo: 1 + ) + end end describe "protect/3" do