Fix test suite

COUCHDB-2796
diff --git a/src/couch_epi_data_gen.erl b/src/couch_epi_data_gen.erl
index 75601cf..e4e0889 100644
--- a/src/couch_epi_data_gen.erl
+++ b/src/couch_epi_data_gen.erl
@@ -220,66 +220,76 @@
 -include_lib("eunit/include/eunit.hrl").
 
 basic_test() ->
-    Module = foo_bar_baz_bugz,
-    Data1 = [some_nice_data],
-    Data2 = "other data",
-    Data3 = {"even more data"},
-    Defs1 = [{foo, Data1}],
-    Defs2 = lists:usort([{foo, Data2}, {bar, Data3}]),
+    try
+        Module = foo_bar_baz_bugz,
 
-    set(Module, app1, Defs1),
-    set(Module, app2, Defs2),
+        meck:new(couch_epi_module_keeper, [passthrough]),
+        meck:expect(couch_epi_module_keeper, save, fun
+            (Handle, Source, Modules) -> save(Handle, Source, Modules)
+        end),
 
-    ?assertEqual([bar, foo], lists:usort(Module:keys())),
-    ?assertEqual([app1, app2], lists:usort(Module:subscribers())),
+        Data1 = [some_nice_data],
+        Data2 = "other data",
+        Data3 = {"even more data"},
+        Defs1 = [{foo, Data1}],
+        Defs2 = lists:usort([{foo, Data2}, {bar, Data3}]),
 
-    ?assertEqual(Data1, Module:get(app1, foo)),
-    ?assertEqual(Data2, Module:get(app2, foo)),
-    ?assertEqual(Data3, Module:get(app2, bar)),
+        set(Module, app1, Defs1),
+        set(Module, app2, Defs2),
 
-    ?assertEqual(undefined, Module:get(bad, key)),
-    ?assertEqual(undefined, Module:get(source, bad)),
+        ?assertEqual([bar, foo], lists:usort(Module:keys())),
+        ?assertEqual([app1, app2], lists:usort(Module:subscribers())),
 
-    ?assertEqual("3KZ4EG4WBF4J683W8GSDDPYR3", Module:version(app1)),
-    ?assertEqual("4EFUU47W9XDNMV9RMZSSJQU3Y", Module:version(app2)),
+        ?assertEqual(Data1, Module:get(app1, foo)),
+        ?assertEqual(Data2, Module:get(app2, foo)),
+        ?assertEqual(Data3, Module:get(app2, bar)),
 
-    ?assertEqual({error,{unknown,bad}}, Module:version(bad)),
+        ?assertEqual(undefined, Module:get(bad, key)),
+        ?assertEqual(undefined, Module:get(source, bad)),
 
-    ?assertEqual(
-        [{app1,"3KZ4EG4WBF4J683W8GSDDPYR3"},
-         {app2,"4EFUU47W9XDNMV9RMZSSJQU3Y"}], lists:usort(Module:version())),
+        ?assertEqual("3KZ4EG4WBF4J683W8GSDDPYR3", Module:version(app1)),
+        ?assertEqual("4EFUU47W9XDNMV9RMZSSJQU3Y", Module:version(app2)),
 
-    ?assertEqual(
-       [{app1,[some_nice_data]},{app2,"other data"}],
-       lists:usort(Module:by_key(foo))),
+        ?assertEqual({error,{unknown,bad}}, Module:version(bad)),
 
-    ?assertEqual([], lists:usort(Module:by_key(bad))),
+        ?assertEqual(
+            [{app1,"3KZ4EG4WBF4J683W8GSDDPYR3"},
+             {app2,"4EFUU47W9XDNMV9RMZSSJQU3Y"}], lists:usort(Module:version())),
 
