Implement wait for a number of calls feature (Fixes #81)

src/meck.erl

@@ -48,6 +48,9 @@
 -export([reset/1]).
 -export([capture/5]).
 -export([capture/6]).
+-export([wait/4]).
+-export([wait/5]).
+-export([wait/6]).
 
 %% Syntactic sugar
 -export([loop/1]).
@@ -435,6 +438,62 @@ num_calls(Mod, OptFun, OptArgsSpec) ->
 num_calls(Mod, OptFun, OptArgsSpec, OptPid) ->
     meck_history:num_calls(OptPid, Mod, OptFun, OptArgsSpec).
 
+%% @doc Blocks until either function `Mod:Func' is called at least once with
+%% arguments matching `OptArgsSpec', or `Timeout' has elapsed. In the latter
+%% case the call fails with `error:timeout'.
+%%
+%% The number of calls is counted starting from the most resent call to
+%% {@link reset/1} on the mock or from the mock creation, whichever occurred
+%% latter. If a matching call has already occurred, then the function returns
+%% `ok' immediately.
+%%
+%% @equiv wait(1, Mod, OptFunc, OptArgsSpec, '_', Timeout)
+-spec wait(Mod, OptFunc, OptArgsSpec, Timeout) -> ok when
+      Mod :: atom(),
+      OptFunc :: '_' | atom(),
+      OptArgsSpec :: '_' | args_spec(),
+      Timeout :: timeout().
+wait(Mod, OptFunc, OptArgsSpec, Timeout) ->
+    wait(1, Mod, OptFunc, OptArgsSpec, '_', Timeout).
+
+%% @doc Blocks until either function `Mod:Func' is called at least `Times' with
+%% arguments matching `OptArgsSpec', or `Timeout' has elapsed. In the latter
+%% case the call fails with `error:timeout'.
+%%
+%% The number of calls is counted starting from the most resent call to
+%% {@link reset/1} on the mock or from the mock creation, whichever occurred
+%% latter. If `Times' number of matching calls has already occurred, then the
+%% function returns `ok' immediately.
+%%
+%% @equiv wait(Times, Mod, OptFunc, OptArgsSpec, '_', Timeout)
+-spec wait(Times, Mod, OptFunc, OptArgsSpec, Timeout) -> ok when
+      Times :: pos_integer(),
+      Mod :: atom(),
+      OptFunc :: '_' | atom(),
+      OptArgsSpec :: '_' | args_spec(),
+      Timeout :: timeout().
+wait(Times, Mod, OptFunc, OptArgsSpec, Timeout) ->
+    wait(Times, Mod, OptFunc, OptArgsSpec, '_', Timeout).
+
+%% @doc Blocks until either function `Mod:Func' is called at least `Times' with
+%% arguments matching `OptArgsSpec' by process `OptCallerPid', or `Timeout' has
+%% elapsed. In the latter case the call fails with `error:timeout'.
+%%
+%% The number of calls is counted starting from the most resent call to
+%% {@link reset/1} on the mock or from the mock creation, whichever occurred
+%% latter. If `Times' number of matching call has already occurred, then the
+%% function returns `ok' immediately.
+-spec wait(Times, Mod, OptFunc, OptArgsSpec, OptCallerPid, Timeout) -> ok when
+      Times :: pos_integer(),
+      Mod :: atom(),
+      OptFunc :: '_' | atom(),
+      OptArgsSpec :: '_' | args_spec(),
+      OptCallerPid :: '_' | pid(),
+      Timeout :: timeout().
+wait(Times, Mod, OptFunc, OptArgsSpec, OptCallerPid, Timeout) ->
+    ArgsMatcher = meck_args_matcher:new(OptArgsSpec),
+    meck_proc:wait(Mod, Times, OptFunc, ArgsMatcher, OptCallerPid, Timeout).
+
 %% @doc Erases the call history for a mocked module or a list of mocked modules.
 %%
 %% This function will erase all calls made heretofore from the history of the

src/meck_history.erl

@@ -29,6 +29,7 @@
 -export([get_history/2]).
 -export([num_calls/4]).
 -export([capture/6]).
+-export([new_filter/3]).
 
 %%%============================================================================
 %%% Types
@@ -93,10 +94,6 @@ capture(Occur, OptCallerPid, Mod, Func, OptArgsSpec, ArgNum) ->
             lists:nth(ArgNum, Args)
     end.
 
-%%%============================================================================
-%%% Internal functions
-%%%============================================================================
-
 -spec new_filter(opt_pid(), opt_func(), meck_args_matcher:args_matcher()) ->
         fun((history_record()) -> boolean()).
 new_filter(TheCallerPid, TheFunc, ArgsMatcher) ->
@@ -112,6 +109,10 @@ new_filter(TheCallerPid, TheFunc, ArgsMatcher) ->
             false
     end.
 
+%%%============================================================================
+%%% Internal functions
+%%%============================================================================
+
 -spec nth_record(Occur::pos_integer(), history()) -> history_record() |
                                                      not_found.
 nth_record(Occur, History) ->

src/meck_proc.erl

@@ -24,6 +24,7 @@
 -export([set_expect/2]).
 -export([delete_expect/3]).
 -export([get_history/1]).
+-export([wait/6]).
 -export([reset/1]).
 -export([validate/1]).
 -export([stop/1]).
@@ -53,7 +54,21 @@
                 original :: term(),
                 was_sticky = false :: boolean(),
                 reload :: {Compiler::pid(), {From::pid(), Tag::any()}} |
-                undefined}).
+                          undefined,
+                tracker :: tracker()}).
+
+-record(tracker, {opt_func :: '_' | atom(),
+                  args_matcher :: meck_args_matcher:args_matcher(),
+                  opt_caller_pid :: '_' | pid(),
+                  countdown :: non_neg_integer(),
+                  timer_ref :: reference(),
+                  reply_to :: {Caller::pid(), Tag::any()}}).
+
+%%%============================================================================
+%%% Types
+%%%============================================================================
+
+-type tracker() :: #tracker{}.
 
 %%%============================================================================
 %%% API
@@ -114,6 +129,30 @@ add_history(Mod, CallerPid, Func, Args, Result) ->
 get_history(Mod) ->
     gen_server(call, Mod, get_history).
 
+-spec wait(Mod::atom(),
+           Times::non_neg_integer(),
+           OptFunc::'_' | atom(),
+           meck_args_matcher:args_matcher(),
+           OptCallerPid::'_' | pid(),
+           timeout()) ->
+        ok.
+wait(Mod, Times, OptFunc, ArgsMatcher, OptCallerPid, Timeout)
+  when erlang:is_integer(Times) andalso Times > 0 andalso
+       erlang:is_integer(Timeout) andalso Timeout >= 0 ->
+    Name = meck_util:proc_name(Mod),
+    try gen_server:call(Name, {wait, Times, OptFunc, ArgsMatcher, OptCallerPid,
+                               Timeout},
+                        infinity)
+    of
+        ok ->
+            ok;
+        {error, timeout} ->
+            erlang:error(timeout)
+    catch
+        exit:_Reason ->
+            erlang:error({not_mocked, Mod})
+    end.
+
 -spec reset(Mod::atom()) -> ok.
 reset(Mod) ->
     gen_server(call, Mod, reset).
@@ -190,6 +229,26 @@ handle_call(get_history, _From, S = #state{history = undefined}) ->
     {reply, [], S};
 handle_call(get_history, _From, S) ->
     {reply, lists:reverse(S#state.history), S};
+handle_call({wait, Times, OptFunc, ArgsMatcher, OptCallerPid, Timeout}, From,
+            S = #state{history = History, tracker = undefined}) ->
+    case times_called(OptFunc, ArgsMatcher, OptCallerPid, History) of
+        CalledSoFar when CalledSoFar >= Times ->
+            {reply, ok, S};
+        _CalledSoFar when Timeout =:= 0 ->
+            {reply, {error, timeout}, S};
+        CalledSoFar ->
+            TimerRef = erlang:start_timer(Timeout, erlang:self(), tracker),
+            Tracker = #tracker{opt_func = OptFunc,
+                               args_matcher = ArgsMatcher,
+                               opt_caller_pid = OptCallerPid,
+                               countdown = Times - CalledSoFar,
+                               timer_ref = TimerRef,
+                               reply_to = From},
+            {noreply, S#state{tracker = Tracker}}
+    end;
+handle_call({wait, _Times, _OptFunc, _ArgsMatcher, _OptCallerPid, _Timeout},
+            _From, S) ->
+    {reply, {error, concurrent_wait}, S};
 handle_call(reset, _From, S) ->
     {reply, ok, S#state{history = []}};
 handle_call(invalidate, _From, S) ->
@@ -200,12 +259,18 @@ handle_call(stop, _From, S) ->
     {stop, normal, ok, S}.
 
 %% @hidden
-handle_cast({add_history, _Item}, S = #state{history = undefined}) ->
-    {noreply, S};
-handle_cast({add_history, Item}, S = #state{reload = Reload}) ->
+handle_cast({add_history, HistoryRecord}, S = #state{history = undefined,
+                                                     tracker = Tracker}) ->
+    UpdTracker = update_tracker(HistoryRecord, Tracker),
+    {noreply, S#state{tracker = UpdTracker}};
+handle_cast({add_history, HistoryRecord}, S = #state{history = History,
+                                                     tracker = Tracker,
+                                                     reload = Reload}) ->
     case Reload of
         undefined ->
-            {noreply, S#state{history = [Item | S#state.history]}};
+            UpdTracker = update_tracker(HistoryRecord, Tracker),
+            {noreply, S#state{history = [HistoryRecord | History],
+                              tracker = UpdTracker}};
         _ ->
             % Skip Item if the mocked module compiler is running.
             {noreply, S}
@@ -222,6 +287,11 @@ handle_info({'EXIT', Pid, _Reason}, S = #state{reload = Reload}) ->
         _ ->
             {noreply, S}
     end;
+handle_info({timeout, TimerRef, tracker},
+            #state{tracker = #tracker{timer_ref = TimerRef,
+                                      reply_to = ReplyTo}} = S) ->
+    gen_server:reply(ReplyTo, {error, timeout}),
+    {noreply, S#state{tracker = undefined}};
 handle_info(_Info, S) ->
     {noreply, S}.
 
@@ -484,3 +554,52 @@ cleanup(Mod) ->
     code:delete(Mod),
     code:purge(meck_util:original_name(Mod)),
     code:delete(meck_util:original_name(Mod)).
+
+-spec times_called(OptFunc::'_' | atom(),
+                   meck_args_matcher:args_matcher(),
+                   OptCallerPid::'_' | pid(),
+                   meck_history:history()) ->
+        non_neg_integer().
+times_called(OptFunc, ArgsMatcher, OptCallerPid, History) ->
+    Filter = meck_history:new_filter(OptCallerPid, OptFunc, ArgsMatcher),
+    lists:foldl(fun(HistoryRec, Acc) ->
+                        case Filter(HistoryRec) of
+                            true ->
+                                Acc + 1;
+                            _Else ->
+                                Acc
+                        end
+                end, 0, History).
+
+-spec update_tracker(meck_history:history_record(), tracker() | undefined) ->
+        UpdTracker::tracker() | undefined.
+update_tracker(_HistoryRecord, undefined) ->
+    undefined;
+update_tracker(HistoryRecord, Tracker) ->
+    CallerPid = erlang:element(1, HistoryRecord),
+    {_Mod, Func, Args} = erlang:element(2, HistoryRecord),
+    update_tracker(Func, Args, CallerPid, Tracker).
+
+-spec update_tracker(Func::atom(), Args::[any()], Caller::pid(), tracker()) ->
+    UpdTracker::tracker() | undefined.
+update_tracker(Func, Args, CallerPid,
+               #tracker{opt_func = OptFunc,
+                        args_matcher = ArgsMatcher,
+                        opt_caller_pid = OptCallerPid,
+                        countdown = Countdown,
+                        timer_ref = TimerRef,
+                        reply_to = ReplyTo} = Tracker)
+  when (OptFunc =:= '_' orelse Func =:= OptFunc) andalso
+       (OptCallerPid =:= '_' orelse CallerPid =:= OptCallerPid) ->
+    case meck_args_matcher:match(Args, ArgsMatcher) of
+        false ->
+            Tracker;
+        true when Countdown == 1 ->
+            erlang:cancel_timer(TimerRef),
+            gen_server:reply(ReplyTo, ok),
+            undefined;
+        true ->
+            Tracker#tracker{countdown = Countdown - 1}
+    end;
+update_tracker(_Func, _Args, _CallerPid, Tracker) ->
+    Tracker.

test/meck_tests.erl

@@ -1241,6 +1241,72 @@ meck_implicit_new_test()->
     ?assertMatch(foo, meck_test_module:c(1, 1)),
     meck:unload().
 
+wait_already_called_test() ->
+    %% Given
+    meck:new(test, [non_strict]),
+    meck:expect(test, foo, 2, ok),
+    %% When
+    test:foo(1, 2),
+    test:foo(1, 2),
+    %% Then
+    ?assertMatch(ok, meck:wait(2, test, foo, [1, '_'], 100)),
+    %% Clean
+    meck:unload().
+
+wait_not_called_zero_timeout_test() ->
+    %% Given
+    meck:new(test, [non_strict]),
+    meck:expect(test, foo, 2, ok),
+    %% When
+    test:foo(1, 2),
+    test:foo(1, 2),
+    %% Then
+    ?assertError(timeout, meck:wait(3, test, foo, [1, '_'], 0)),
+    %% Clean
+    meck:unload().
+
+wait_not_called_another_proc_test() ->
+    %% Given
+    meck:new(test, [non_strict]),
+    meck:expect(test, foo, 2, ok),
+    %% When
+    test:foo(1, 2), % Called, but not by the expected proc.
+    Pid = erlang:spawn(fun() ->
+                              timer:sleep(50),
+                              test:foo(2, 2) % Unexpected first argument
+                       end),
+    %% Then
+    ?assertError(timeout, meck:wait(1, test, foo, [1, '_'], Pid, 0)),
+    %% Clean
+    meck:unload().
+
+wait_called_another_proc_test() ->
+    %% Given
+    meck:new(test, [non_strict]),
+    meck:expect(test, foo, 2, ok),
+    %% When
+    Pid = erlang:spawn(fun() ->
+                              timer:sleep(50),
+                              test:foo(1, 2),
+                              test:foo(2, 2), % Unexpected first argument
+                              test:foo(1, 2)
+                       end),
+    %% Then
+    ?assertMatch(ok, meck:wait(2, test, foo, [1, '_'], Pid, 500)),
+    %% Clean
+    meck:unload().
+
+wait_timeout_test() ->
+    %% Given
+    meck:new(test, [non_strict]),
+    meck:expect(test, foo, 2, ok),
+    %% When
+    test:foo(1, 2),
+    %% Then
+    ?assertError(timeout, meck:wait(2, test, foo, [1, '_'], '_', 10)),
+    %% Clean
+    meck:unload().
+
 %%=============================================================================
 %% Internal Functions
 %%=============================================================================