1 files changed, 52 insertions, 22 deletions

diff --git a/src/rebar_prv_eunit.erl b/src/rebar_prv_eunit.erl
index 8eaa926..3ad593d 100644
--- a/src/rebar_prv_eunit.erl
+++ b/src/rebar_prv_eunit.erl
@@ -37,7 +37,8 @@ init(State) ->
 -spec do(rebar_state:t()) -> {ok, rebar_state:t()} | {error, string()}.
 do(State) ->
     ?INFO("Performing EUnit tests...", []),
-    code:add_pathsa(rebar_state:code_paths(State, all_deps)),
+    rebar_utils:update_code(rebar_state:code_paths(State, all_deps)),
+
     %% Run eunit provider prehooks
     Providers = rebar_state:providers(State),
     Cwd = rebar_dir:get_cwd(),
@@ -49,7 +50,7 @@ do(State) ->
                 {ok, State1} ->
                     %% Run eunit provider posthooks
                     rebar_hooks:run_all_hooks(Cwd, post, ?PROVIDER, Providers, State1),
-                    rebar_utils:cleanup_code_path(rebar_state:code_paths(State1, default)),
+                    rebar_utils:cleanup_code_path(rebar_state:code_paths(State, default)),
                     {ok, State1};
                 Error ->
                     rebar_utils:cleanup_code_path(rebar_state:code_paths(State, default)),
@@ -84,8 +85,7 @@ format_error({error_running_tests, Reason}) ->
 test_state(State) ->
     ErlOpts = rebar_state:get(State, eunit_compile_opts, []),
     TestOpts = safe_define_test_macro(ErlOpts),
-    TestDir = [{extra_src_dirs, ["test"]}],
-    first_files(State) ++ [{erl_opts, TestOpts ++ TestDir}].
+    first_files(State) ++ [{erl_opts, TestOpts}].
 
 safe_define_test_macro(Opts) ->
     %% defining a compile macro twice results in an exception so
@@ -106,38 +106,56 @@ first_files(State) ->
 
 prepare_tests(State) ->
     {RawOpts, _} = rebar_state:command_parsed_args(State),
-    ok = maybe_cover_compile(State, RawOpts),
-    ProjectApps = project_apps(State),
-    resolve_apps(ProjectApps, RawOpts).
-
-maybe_cover_compile(State, Opts) ->
-    State1 = case proplists:get_value(cover, Opts, false) of
-        true  -> rebar_state:set(State, cover_enabled, true);
-        false -> State
-    end,
-    rebar_prv_cover:maybe_cover_compile(State1).
+    resolve_apps(State, RawOpts).
 
-resolve_apps(ProjectApps, RawOpts) ->
+resolve_apps(State, RawOpts) ->
     case proplists:get_value(app, RawOpts) of
-        undefined -> resolve_suites(ProjectApps, RawOpts);
+        undefined -> resolve_suites(State, RawOpts);
         %% convert app name strings to `rebar_app_info` objects
         Apps      -> AppNames = string:tokens(Apps, [$,]),
+                     ProjectApps = project_apps(State),
                      case filter_apps_by_name(AppNames, ProjectApps) of
-                         {ok, TestApps} -> resolve_suites(TestApps, RawOpts);
+                         {ok, TestApps} -> resolve_suites(State, TestApps, RawOpts);
                          Error          -> Error
                      end
     end.
 
-resolve_suites(Apps, RawOpts) ->
+resolve_suites(State, RawOpts) -> resolve_suites(State, project_apps(State), RawOpts).
+
+resolve_suites(State, Apps, RawOpts) ->
     case proplists:get_value(suite, RawOpts) of
-        undefined  -> test_set(Apps, all);
+        undefined  -> compile_tests(State, Apps, all, RawOpts);
         Suites     -> SuiteNames = string:tokens(Suites, [$,]),
                       case filter_suites_by_apps(SuiteNames, Apps) of
-                          {ok, S} -> test_set(Apps, S);
+                          {ok, S} -> compile_tests(State, Apps, S, RawOpts);
                           Error   -> Error
                       end
     end.
 
+compile_tests(State, TestApps, Suites, RawOpts) ->
+    F = fun(AppInfo) ->
+        AppDir = rebar_app_info:dir(AppInfo),
+        S = case rebar_app_info:state(AppInfo) of
+            undefined ->
+                C = rebar_config:consult(AppDir),
+                rebar_state:new(State, C, AppDir);
+            AppState ->
+                AppState
+        end,
+        ok = rebar_erlc_compiler:compile(replace_src_dirs(S),
+                                         ec_cnv:to_list(rebar_app_info:out_dir(AppInfo)))
+    end,
+    lists:foreach(F, TestApps),
+    ok = maybe_cover_compile(State, RawOpts),
+    {ok, test_set(TestApps, Suites)}.
+
+maybe_cover_compile(State, Opts) ->
+    State1 = case proplists:get_value(cover, Opts, false) of
+        true  -> rebar_state:set(State, cover_enabled, true);
+        false -> State
+    end,
+    rebar_prv_cover:maybe_cover_compile(State1).
+
 project_apps(State) ->
     filter_checkouts(rebar_state:project_apps(State)).
 
@@ -204,8 +222,20 @@ app_modules([App|Rest], Acc) ->
             app_modules(Rest, NewAcc)
     end.
 
-test_set(Apps, all) -> {ok, set_apps(Apps, [])};
-test_set(_Apps, Suites) -> {ok, set_suites(Suites, [])}.
+replace_src_dirs(State) ->
+    %% replace any `src_dirs` with the test dirs
+    ErlOpts = rebar_state:get(State, erl_opts, []),
+    StrippedOpts = filter_src_dirs(ErlOpts),
+    case rebar_dir:extra_src_dirs(State) of
+        [] -> rebar_state:set(State, erl_opts, [{src_dirs, ["test"]}|StrippedOpts]);
+        _  -> rebar_state:set(State, erl_opts, StrippedOpts)
+    end.
+
+filter_src_dirs(ErlOpts) ->
+    lists:filter(fun({src_dirs, _}) -> false; (_) -> true end, ErlOpts).
+
+test_set(Apps, all) -> set_apps(Apps, []);
+test_set(_Apps, Suites) -> set_suites(Suites, []).
 
 set_apps([], Acc) -> lists:reverse(Acc);
 set_apps([App|Rest], Acc) ->