diff --git a/test/ejabberd_SUITE.erl b/test/ejabberd_SUITE.erl
index dc528d79a..470055654 100644
--- a/test/ejabberd_SUITE.erl
+++ b/test/ejabberd_SUITE.erl
@@ -1663,7 +1663,11 @@ mam_query_all(Config, NS) ->
     QID = randoms:get_string(),
     MyJID = my_jid(Config),
     Peer = ?config(slave, Config),
-    I = send(Config, #iq{type = get, sub_els = [#mam_query{xmlns = NS, id = QID}]}),
+    Type = case NS of
+	       ?NS_MAM_TMP -> get;
+	       _ -> set
+	   end,
+    I = send(Config, #iq{type = Type, sub_els = [#mam_query{xmlns = NS, id = QID}]}),
     maybe_recv_iq_result(NS, I),
     Iter = if NS == ?NS_MAM_TMP -> lists:seq(1, 5);
 	      true -> lists:seq(1, 5) ++ lists:seq(1, 5)
@@ -1692,15 +1696,15 @@ mam_query_all(Config, NS) ->
 mam_query_with(Config, JID, NS) ->
     MyJID = my_jid(Config),
     Peer = ?config(slave, Config),
-    Query = if NS == ?NS_MAM_TMP ->
-		    #mam_query{xmlns = NS, with = JID};
+    {Query, Type} = if NS == ?NS_MAM_TMP ->
+		    {#mam_query{xmlns = NS, with = JID}, get};
 	       true ->
 		    Fs = [#xdata_field{var = <<"jid">>,
 				       values = [jlib:jid_to_string(JID)]}],
-		    #mam_query{xmlns = NS,
-			       xdata = #xdata{type = submit, fields = Fs}}
+		    {#mam_query{xmlns = NS,
+			       xdata = #xdata{type = submit, fields = Fs}}, set}
 	    end,
-    I = send(Config, #iq{type = get, sub_els = [Query]}),
+    I = send(Config, #iq{type = Type, sub_els = [Query]}),
     Iter = if NS == ?NS_MAM_TMP -> lists:seq(1, 5);
 	      true -> lists:seq(1, 5) ++ lists:seq(1, 5)
 	   end,
@@ -1733,9 +1737,13 @@ maybe_recv_iq_result(_, _) ->
 mam_query_rsm(Config, NS) ->
     MyJID = my_jid(Config),
     Peer = ?config(slave, Config),
+    Type = case NS of
+	       ?NS_MAM_TMP -> get;
+	       _ -> set
+	   end,
     %% Get the first 3 items out of 5
     I1 = send(Config,
-              #iq{type = get,
+              #iq{type = Type,
                   sub_els = [#mam_query{xmlns = NS, rsm = #rsm_set{max = 3}}]}),
     maybe_recv_iq_result(NS, I1),
     lists:foreach(
@@ -1764,7 +1772,7 @@ mam_query_rsm(Config, NS) ->
     %% Get the next items starting from the `Last`.
     %% Limit the response to 2 items.
     I2 = send(Config,
-              #iq{type = get,
+              #iq{type = Type,
                   sub_els = [#mam_query{xmlns = NS,
 					rsm = #rsm_set{max = 2,
                                                        'after' = Last}}]}),
@@ -1800,7 +1808,7 @@ mam_query_rsm(Config, NS) ->
     end,
     %% Paging back. Should receive 2 elements: 2, 3.
     I3 = send(Config,
-              #iq{type = get,
+              #iq{type = Type,
                   sub_els = [#mam_query{xmlns = NS,
 					rsm = #rsm_set{max = 2,
                                                        before = First}}]}),
@@ -1829,7 +1837,7 @@ mam_query_rsm(Config, NS) ->
     end,
     %% Getting the item count. Should be 5 (or 10).
     I4 = send(Config,
-	      #iq{type = get,
+	      #iq{type = Type,
 		  sub_els = [#mam_query{xmlns = NS,
 					rsm = #rsm_set{max = 0}}]}),
     maybe_recv_iq_result(NS, I4),