From 1da8bee5cad9816165f57492d49e151e1da4cf2b Mon Sep 17 00:00:00 2001
From: Nelson Vides <nelson.vides@erlang-solutions.com>
Date: Fri, 23 Aug 2019 16:08:50 +0200
Subject: [PATCH] Make wait_until more flexible

---
 test/async_helper.erl | 38 ++++++++++++++++++++------------------
 1 file changed, 20 insertions(+), 18 deletions(-)

diff --git a/test/async_helper.erl b/test/async_helper.erl
index 31cd60a8ba0..43a52bce08b 100644
--- a/test/async_helper.erl
+++ b/test/async_helper.erl
@@ -49,38 +49,40 @@ helper_loop() ->
     end.
 
 % @doc Waits `TimeLeft` for `Fun` to return `Expected Value`, then returns `ExpectedValue`
-% If no value is returned or the result doesn't match  `ExpetedValue` error is raised
+% If no value is returned or the result doesn't match  `ExpectedValue` error is raised
 
 wait_until(Fun, ExpectedValue) ->
     wait_until(Fun, ExpectedValue, #{}).
 
 %% Example: wait_until(fun () -> ... end, SomeVal, #{time_left => timer:seconds(2)})
 wait_until(Fun, ExpectedValue, Opts) ->
-    Defaults = #{time_left => timer:seconds(5),
+    DefValidator = fun(NewValue) -> ExpectedValue =:= NewValue end,
+    Defaults = #{validator => DefValidator,
+                 time_left => timer:seconds(5),
                  sleep_time => 100,
                  history => []},
-    do_wait_until(Fun, ExpectedValue, maps:merge(Defaults, Opts)).
+    do_wait_until(Fun, maps:merge(Defaults, Opts)).
 
-do_wait_until(_Fun, _ExpectedValue, #{
-                                      time_left := TimeLeft,
-                                      history := History
-                                     }) when TimeLeft =< 0 ->
+do_wait_until(_Fun, #{
+                time_left := TimeLeft,
+                history := History
+               }) when TimeLeft =< 0 ->
     error({badmatch, lists:reverse(History)});
 
-do_wait_until(Fun, ExpectedValue, Opts) ->
+do_wait_until(Fun, #{validator := Validator} = Opts) ->
     try Fun() of
-        ExpectedValue ->
-            {ok, ExpectedValue};
-        OtherValue ->
-            wait_and_continue(Fun, ExpectedValue, OtherValue, Opts)
+        Value -> case Validator(Value) of
+                   true -> {ok, Value};
+                   _ -> wait_and_continue(Fun, Value, Opts)
+               end
     catch Error:Reason ->
-            wait_and_continue(Fun, ExpectedValue, {Error, Reason}, Opts)
+            wait_and_continue(Fun, {Error, Reason}, Opts)
     end.
 
-wait_and_continue(Fun, ExpectedValue, FunResult, #{time_left := TimeLeft,
-                                                   sleep_time := SleepTime,
-                                                   history := History} = Opts) ->
+wait_and_continue(Fun, FunResult, #{time_left := TimeLeft,
+                                    sleep_time := SleepTime,
+                                    history := History} = Opts) ->
     timer:sleep(SleepTime),
-    do_wait_until(Fun, ExpectedValue, Opts#{time_left => TimeLeft - SleepTime,
-                                            history => [FunResult | History]}).
+    do_wait_until(Fun, Opts#{time_left => TimeLeft - SleepTime,
+                             history => [FunResult | History]}).