From 115cffb8bcc18ebd776ba169ca41351c502737d9 Mon Sep 17 00:00:00 2001
From: Arkadiusz Gil <arkgil@users.noreply.github.com>
Date: Tue, 24 Apr 2018 15:29:06 +0200
Subject: [PATCH] Add send_iq_and_wait_for_result/2,3 helpers (#175)

These functions send an IQ stanza and wait for response. They assert
that the response is the IQ of type "result", and that its ID matches the
ID of the sent request.
---
 src/escalus.erl        |  8 ++++++++
 src/escalus_client.erl | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/src/escalus.erl b/src/escalus.erl
index f6279120..361dbaf7 100644
--- a/src/escalus.erl
+++ b/src/escalus.erl
@@ -41,6 +41,8 @@
          wait_for_stanza/2,
          wait_for_stanzas/2,
          wait_for_stanzas/3,
+         send_iq_and_wait_for_result/2,
+         send_iq_and_wait_for_result/3,
          peek_stanzas/1]).
 
 -export_type([client/0,
@@ -165,6 +167,12 @@ wait_for_stanzas(Client, Count, Timeout) ->
 peek_stanzas(Client) ->
     escalus_client:peek_stanzas(Client).
 
+send_iq_and_wait_for_result(Client, Iq) ->
+    escalus_client:send_iq_and_wait_for_result(Client, Iq).
+
+send_iq_and_wait_for_result(Client, Iq, Timeout) ->
+    escalus_client:send_iq_and_wait_for_result(Client, Iq, Timeout).
+
 %% Other functions
 
 override(Config, OverrideName, NewValue) ->
diff --git a/src/escalus_client.erl b/src/escalus_client.erl
index 93535829..bf658da3 100644
--- a/src/escalus_client.erl
+++ b/src/escalus_client.erl
@@ -28,6 +28,7 @@
          peek_stanzas/1, has_stanzas/1,
          wait_for_stanzas/2, wait_for_stanzas/3,
          wait_for_stanza/1, wait_for_stanza/2,
+         send_iq_and_wait_for_result/2, send_iq_and_wait_for_result/3,
          is_client/1,
          full_jid/1,
          short_jid/1,
@@ -150,6 +151,38 @@ send_and_wait(Client, Packet) ->
     ok = send(Client, Packet),
     wait_for_stanza(Client).
 
+-spec send_iq_and_wait_for_result(client(), exml:element()) -> exml:element() | no_return().
+send_iq_and_wait_for_result(Client, Iq) ->
+    send_iq_and_wait_for_result(Client, Iq, ?WAIT_FOR_STANZA_TIMEOUT).
+
+-spec send_iq_and_wait_for_result(client(), exml:element(), non_neg_integer()) ->
+    exml:element() | no_return().
+send_iq_and_wait_for_result(Client, #xmlel{name = <<"iq">>} = Req, Timeout) ->
+    ok = send(Client, Req),
+    Resp = #xmlel{name = RespName} = wait_for_stanza(Client, Timeout),
+    RespType = exml_query:attr(Resp, <<"type">>, undefined),
+    RespId = exml_query:attr(Resp, <<"id">>),
+    ReqId = exml_query:attr(Req, <<"id">>),
+    case {RespName, RespType, RespId == ReqId} of
+        {<<"iq">>, <<"result">>, true} ->
+            Resp;
+        {<<"iq">>, <<"result">>, false} ->
+            raise_invalid_iq_resp_error(received_invalid_iq_result_id, ReqId, RespId, Req, Resp);
+        {<<"iq">>, _, _} ->
+            raise_invalid_iq_resp_error(received_invalid_iq_stanza_type, <<"result">>, RespType,
+                                        Req, Resp);
+        {_, _, _} ->
+            raise_invalid_iq_resp_error(received_invalid_stanza, <<"iq">>, RespName, Req, Resp)
+    end.
+
+-spec raise_invalid_iq_resp_error(atom(), term(), term(), exml:element(), exml:element()) ->
+    no_return().
+raise_invalid_iq_resp_error(Reason, Expected, Received, Req, Resp) ->
+    error({Reason, [{expected, Expected},
+                    {received, Received},
+                    {request, Req},
+                    {response, Resp}]}).
+
 -spec is_client(term()) -> boolean().
 is_client(#client{}) ->
     true;