Repository: msz/hammox
Pull request #55 (2020-10-08 16:34:43)

Decorate all functions with protect/2
Relates to https://github.com/msz/hammox/issues/8

I'm not sure what to write in the doc of `protect/2` 😞 
@msz could you help me with that?


Patch:
diff --git a/lib/hammox.ex b/lib/hammox.ex
index 1e11ff1..c5cd51c 100644
--- a/lib/hammox.ex
+++ b/lib/hammox.ex
@@ -112,7 +112,7 @@ defmodule Hammox do
   """
   @spec protect(module_name :: module()) :: map()
   def protect(module_name) when is_atom(module_name) do
-    funs = module_name.__info__(:functions)
+    funs = get_funcs!(module_name)
     protect(module_name, module_name, funs)
   end
 
@@ -152,7 +152,7 @@ defmodule Hammox do
   def protect(mfa, behaviour_name)
 
   @spec protect(module_name :: module(), funs :: [{atom(), arity() | [arity()]}]) :: map()
-  def protect(module_name, funs) when is_atom(module_name and is_list(funs)),
+  def protect(module_name, funs) when is_atom(module_name) and is_list(funs),
     do: protect(module_name, module_name, funs)
 
   @spec protect(mfa :: mfa(), behaviour_name :: module()) :: fun()
@@ -169,6 +169,13 @@ defmodule Hammox do
     protected(code, typespecs, arity)
   end
 
+  @spec protect(implementation_name :: module(), behaviour_name :: module()) :: map()
+  def protect(implementation_name, behaviour_name)
+      when is_atom(implementation_name) and is_atom(behaviour_name) do
+    funs = get_funcs!(behaviour_name)
+    protect(implementation_name, behaviour_name, funs)
+  end
+
   @doc since: "0.1.0"
   @doc """
   Same as `protect/2`, but allows decorating multiple functions at the same
@@ -394,16 +401,7 @@ defmodule Hammox do
   end
 
   defp do_fetch_typespecs(behaviour_module_name, function_name, arity) do
-    callbacks =
-      case Cache.get({:callbacks, behaviour_module_name}) do
-        nil ->
-          {:ok, callbacks} = Code.Typespec.fetch_callbacks(behaviour_module_name)
-          Cache.put({:callbacks, behaviour_module_name}, callbacks)
-          callbacks
-
-        callbacks ->
-          callbacks
-      end
+    callbacks = fetch_callbacks(behaviour_module_name)
 
     callbacks
     |> Enum.find_value([], fn
@@ -415,6 +413,18 @@ defmodule Hammox do
     end)
   end
 
+  defp fetch_callbacks(behaviour_module_name) do
+    case Cache.get({:callbacks, behaviour_module_name}) do
+      nil ->
+        {:ok, callbacks} = Code.Typespec.fetch_callbacks(behaviour_module_name)
+        Cache.put({:callbacks, behaviour_module_name}, callbacks)
+        callbacks
+
+      callbacks ->
+        callbacks
+    end
+  end
+
   defp fetch_typespecs_for_mock(mock_name, function_name, arity)
        when is_atom(mock_name) and is_atom(function_name) and is_integer(arity) do
     mock_name.__mock_for__()
@@ -455,4 +465,14 @@ defmodule Hammox do
         )
     end
   end
+
+  defp get_funcs!(module_name) do
+    module_exist?(module_name)
+
+    module_name
+    |> fetch_callbacks()
+    |> Enum.map(fn {callback, _typespecs} ->
+      callback
+    end)
+  end
 end


Test patch:
diff --git a/test/hammox_test.exs b/test/hammox_test.exs
index 1da8488..eaefc39 100644
--- a/test/hammox_test.exs
+++ b/test/hammox_test.exs
@@ -61,6 +61,11 @@ defmodule HammoxTest do
                  other_foo: 1
                )
     end
+
+    test "decorate all functions" do
+      assert %{foo_0: _, other_foo_0: _, other_foo_1: _} =
+               Hammox.protect(Hammox.Test.SmallImplementation, Hammox.Test.SmallBehaviour)
+    end
   end
 
   describe "protect/3" do