From 41a248cd88a65c1f3d2e3ffbecd5c4c755cd62d1 Mon Sep 17 00:00:00 2001
From: saneery <anatolyniky@gmail.com>
Date: Fri, 2 Oct 2020 21:12:47 +0300
Subject: [PATCH] Add protect/2 for behaviours and implementations in the same
 module

---
 lib/hammox.ex        | 17 ++++++++++++++++-
 test/hammox_test.exs |  8 ++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

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