From e7ccd88719b83f178dc5d74383db74a70fc8d1ab Mon Sep 17 00:00:00 2001
From: JianBo He <heeejianbo@163.com>
Date: Mon, 8 Nov 2021 15:05:54 +0800
Subject: [PATCH] test(props): cover messages headers

---
 apps/emqx_exhook/src/emqx_exhook_handler.erl  | 14 +++++++++-----
 .../emqx_exhook/test/emqx_exhook_demo_svr.erl |  9 +++++++--
 .../test/props/prop_exhook_hooks.erl          | 19 +++++++++++++------
 3 files changed, 29 insertions(+), 13 deletions(-)

diff --git a/apps/emqx_exhook/src/emqx_exhook_handler.erl b/apps/emqx_exhook/src/emqx_exhook_handler.erl
index bb2380e9b..cb0f2467e 100644
--- a/apps/emqx_exhook/src/emqx_exhook_handler.erl
+++ b/apps/emqx_exhook/src/emqx_exhook_handler.erl
@@ -49,6 +49,7 @@
 
 %% Utils
 -export([ message/1
+        , headers/1
         , stringfy/1
         , merge_responsed_bool/2
         , merge_responsed_message/2
@@ -301,11 +302,14 @@ assign_to_message(InMessage = #{qos := Qos, topic := Topic,
     enrich_header(maps:get(headers, InMessage, #{}), NMsg).
 
 enrich_header(Headers, Message) ->
-    AllowPub = case maps:get(<<"allow_publish">>, Headers, <<"true">>) of
-                   <<"false">> -> false;
-                   _ -> true
-               end,
-    emqx_message:set_header(allow_publish, AllowPub, Message).
+    case maps:get(<<"allow_publish">>, Headers, undefined) of
+        <<"false">> ->
+            emqx_message:set_header(allow_publish, false, Message);
+        <<"true">> ->
+            emqx_message:set_header(allow_publish, true, Message);
+        _ ->
+            Message
+    end.
 
 topicfilters(Tfs) when is_list(Tfs) ->
     [#{name => Topic, qos => Qos} || {Topic, #{qos := Qos}} <- Tfs].
diff --git a/apps/emqx_exhook/test/emqx_exhook_demo_svr.erl b/apps/emqx_exhook/test/emqx_exhook_demo_svr.erl
index 0bc4680c9..aa1d02a15 100644
--- a/apps/emqx_exhook/test/emqx_exhook_demo_svr.erl
+++ b/apps/emqx_exhook/test/emqx_exhook_demo_svr.erl
@@ -306,8 +306,8 @@ on_message_publish(#{message := #{from := From} = Msg} = Req, Md) ->
             {ok, #{type => 'STOP_AND_RETURN',
                    value => {message, NMsg}}, Md};
         <<"gooduser">> ->
-            NMsg = Msg#{topic => From,
-                        payload => From},
+            NMsg = allow(Msg#{topic => From,
+                              payload => From}),
             {ok, #{type => 'STOP_AND_RETURN',
                    value => {message, NMsg}}, Md};
         _ ->
@@ -319,6 +319,11 @@ deny(Msg) ->
                        maps:get(headers, Msg, #{})),
     maps:put(headers, NHeader, Msg).
 
+allow(Msg) ->
+    NHeader = maps:put(<<"allow_publish">>, <<"true">>,
+                       maps:get(headers, Msg, #{})),
+    maps:put(headers, NHeader, Msg).
+
 -spec on_message_delivered(emqx_exhook_pb:message_delivered_request(), grpc:metadata())
     -> {ok, emqx_exhook_pb:empty_success(), grpc:metadata()}
      | {error, grpc_cowboy_h:error_response()}.
diff --git a/apps/emqx_exhook/test/props/prop_exhook_hooks.erl b/apps/emqx_exhook/test/props/prop_exhook_hooks.erl
index c7804346f..cbd7a2a2a 100644
--- a/apps/emqx_exhook/test/props/prop_exhook_hooks.erl
+++ b/apps/emqx_exhook/test/props/prop_exhook_hooks.erl
@@ -296,19 +296,24 @@ prop_message_publish() ->
                 _ ->
                     ExpectedOutMsg = case emqx_message:from(Msg) of
                                          <<"baduser">> ->
-                                             MsgMap = emqx_message:to_map(Msg),
+                                             MsgMap = #{headers := Headers}
+                                                    = emqx_message:to_map(Msg),
                                              emqx_message:from_map(
                                                MsgMap#{qos => 0,
                                                        topic => <<"">>,
-                                                       payload => <<"">>
+                                                       payload => <<"">>,
+                                                       headers => maps:put(allow_publish, false, Headers)
                                                       });
                                          <<"gooduser">> = From ->
-                                             MsgMap = emqx_message:to_map(Msg),
+                                             MsgMap = #{headers := Headers}
+                                                    = emqx_message:to_map(Msg),
                                              emqx_message:from_map(
                                                MsgMap#{topic => From,
-                                                       payload => From
+                                                       payload => From,
+                                                       headers => maps:put(allow_publish, true, Headers)
                                                       });
-                                         _ -> Msg
+                                         _ ->
+                                             Msg
                                      end,
                     ?assertEqual(ExpectedOutMsg, OutMsg),
 
@@ -461,7 +466,9 @@ from_message(Msg) ->
       from => stringfy(emqx_message:from(Msg)),
       topic => emqx_message:topic(Msg),
       payload => emqx_message:payload(Msg),
-      timestamp => emqx_message:timestamp(Msg)
+      timestamp => emqx_message:timestamp(Msg),
+      headers => emqx_exhook_handler:headers(
+                  emqx_message:get_headers(Msg))
      }.
 
 %%--------------------------------------------------------------------