-    ?assertEqual(
-        [
-            {bar, [{app2, {"even more data"}}]},
-            {foo, [{app2, "other data"}, {app1, [some_nice_data]}]}
-        ],
-        lists:usort(Module:by_key())),
+        ?assertEqual(
+           [{app1,[some_nice_data]},{app2,"other data"}],
+           lists:usort(Module:by_key(foo))),
+
+        ?assertEqual([], lists:usort(Module:by_key(bad))),
+
+        ?assertEqual(
+            [
+                {bar, [{app2, {"even more data"}}]},
+                {foo, [{app2, "other data"}, {app1, [some_nice_data]}]}
+            ],
+            lists:usort(Module:by_key())),
 
 
-    ?assertEqual(Defs1, lists:usort(Module:by_source(app1))),
-    ?assertEqual(Defs2, lists:usort(Module:by_source(app2))),
+        ?assertEqual(Defs1, lists:usort(Module:by_source(app1))),
+        ?assertEqual(Defs2, lists:usort(Module:by_source(app2))),
 
-    ?assertEqual([], lists:usort(Module:by_source(bad))),
+        ?assertEqual([], lists:usort(Module:by_source(bad))),
 
-    ?assertEqual(
-        [
-            {app1, [{foo, [some_nice_data]}]},
-            {app2, [{foo, "other data"}, {bar, {"even more data"}}]}
-        ],
-        lists:usort(Module:by_source())),
+        ?assertEqual(
+            [
+                {app1, [{foo, [some_nice_data]}]},
+                {app2, [{foo, "other data"}, {bar, {"even more data"}}]}
+            ],
+            lists:usort(Module:by_source())),
 
-    ?assertEqual(
-        lists:usort([Data1, Data2, Data3]), lists:usort(Module:all())),
-    ?assertEqual(lists:usort([Data1, Data2]), lists:usort(Module:all(foo))),
-    ?assertEqual([], lists:usort(Module:all(bad))),
-
+        ?assertEqual(
+            lists:usort([Data1, Data2, Data3]), lists:usort(Module:all())),
+        ?assertEqual(lists:usort([Data1, Data2]), lists:usort(Module:all(foo))),
+        ?assertEqual([], lists:usort(Module:all(bad))),
+        ok
+    after
+        meck:unload(couch_epi_module_keeper)
+    end,
     ok.
 
 -endif.
diff --git a/src/couch_epi_functions_gen.erl b/src/couch_epi_functions_gen.erl
index 6d93787..84ecc5f 100644
--- a/src/couch_epi_functions_gen.erl
+++ b/src/couch_epi_functions_gen.erl
@@ -330,26 +330,36 @@
     [].
 
 basic_test() ->
-    Module = foo_bar_dispatcher,
-    add(Module, app1, [?MODULE]),
+    try
+        Module = foo_bar_dispatcher,
+        meck:new(couch_epi_module_keeper, [passthrough]),
+        meck:expect(couch_epi_module_keeper, save, fun
+            (Handle, Source, Modules) -> save(Handle, Source, Modules)
+        end),
 
-    ?assertMatch([?MODULE], modules(Module, foo, 2)),
+        add(Module, app1, [?MODULE]),
 
-    ?assert(is_list(Module:version(app1))),
+        ?assertMatch([?MODULE], modules(Module, foo, 2)),
 
-    Defs1 = lists:usort(Module:definitions()),
-    ?assertMatch([{app1, [{?MODULE, _}]}], Defs1),
-    [{app1, [{?MODULE, Exports}]}] = Defs1,
-    ?assert(lists:member({bar, 0}, Exports)),
+        ?assert(is_list(Module:version(app1))),
 
-    add(Module, app2, [?MODULE]),
-    Defs2 = lists:usort(Module:definitions()),
-    ?assertMatch([{app1, [{?MODULE, _}]}, {app2, [{?MODULE, _}]}], Defs2),
+        Defs1 = lists:usort(Module:definitions()),
+        ?assertMatch([{app1, [{?MODULE, _}]}], Defs1),
+        [{app1, [{?MODULE, Exports}]}] = Defs1,
+        ?assert(lists:member({bar, 0}, Exports)),
 
-    ?assertMatch([{app1, Hash}, {app2, Hash}], Module:version()),
+        add(Module, app2, [?MODULE]),
+        Defs2 = lists:usort(Module:definitions()),
+        ?assertMatch([{app1, [{?MODULE, _}]}, {app2, [{?MODULE, _}]}], Defs2),
 
