Add --generator option for eunit

2 files changed, 67 insertions, 12 deletions

diff --git a/src/rebar_prv_eunit.erl b/src/rebar_prv_eunit.erl
index f120926..d8136a2 100644
--- a/src/rebar_prv_eunit.erl
+++ b/src/rebar_prv_eunit.erl
@@ -40,14 +40,18 @@ init(State) ->
 
 -spec do(rebar_state:t()) -> {ok, rebar_state:t()} | {error, string()}.
 do(State) ->
-    Tests = prepare_tests(State),
-    %% inject `eunit_first_files`, `eunit_compile_opts` and any
-    %% directories required by tests into the applications
-    NewState = inject_eunit_state(State, Tests),
-    case compile(NewState) of
-        %% successfully compiled apps
-        {ok, S} -> do(S, Tests);
-        Error   -> Error
+    try
+        Tests = prepare_tests(State),
+        %% inject `eunit_first_files`, `eunit_compile_opts` and any
+        %% directories required by tests into the applications
+        NewState = inject_eunit_state(State, Tests),
+        case compile(NewState) of
+            %% successfully compiled apps
+            {ok, S} -> do(S, Tests);
+            Error   -> Error
+        end
+    catch
+        throw:Reason -> {error, Reason}
     end.
 
 do(State, Tests) ->
@@ -134,17 +138,28 @@ resolve_tests(State) ->
     Files        = resolve(file, RawOpts),
     Modules      = resolve(module, RawOpts),
     Suites       = resolve(suite, module, RawOpts),
-    Apps ++ Applications ++ Dirs ++ Files ++ Modules ++ Suites.
+    Generator    = resolve(generator, RawOpts),
+    Apps ++ Applications ++ Dirs ++ Files ++ Modules ++ Suites ++ Generator.
 
 resolve(Flag, RawOpts) -> resolve(Flag, Flag, RawOpts).
 
 resolve(Flag, EUnitKey, RawOpts) ->
     case proplists:get_value(Flag, RawOpts) of
         undefined -> [];
-        Args      -> lists:map(fun(Arg) -> normalize(EUnitKey, Arg) end,
-                               rebar_string:lexemes(Args, [$,]))
+        Args      ->
+            lists:flatten(lists:map(fun(Arg) -> normalize(EUnitKey, Arg) end,
+                                    rebar_string:lexemes(Args, [$,])))
     end.
 
+normalize(generator, Value) ->
+    case string:tokens(Value, [$:]) of
+        [Module0, Functions] ->
+            Module = list_to_atom(Module0),
+            lists:map(fun(F) -> {generator, Module, list_to_atom(F)} end,
+                      string:tokens(Functions, [$;]));
+        _ ->
+            throw(lists:concat(["Generator `", Value, "' is invalid format."]))
+    end;
 normalize(Key, Value) when Key == dir; Key == file -> {Key, Value};
 normalize(Key, Value) -> {Key, list_to_atom(Value)}.
 
@@ -353,6 +368,8 @@ validate(State, {module, Module}) ->
     validate_module(State, Module);
 validate(State, {suite, Module}) ->
     validate_module(State, Module);
+validate(State, {generator, Module, Function}) ->
+    validate_generator(State, Module, Function);
 validate(State, Module) when is_atom(Module) ->
     validate_module(State, Module);
 validate(State, Path) when is_list(Path) ->
@@ -395,6 +412,9 @@ validate_module(_State, Module) ->
         _            -> ok
     end.
 
+validate_generator(State, Module, _Function) ->
+    validate_module(State, Module).
+
 resolve_eunit_opts(State) ->
     {Opts, _} = rebar_state:command_parsed_args(State),
     EUnitOpts = rebar_state:get(State, eunit_opts, []),
@@ -490,6 +510,7 @@ eunit_opts(_State) ->
      {file, $f, "file", string, help(file)},
      {module, $m, "module", string, help(module)},
      {suite, $s, "suite", string, help(module)},
+     {generator, $g, "generator", string, help(generator)},
      {verbose, $v, "verbose", boolean, help(verbose)},
      {name, undefined, "name", atom, help(name)},
      {sname, undefined, "sname", atom, help(sname)},
@@ -501,6 +522,7 @@ help(cover_export_name) -> "Base name of the coverdata file to write";
 help(dir)       -> "Comma separated list of dirs to load tests from. Equivalent to `[{dir, Dir}]`.";
 help(file)      -> "Comma separated list of files to load tests from. Equivalent to `[{file, File}]`.";
 help(module)    -> "Comma separated list of modules to load tests from. Equivalent to `[{module, Module}]`.";
