Merge pull request #131 from rlipscombe/rl-hide-on-load

Allow hiding 'on_load' attribute.
diff --git a/src/meck_code.erl b/src/meck_code.erl
index 895ded0..806a721 100644
--- a/src/meck_code.erl
+++ b/src/meck_code.erl
@@ -28,6 +28,7 @@
 -export([compile_and_load_forms/1]).
 -export([compile_and_load_forms/2]).
 -export([compile_options/1]).
+-export([enable_on_load/2]).
 -export([rename_module/2]).
 
 %% Types
@@ -89,6 +90,14 @@
 compile_options(Module) ->
   filter_options(proplists:get_value(options, Module:module_info(compile))).
 
+enable_on_load(Forms, false) ->
+    Map = fun({attribute,L,on_load,{F,A}}) -> {attribute,L,export,[{F,A}]};
+             (Other) -> Other
+          end,
+    lists:map(Map, Forms);
+enable_on_load(Forms, _) ->
+    Forms.
+
 -spec rename_module(erlang_form(), module()) -> erlang_form().
 rename_module([{attribute, Line, module, OldAttribute}|T], NewName) ->
     case OldAttribute of
diff --git a/src/meck_proc.erl b/src/meck_proc.erl
index 523c6f9..58c5603 100644
--- a/src/meck_proc.erl
+++ b/src/meck_proc.erl
@@ -200,7 +200,8 @@
         _    -> false
     end,
     NoPassCover = proplists:get_bool(no_passthrough_cover, Options),
-    Original = backup_original(Mod, NoPassCover),
+    EnableOnLoad = proplists:get_bool(enable_on_load, Options),
+    Original = backup_original(Mod, NoPassCover, EnableOnLoad),
     NoHistory = proplists:get_bool(no_history, Options),
     History = if NoHistory -> undefined; true -> [] end,
     CanExpect = resolve_can_expect(Mod, Exports, Options),
@@ -338,16 +339,17 @@
         false -> normal
     end.
 
--spec backup_original(Mod::atom(), NoPassCover::boolean()) ->
+-spec backup_original(Mod::atom(), NoPassCover::boolean(), EnableOnLoad::boolean()) ->
     {Cover:: false |
              {File::string(), Data::string(), CompiledOptions::[any()]},
      Binary:: no_binary |
               no_passthrough_cover |
               binary()}.
-backup_original(Mod, NoPassCover) ->
+backup_original(Mod, NoPassCover, EnableOnLoad) ->
     Cover = get_cover_state(Mod),
     try
-        Forms = meck_code:abstract_code(meck_code:beam_file(Mod)),
+        Forms0 = meck_code:abstract_code(meck_code:beam_file(Mod)),
+        Forms = meck_code:enable_on_load(Forms0, EnableOnLoad),
         NewName = meck_util:original_name(Mod),
         CompileOpts = meck_code:compile_options(meck_code:beam_file(Mod)),
         Renamed = meck_code:rename_module(Forms, NewName),
diff --git a/test/meck_on_load_module.erl b/test/meck_on_load_module.erl
new file mode 100644
index 0000000..3519dde
--- /dev/null
+++ b/test/meck_on_load_module.erl
@@ -0,0 +1,10 @@
+-module(meck_on_load_module).
+-on_load(on_load/0).
+-export([ping/0]).
+
+on_load() ->
+    % Assumes that there's an on_load_listener.
+    catch (on_load_listener ! on_load_called),
+    ok.
+
+ping() -> pong.
diff --git a/test/meck_on_load_tests.erl b/test/meck_on_load_tests.erl
new file mode 100644
index 0000000..0e52411
--- /dev/null
+++ b/test/meck_on_load_tests.erl
@@ -0,0 +1,40 @@
+-module(meck_on_load_tests).
+
+-include_lib("eunit/include/eunit.hrl").
+
+on_load_test_() ->
+    {foreach, fun setup/0, fun teardown/1,
+     [fun no_enable_on_load/0,
+      fun enable_on_load/0]}.
+
+setup() ->
+    ok.
+
+teardown(_) ->
+    meck:unload().
+
+no_enable_on_load() ->
+    % We _don't_ want on_load to be called. Listen out for it.
+    register(on_load_listener, self()),
+    meck:new(meck_on_load_module, [passthrough]),
+    ?assertEqual(pong, meck_on_load_module:ping()),
+    receive
+        on_load_called -> erlang:error(unexpected_call_to_on_load)
+    after 100 ->
+              % Use a relatively short timeout, because the happy path goes
+              % through here.
+              ok
+    end.
+
+enable_on_load() ->
+    % We _do_ want on_load to be called.
+    register(on_load_listener, self()),
+    meck:new(meck_on_load_module, [passthrough, enable_on_load]),
+    ?assertEqual(pong, meck_on_load_module:ping()),
+    receive
+        on_load_called -> ok
+    after 200 ->
+              % Use a longer timeout, because testing for not-called is harder,
+              % and this is the sad path.
+              erlang:error(expected_call_to_on_load)
+    end.