-    ?assertMatch([], Module:dispatch(?MODULE, bar, [])),
-    ?assertMatch({1, 2}, Module:dispatch(?MODULE, foo, [1, 2])),
+        ?assertMatch([{app1, Hash}, {app2, Hash}], Module:version()),
+
+        ?assertMatch([], Module:dispatch(?MODULE, bar, [])),
+        ?assertMatch({1, 2}, Module:dispatch(?MODULE, foo, [1, 2])),
+        ok
+    after
+        meck:unload(couch_epi_module_keeper)
+    end,
 
     ok.
 
diff --git a/test/couch_epi_data_source_tests.erl b/test/couch_epi_data_source_tests.erl
deleted file mode 100644
index f5d701f..0000000
--- a/test/couch_epi_data_source_tests.erl
+++ /dev/null
@@ -1,90 +0,0 @@
-% Licensed under the Apache License, Version 2.0 (the "License"); you may not
-% use this file except in compliance with the License. You may obtain a copy of
-% the License at
-%
-% http://www.apache.org/licenses/LICENSE-2.0
-%
-% Unless required by applicable law or agreed to in writing, software
-% distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
-% WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
-% License for the specific language governing permissions and limitations under
-% the License.
-
--module(couch_epi_data_source_tests).
-
--include_lib("couch/include/couch_eunit.hrl").
-
--define(DATA_FILE1, ?ABS_PATH("test/fixtures/app_data1.cfg")).
--define(DATA_FILE2, ?ABS_PATH("test/fixtures/app_data2.cfg")).
-
--record(ctx, {file, handle, pid}).
-
-setup() ->
-    Key = {test_app, descriptions},
-    File = ?tempfile(),
-    {ok, _} = file:copy(?DATA_FILE1, File),
-    {ok, Pid} = couch_epi_data_source:start_link(
-        test_app, {epi_key, Key}, {file, File}, [{interval, 100}]),
-    ok = couch_epi_data_source:wait(Pid),
-    #ctx{
-        pid = Pid,
-        file = File,
-        handle = couch_epi_data_gen:get_handle(Key)}.
-
-
-teardown(#ctx{pid = Pid, file = File}) ->
-    file:delete(File),
-    couch_epi_data_source:stop(Pid),
-    catch meck:unload(compile),
-    ok.
-
-
-epi_data_source_reload_test_() ->
-    {
-        "data_source reload tests",
-        {
-            foreach,
-            fun setup/0,
-            fun teardown/1,
-            [
-                fun ensure_reload_if_manually_triggered/1,
-                fun ensure_reload_if_changed/1,
-                fun ensure_no_reload_when_no_change/1
-            ]
-        }
-    }.
-
-ensure_reload_if_manually_triggered(#ctx{pid = Pid, file = File}) ->
-    ?_test(begin
-        ok = meck:new(compile, [passthrough, unstick]),
-        ok = meck:expect(compile, forms, fun(_, _) -> {error, reload} end),
-        {ok, _} = file:copy(?DATA_FILE2, File),
-        Result = couch_epi_data_source:reload(Pid),
-        ?assertMatch({error,{badmatch,{error,reload}}}, Result)
-    end).
-
-ensure_reload_if_changed(#ctx{file = File, handle = Handle}) ->
-    ?_test(begin
-        ?assertMatch(
-            [[{type,counter},{desc,foo}]],
-            couch_epi_data_gen:get(Handle, [complex, key, 1])),
-        {ok, _} = file:copy(?DATA_FILE2, File),
-        timer:sleep(150),
-        ?assertMatch(
-            [[{type,counter},{desc,bar}]],
-            couch_epi_data_gen:get(Handle, [complex, key, 2]))
-    end).
-
-ensure_no_reload_when_no_change(#ctx{handle = Handle}) ->
-    ok = meck:new(compile, [passthrough, unstick]),
-    ok = meck:expect(compile, forms, fun(_, _) ->
-        {error, compile_should_not_be_called} end),
-    ?_test(begin
-        ?assertMatch(
-            [[{type,counter},{desc,foo}]],
-            couch_epi_data_gen:get(Handle, [complex, key, 1])),
-        timer:sleep(200),
-        ?assertMatch(
-            [],
-            couch_epi_data_gen:get(Handle, [complex, key, 2]))
-    end).
diff --git a/test/couch_epi_functions_tests.erl b/test/couch_epi_functions_tests.erl
deleted file mode 100644
index 6b035b9..0000000
--- a/test/couch_epi_functions_tests.erl
+++ /dev/null
@@ -1,126 +0,0 @@
-% Licensed under the Apache License, Version 2.0 (the "License"); you may not
-% use this file except in compliance with the License. You may obtain a copy of
-% the License at
-%
-% http://www.apache.org/licenses/LICENSE-2.0
-%
-% Unless required by applicable law or agreed to in writing, software
-% distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
-% WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
-% License for the specific language governing permissions and limitations under
-% the License.
-
--module(couch_epi_functions_tests).
-
--include_lib("couch/include/couch_eunit.hrl").
-
--define(MODULE1(Name), "
-    -export([foo/2, bar/0, inc/1]).
-    foo(A1, A2) ->
-        {A1, A2}.
-
-    bar() ->
-        [].
-
-    inc(A) ->
-        A + 1.
-").
-
--define(MODULE2(Name), "
-    -export([baz/1, inc/1]).
-    baz(A1) ->
-        A1.
-
-    inc(A) ->
-        A + 1.
-").
-
-setup() ->
-    setup([{interval, 100}]).
-
-setup(Opts) ->
-    ServiceId = my_service,
-    Module = my_test_module,
-    ok = generate_module(Module, ?MODULE1(Module)),
-    {ok, Pid} = couch_epi_functions:start_link(
-        test_app, {epi_key, ServiceId}, {modules, [Module]}, Opts),
-    ok = couch_epi_functions:wait(Pid),
-    {Pid, Module, ServiceId, couch_epi_functions_gen:get_handle(ServiceId)}.
-
-teardown({Pid, Module, _, _Handle}) ->
-    code:purge(Module),
-    couch_epi_functions:stop(Pid),
-    catch meck:unload(compile),
-    ok.
-
-generate_module(Name, Body) ->
-    Tokens = couch_epi_codegen:scan(Body),
-    couch_epi_codegen:generate(Name, Tokens).
-
-upgrade_release(Pid) ->
-    sys:suspend(Pid),
-    'ok' = sys:change_code(Pid, couch_epi_functions, 'undefined', []),
-    sys:resume(Pid),
-    ok.
-
-epi_functions_test_() ->
-    {
-        "functions reload tests",
-        {
-            foreach,
-            fun setup/0,
-            fun teardown/1,
-            [
-                fun ensure_reload_if_changed/1,
-                fun ensure_no_reload_when_no_change/1
-            ]
-        }
-    }.
-
-epi_functions_manual_reload_test_() ->
-    {
-        "functions manual reload tests",
-        {
-            foreach,
-            fun() -> setup([{interval, 10000}]) end,
-            fun teardown/1,
-            [
-                fun ensure_reload_if_manually_triggered/1
-            ]
-        }
-    }.
-
-ensure_reload_if_manually_triggered({Pid, Module, _ServiceId, _Handle}) ->
-    ?_test(begin
-        ok = generate_module(Module, ?MODULE2(Module)),
-        ok = meck:new(compile, [passthrough, unstick]),
-        ok = meck:expect(compile, forms, fun(_, _) -> {error, reload} end),
-        Result = couch_epi_functions:reload(Pid),
-        ?assertMatch({error,{badmatch,{error,reload}}}, Result)
-    end).
-
-ensure_reload_if_changed({Pid, Module, ServiceId, _Handle}) ->
-    ?_test(begin
-        ?assertMatch(
-            [{1, 2}],
-            couch_epi_functions_gen:apply(ServiceId, foo, [1, 2], [])),
-        ok = generate_module(Module, ?MODULE2(Module)),
-        upgrade_release(Pid),
-        ?assertMatch(
-            [3],
-            couch_epi_functions_gen:apply(ServiceId, baz, [3], []))
-    end).
-
-ensure_no_reload_when_no_change({Pid, _Module, ServiceId, _Handle}) ->
-    ok = meck:new(compile, [passthrough, unstick]),
-    ok = meck:expect(compile, forms, fun(_, _) ->
-        {error, compile_should_not_be_called} end),
-    ?_test(begin
-        ?assertMatch(
-            [{1, 2}],
-            couch_epi_functions_gen:apply(ServiceId, foo, [1, 2], [])),
-        upgrade_release(Pid),
-        ?assertMatch(
-            [],
-            couch_epi_functions_gen:apply(ServiceId, baz, [3], []))
-    end).
diff --git a/test/couch_epi_tests.erl b/test/couch_epi_tests.erl
index 227a79f..78eab4d 100644
--- a/test/couch_epi_tests.erl
+++ b/test/couch_epi_tests.erl
@@ -229,16 +229,22 @@
         couch_epi_data_source,
         couch_epi_functions
     ],
+    Funs = [
+        fun ensure_reload_if_manually_triggered/2,
+        fun ensure_reload_if_changed/2,
+        fun ensure_no_reload_when_no_change/2
+    ],
     {
         "epi reload tests",
         {
             foreachx,
             fun setup/1,
             fun teardown/2,
-            [{M, fun ensure_reloaded/2} || M <- Modules]
+            [{M, Fun} || M <- Modules, Fun <- Funs]
         }
     }.
 
+
 apply_options_test_() ->
     Funs = [fun ensure_apply_is_called/2],
     make_case("Apply with options: ", valid_options_permutations(), Funs).
@@ -416,16 +422,60 @@
     ?_assertMatch([test_app], couch_epi:subscribers(Handle)).
 
 
-ensure_reloaded(Module, #ctx{pid = Pid, key = Key} = Ctx) ->
+ensure_reload_if_manually_triggered(Module, #ctx{pid = Pid, key = Key} = Ctx) ->
     ?_test(begin
         subscribe(Ctx, test_app, Key),
         update_definitions(Module, Ctx),
         Module:reload(Pid),
-        timer:sleep(200),
+        timer:sleep(50),
         Result = get(Ctx, is_called),
         ?assertNotMatch(error, Result)
     end).
 
+ensure_reload_if_changed(couch_epi_data_source =  Module,
+        #ctx{key = Key, handle = Handle} = Ctx) ->
+    ?_test(begin
+        Version = Handle:version(),
+        subscribe(Ctx, test_app, Key),
+        update_definitions(Module, Ctx),
+        timer:sleep(250),
+        ?assertNotEqual(Version, Handle:version()),
+        Result = get(Ctx, is_called),
+        ?assertNotMatch(error, Result)
+    end);
+ensure_reload_if_changed(Module,
+        #ctx{key = Key, handle = Handle} = Ctx) ->
+    ?_test(begin
+        Version = Handle:version(),
+        subscribe(Ctx, test_app, Key),
+        update(Module, Ctx),
+        ?assertNotEqual(Version, Handle:version()),
+        timer:sleep(100), %% Allow some time for notify to be called
+        Result = get(Ctx, is_called),
+        ?assertNotMatch(error, Result)
+    end).
+
+ensure_no_reload_when_no_change(couch_epi_functions = Module,
+        #ctx{pid = Pid, key = Key, handle = Handle} = Ctx) ->
+    ?_test(begin
+        Version = Handle:version(),
+        subscribe(Ctx, test_app, Key),
+        upgrade_release(Pid, Module),
+        ?assertEqual(Version, Handle:version()),
+        Result = get(Ctx, is_called),
+        ?assertMatch(error, Result)
+    end);
+ensure_no_reload_when_no_change(Module,
+        #ctx{key = Key, handle = Handle} = Ctx) ->
+    ?_test(begin
+        Version = Handle:version(),
+        subscribe(Ctx, test_app, Key),
+        timer:sleep(450),
+        ?assertEqual(Version, Handle:version()),
+        Result = get(Ctx, is_called),
+        ?assertMatch(error, Result)
+    end).
+
 
 %% ------------------------------------------------------------------
 %% Internal Function Definitions