Hooks rework

big_tests/tests/mod_ping_SUITE.erl

@@ -107,24 +107,30 @@ start_mod_ping(Opts) ->
 setup_pong_hook(Config) ->
     Pid = self(),
     HostType = domain_helper:host_type(mim),
-    Handler = mongoose_helper:successful_rpc(?MODULE, setup_pong_hook, [HostType, Pid]),
-    [{pong_handler, Handler} | Config].
+    mongoose_helper:successful_rpc(?MODULE, setup_pong_hook, [HostType, Pid]),
+    [{pid, Pid} | Config].
 
 setup_pong_hook(HostType, Pid) ->
-    Handler = fun(_Acc, _HostType, JID, _Response, _TDelta) ->
-                      Pid ! {pong, jid:to_binary(jid:to_lower(JID))}
-              end,
-    ejabberd_hooks:add(user_ping_response, HostType, Handler, 50),
-    Handler.
+    gen_hook:add_handler(user_ping_response, HostType,
+                         fun ?MODULE:pong_hook_handler/3,
+                         #{pid => Pid}, 50).
+
+pong_hook_handler(Acc,
+                  #{jid := JID} = _Params,
+                  #{pid := Pid} = _Extra) ->
+    Pid ! {pong, jid:to_binary(jid:to_lower(JID))},
+    {ok, Acc}.
 
 clear_pong_hook(Config) ->
-    {value, {_, Handler}, NConfig} = lists:keytake(pong_handler, 1, Config),
+    {value, {_, Pid}, NConfig} = lists:keytake(pid, 1, Config),
     HostType = domain_helper:host_type(mim),
-    mongoose_helper:successful_rpc(?MODULE, clear_pong_hook, [HostType, Handler]),
+    mongoose_helper:successful_rpc(?MODULE, clear_pong_hook, [HostType, Pid]),
     NConfig.
 
-clear_pong_hook(HostType, Handler) ->
-    ejabberd_hooks:delete(user_ping_response, HostType, Handler, 50).
+clear_pong_hook(HostType, Pid) ->
+    gen_hook:delete_handler(user_ping_response, HostType,
+                            fun ?MODULE:pong_hook_handler/3,
+                            #{pid => Pid}, 50).
 
 %%--------------------------------------------------------------------
 %% Ping tests

big_tests/tests/push_SUITE.erl

@@ -134,14 +134,14 @@ init_per_testcase(CaseName = push_notifications_not_listed_disco_when_not_availa
     escalus:init_per_testcase(CaseName, Config);
 init_per_testcase(CaseName, Config0) ->
     Config1 = escalus_fresh:create_users(Config0, [{bob, 1}, {alice, 1}, {kate, 1}]),
-    Config = add_pubsub_jid([{case_name, CaseName} | Config1]),
+    Config2 = add_pubsub_jid([{case_name, CaseName} | Config1]),
 
-    case ?config(pubsub_host, Config0) of
-        virtual ->
-            start_hook_listener(Config);
-        _ ->
-            start_route_listener(Config)
-    end,
+    Config = case ?config(pubsub_host, Config0) of
+                 virtual ->
+                     start_hook_listener(Config2);
+                 _ ->
+                     start_route_listener(Config2)
+             end,
 
     escalus:init_per_testcase(CaseName, Config).
 
@@ -151,7 +151,12 @@ end_per_testcase(CaseName = push_notifications_listed_disco_when_available, Conf
 end_per_testcase(CaseName = push_notifications_not_listed_disco_when_not_available, Config) ->
     escalus:end_per_testcase(CaseName, Config);
 end_per_testcase(CaseName, Config) ->
-    rpc(ejabberd_router, unregister_route, [atom_to_binary(CaseName, utf8)]),
+    case ?config(pubsub_host, Config) of
+        virtual ->
+            stop_hook_listener(Config);
+        _ ->
+            stop_route_listener(Config)
+    end,
     escalus:end_per_testcase(CaseName, Config).
 
 %% --------------------- Helpers ------------------------
@@ -692,7 +697,12 @@ start_route_listener(Config) ->
                push_form_ns => push_helper:push_form_type() },
     Handler = rpc(mongoose_packet_handler, new, [?MODULE, #{state => State}]),
     Domain = pubsub_domain(Config),
-    rpc(ejabberd_router, register_route, [Domain, Handler]).
+    rpc(ejabberd_router, register_route, [Domain, Handler]),
+    Config.
+
+stop_route_listener(Config) ->
+    Domain = pubsub_domain(Config),
+    rpc(ejabberd_router, unregister_route, [Domain]).
 
 process_packet(_Acc, _From, To, El, #{state := State}) ->
     #{ pid := TestCasePid, pub_options_ns := PubOptionsNS, push_form_ns := PushFormNS } = State,
@@ -739,26 +749,41 @@ valid_ns_if_defined(NS, FormProplist) ->
 start_hook_listener(Config) ->
     TestCasePid = self(),
     PubSubJID = pubsub_jid(Config),
-    rpc(?MODULE, rpc_start_hook_handler, [TestCasePid, PubSubJID]).
+    rpc(?MODULE, rpc_start_hook_handler, [TestCasePid, PubSubJID]),
+    [{pid, TestCasePid}, {jid, PubSubJID} | Config].
+
+stop_hook_listener(Config) ->
+    TestCasePid = proplists:get_value(pid, Config),
+    PubSubJID = proplists:get_value(jid, Config),
+    rpc(?MODULE, rpc_stop_hook_handler, [TestCasePid, PubSubJID]).
 
 rpc_start_hook_handler(TestCasePid, PubSubJID) ->
-    Handler = fun(Acc, _Host, [PayloadMap], OptionMap) ->
-                  try jid:to_binary(mongoose_acc:get(push_notifications, pubsub_jid, Acc)) of
-                      PubSubJIDBin when PubSubJIDBin =:= PubSubJID->
-                          TestCasePid ! push_notification(PubSubJIDBin,
-                                                          maps:to_list(PayloadMap),
-                                                          maps:to_list(OptionMap));
-                      _ -> ok
-                  catch
-                      C:R:S ->
-                          TestCasePid ! #{ event => handler_error,
-                                           class => C,
-                                           reason => R,
-                                           stacktrace => S }
-                  end,
-                  Acc
-              end,
-    ejabberd_hooks:add(push_notifications, <<"localhost">>, Handler, 50).
+    gen_hook:add_handler(push_notifications,  <<"localhost">>,
+                         fun ?MODULE:hook_handler_fn/3,
+                         #{pid => TestCasePid, jid => PubSubJID}, 50).
+
+hook_handler_fn(Acc,
+                #{notification_forms := [PayloadMap], options := OptionMap} = _Params,
+                #{pid := TestCasePid, jid := PubSubJID} = _Extra) ->
+    try jid:to_binary(mongoose_acc:get(push_notifications, pubsub_jid, Acc)) of
+        PubSubJIDBin when PubSubJIDBin =:= PubSubJID ->
+            TestCasePid ! push_notification(PubSubJIDBin,
+                                            maps:to_list(PayloadMap),
+                                            maps:to_list(OptionMap));
+        _ -> ok
+    catch
+        C:R:S ->
+            TestCasePid ! #{event => handler_error,
+                            class => C,
+                            reason => R,
+                            stacktrace => S}
+    end,
+    {ok, Acc}.
+
+rpc_stop_hook_handler(TestCasePid, PubSubJID) ->
+    gen_hook:delete_handler(push_notifications, <<"localhost">>,
+                            fun ?MODULE:hook_handler_fn/3,
+                            #{pid => TestCasePid, jid => PubSubJID}, 50).
 
 %%--------------------------------------------------------------------
 %% Test helpers

big_tests/tests/sm_SUITE.erl

@@ -865,7 +865,7 @@ unacknowledged_message_hook_common(RestartConnectionFN, Config) ->
     AliceSpec0 = escalus_fresh:create_fresh_user(Config, alice),
     Resource = proplists:get_value(username, AliceSpec0),
     AliceSpec = [{resource, Resource} | AliceSpec0],
-    start_hook_listener(Resource),
+    HookHandlerExtra = start_hook_listener(Resource),
     {ok, Alice, _} = escalus_connection:start(AliceSpec, ConnSteps ++ [stream_resumption]),
     SMID = proplists:get_value(smid, Alice#client.props),
     escalus_connection:send(Alice, escalus_stanza:presence(<<"available">>)),
@@ -912,7 +912,7 @@ unacknowledged_message_hook_common(RestartConnectionFN, Config) ->
     escalus:assert(is_chat_message, [<<"msg-3">>], wait_for_unacked_msg_hook(1, NewResource, 100)),
     escalus:assert(is_chat_message, [<<"msg-4">>], wait_for_unacked_msg_hook(1, NewResource, 100)),
     ?assertEqual(timeout, wait_for_unacked_msg_hook(0, Resource, 100)),
-
+    stop_hook_listener(HookHandlerExtra),
     escalus_connection:stop(Bob).
 
 resume_session(Config) ->
@@ -1323,20 +1323,34 @@ start_hook_listener(Resource) ->
     TestCasePid = self(),
     rpc(mim(), ?MODULE, rpc_start_hook_handler, [TestCasePid, Resource, host_type()]).
 
+stop_hook_listener(HookExtra) ->
+    rpc(mim(), ?MODULE, rpc_stop_hook_handler, [HookExtra, host_type()]).
+
 rpc_start_hook_handler(TestCasePid, User, HostType) ->
-    LUser=jid:nodeprep(User),
-    Handler = fun(Acc, Jid) ->
-                {U, _S, R} = jid:to_lower(Jid),
-                case U of
-                    LUser ->
-                        Counter = mongoose_acc:get(sm_test, counter, 0, Acc),
-                        El = mongoose_acc:element(Acc),
-                        TestCasePid ! {sm_test, Counter, R, El},
-                        mongoose_acc:set_permanent(sm_test, counter, Counter + 1, Acc);
-                    _ -> Acc
-                end
-              end,
-    ejabberd_hooks:add(unacknowledged_message, HostType, Handler, 50).
+    LUser = jid:nodeprep(User),
+    Extra = #{luser => LUser, pid => TestCasePid},
+    gen_hook:add_handler(unacknowledged_message, HostType,
+                         fun ?MODULE:hook_handler_fn/3,
+                         Extra, 50),
+    Extra.
+
+rpc_stop_hook_handler(HookExtra, HostType) ->
+    gen_hook:delete_handler(unacknowledged_message, HostType,
+                            fun ?MODULE:hook_handler_fn/3,
+                            HookExtra, 50).
+
+hook_handler_fn(Acc,
+                #{args := [Jid]} = _Params,
+                #{luser := LUser, pid := TestCasePid} = _Extra) ->
+    {U, _S, R} = jid:to_lower(Jid),
+    case U of
+        LUser ->
+            Counter = mongoose_acc:get(sm_test, counter, 0, Acc),
+            El = mongoose_acc:element(Acc),
+            TestCasePid ! {sm_test, Counter, R, El},
+            {ok, mongoose_acc:set_permanent(sm_test, counter, Counter + 1, Acc)};
+        _ -> {ok, Acc}
+    end.
 
 wait_for_unacked_msg_hook(Counter, Res, Timeout) ->
     receive