+help(generator) -> "Comma separated list of generators (the format is `module:function`) to load tests from. Equivalent to `[{generator, Module, Function}]`.";
 help(verbose)   -> "Verbose output. Defaults to false.";
 help(name)      -> "Gives a long name to the node";
 help(sname)     -> "Gives a short name to the node";
diff --git a/test/rebar_eunit_SUITE.erl b/test/rebar_eunit_SUITE.erl
index 1a8bade..87dd3ed 100644
--- a/test/rebar_eunit_SUITE.erl
+++ b/test/rebar_eunit_SUITE.erl
@@ -13,6 +13,7 @@
 -export([single_application_arg/1, multi_application_arg/1, missing_application_arg/1]).
 -export([single_module_arg/1, multi_module_arg/1, missing_module_arg/1]).
 -export([single_suite_arg/1, multi_suite_arg/1, missing_suite_arg/1]).
+-export([single_generator_arg/1, multi_generator_arg/1, missing_generator_arg/1]).
 -export([single_file_arg/1, multi_file_arg/1, missing_file_arg/1]).
 -export([single_dir_arg/1, multi_dir_arg/1, missing_dir_arg/1]).
 -export([multiple_arg_composition/1, multiple_arg_errors/1]).
@@ -47,6 +48,7 @@ groups() ->
                           single_application_arg, multi_application_arg, missing_application_arg,
                           single_module_arg, multi_module_arg, missing_module_arg,
                           single_suite_arg, multi_suite_arg, missing_suite_arg,
+                          single_generator_arg, multi_generator_arg, missing_generator_arg,
                           single_file_arg, multi_file_arg, missing_file_arg,
                           single_dir_arg, multi_dir_arg, missing_dir_arg,
                           multiple_arg_composition, multiple_arg_errors]}].
@@ -239,7 +241,7 @@ multi_app_testset(Config) ->
     Set = {ok, [{application, multi_app_baz},
                 {application, multi_app_bar},
                 {module, multi_app_bar_tests_helper},
-                {module, multi_app_baz_tests_helper}, 
+                {module, multi_app_baz_tests_helper},
                 {module, multi_app_tests},
                 {module, multi_app_tests_helper}]},
     Set = rebar_prv_eunit:prepare_tests(Result).
@@ -404,6 +406,37 @@ missing_suite_arg(Config) ->
     Error = {error, {rebar_prv_eunit, {eunit_test_errors, ["Module `missing_app' not found in project."]}}},
     Error = rebar_prv_eunit:validate_tests(State, rebar_prv_eunit:prepare_tests(State)).
 
+%% check that the --generator cmd line opt generates the correct test set
+single_generator_arg(Config) ->
+    S = ?config(result, Config),
+
+    {ok, Args} = getopt:parse(rebar_prv_eunit:eunit_opts(S), ["--generator=module_name:function_name"]),
+    State = rebar_state:command_parsed_args(S, Args),
+
+    {ok, [{generator, module_name, function_name}]} = rebar_prv_eunit:prepare_tests(State).
+
+multi_generator_arg(Config) ->
+    S = ?config(result, Config),
+
+    {ok, Args} = getopt:parse(rebar_prv_eunit:eunit_opts(S), ["--generator=module1:func1;func2,module2:func1;func2"]),
+    State = rebar_state:command_parsed_args(S, Args),
+
+    Generators = [{generator, module1, func1},
+                  {generator, module1, func2},
+                  {generator, module2, func1},
+                  {generator, module2, func2}],
+    {ok, Generators} = rebar_prv_eunit:prepare_tests(State).
+
+%% check that an invalid --suite cmd line opt generates an error
+missing_generator_arg(Config) ->
+    S = ?config(result, Config),
+
+    {ok, Args} = getopt:parse(rebar_prv_eunit:eunit_opts(S), ["--generator=missing_module:func1"]),
+    State = rebar_state:command_parsed_args(S, Args),
+
+    Error = {error, {rebar_prv_eunit, {eunit_test_errors, ["Module `missing_module' not found in project."]}}},
+    Error = rebar_prv_eunit:validate_tests(State, rebar_prv_eunit:prepare_tests(State)).
+
 %% check that the --file cmd line opt generates the correct test set
 single_file_arg(Config) ->
     S = ?config(result, Config),