From 1ed7933038d709faae7d7fa8152e2c07ad2d6e50 Mon Sep 17 00:00:00 2001
From: Filipe Cristovao <filipe.cristovao@klarna.com>
Date: Tue, 12 May 2020 22:13:51 +0200
Subject: [PATCH] Add introspection of which modules are currently mocked

There's already an `meck:expects/[1,2]`, allowing introspection into
what kind of expectations are set.
This allows one to also check which modules is meck handling at any
point.

Introduced a `foldl_mocks` internal function to reuse the same logic
for both `unload/0` and `mocked/0`
---
 src/meck.erl        | 47 ++++++++++++++++++++++++++-------------------
 test/meck_tests.erl | 21 ++++++++++++++++++++
 2 files changed, 48 insertions(+), 20 deletions(-)

diff --git a/src/meck.erl b/src/meck.erl
index 9f4efc2..8eb0365 100644
--- a/src/meck.erl
+++ b/src/meck.erl
@@ -54,6 +54,7 @@
 -export([wait/4]).
 -export([wait/5]).
 -export([wait/6]).
+-export([mocked/0]).
 
 %% Syntactic sugar
 -export([loop/1]).
@@ -460,7 +461,15 @@ history(Mod, OptCallerPid)
 -spec unload() -> Unloaded when
       Unloaded :: [Mod],
       Mod :: atom().
-unload() -> lists:foldl(fun unload_if_mocked/2, [], registered()).
+unload() ->
+    foldl_mocks(fun(Mod, Acc) ->
+        try
+            unload(Mod),
+            [Mod | Acc]
+        catch error:{not_mocked, Mod} ->
+            Acc
+        end
+    end, []).
 
 %% @doc Unload a mocked module or a list of mocked modules.
 %%
@@ -723,6 +732,11 @@ capture(Occur, Mod, Func, OptArgsSpec, ArgNum, OptCallerPid) ->
 capture(Occur, Mod, Func, OptArgsSpec, ArgNum) ->
     meck_history:capture(Occur, '_', Mod, Func, OptArgsSpec, ArgNum).
 
+%% @doc Returns the currently mocked modules.
+-spec mocked() -> list(atom()).
+mocked() ->
+    foldl_mocks(fun(M, Acc) -> [M | Acc] end, []).
+
 %%%============================================================================
 %%% Internal functions
 %%%============================================================================
@@ -732,25 +746,18 @@ wait_for_exit(Mod) ->
     MonitorRef = erlang:monitor(process, meck_util:proc_name(Mod)),
     receive {'DOWN', MonitorRef, _Type, _Object, _Info} -> ok end.
 
--spec unload_if_mocked(Mod::atom() | string(), Unloaded::[atom()]) ->
-        NewUnloaded::[atom()].
-unload_if_mocked(Mod, Unloaded) when is_atom(Mod) ->
-    unload_if_mocked(atom_to_list(Mod), Unloaded);
-unload_if_mocked(ModName, Unloaded) when length(ModName) > 5 ->
-    case lists:split(length(ModName) - 5, ModName) of
-        {Name, "_meck"} ->
-            Mocked = erlang:list_to_existing_atom(Name),
-            try
-                unload(Mocked)
-            catch error:{not_mocked, Mocked} ->
-                    ok
-            end,
-            [Mocked | Unloaded];
-        _Else ->
-            Unloaded
-    end;
-unload_if_mocked(_P, Unloaded) ->
-    Unloaded.
+-spec foldl_mocks(Fun, AccIn) -> AccOut when
+    Fun :: fun((Elem :: module(), AccIn) -> AccOut),
+    AccIn :: term(),
+    AccOut :: term().
+foldl_mocks(Fun, Acc0) when is_function(Fun, 2) ->
+    lists:foldl(fun(Mod, Acc)  ->
+        ModName = atom_to_list(Mod),
+        case lists:split(max(length(ModName) - 5, 0), ModName) of
+            {Name, "_meck"} -> Fun(erlang:list_to_existing_atom(Name), Acc);
+            _Else -> Acc
+        end
+    end, Acc0, erlang:registered()).
 
 -spec check_expect_result(ok | {error, Reason::any()}) -> ok.
 check_expect_result(ok) -> ok;
diff --git a/test/meck_tests.erl b/test/meck_tests.erl
index f17223c..988ded7 100644
--- a/test/meck_tests.erl
+++ b/test/meck_tests.erl
@@ -1547,6 +1547,27 @@ wait_purge_expired_tracker_test() ->
     %% Clean
     meck:unload().
 
+mocked_test() ->
+  %% At start, no modules should be mocked:
+  [] = meck:mocked(),
+
+  %% Mocking a new module adds it to the list of mocked modules:
+  meck:new(mocked_test, [non_strict]),
+  [mocked_test] = meck:mocked(),
+
+  %% Mocking with implicit `meck:new` also adds to the list of mocked modules:
+  meck:expect(meck_test_module, c, 2, ok),
+  [meck_test_module, mocked_test] = lists:sort(meck:mocked()),
+
+  meck:new([foo, bar], [non_strict]),
+  [bar, foo, meck_test_module, mocked_test] = lists:sort(meck:mocked()),
+
+  ok = meck:unload([foo, meck_test_module]),
+  [bar, mocked_test] = lists:sort(meck:mocked()),
+
+  meck:unload(),
+  %% After unload all, no modules should be mocked again
+  [] = meck:mocked().
 
 meck_passthrough_test_() ->
     {foreach, fun setup_passthrough/0, fun teardown/1,
-- 
2.42.0.windows.2