Update the test suite to support new XMPP socket API

1 files changed, 72 insertions, 30 deletions

diff --git a/test/suite.erl b/test/suite.erl
index ce759446e..f66401c5e 100644
--- a/test/suite.erl
+++ b/test/suite.erl
@@ -117,7 +117,7 @@ init_config(Config) ->
      {rosterver, false},
      {lang, <<"en">>},
      {base_dir, BaseDir},
-     {socket, undefined},
+     {receiver, undefined},
      {pubsub_node, <<"node!@#$%^&*()'\"`~<>+-/;:_=[]{}|\\">>},
      {pubsub_node_title, <<"title!@#$%^&*()'\"`~<>+-/;:_=[]{}|\\">>},
      {resource, <<"resource!@#$%^&*()'\"`~<>+-/;:_=[]{}|\\">>},
@@ -201,7 +201,7 @@ connect(Config) ->
     end.
 
 tcp_connect(Config) ->
-    case ?config(socket, Config) of
+    case ?config(receiver, Config) of
 	undefined ->
 	    Owner = self(),
 	    NS = case ?config(type, Config) of
@@ -209,13 +209,12 @@ tcp_connect(Config) ->
 		     server -> ?NS_SERVER;
 		     component -> ?NS_COMPONENT
 		 end,
-	    ReceiverPid = spawn(fun() -> receiver(NS, Owner) end),
-	    {ok, Sock} = ejabberd_socket:connect(
-			   ?config(server_host, Config),
-			   ?config(server_port, Config),
-			   [binary, {packet, 0}, {active, false}],
-			   infinity, ReceiverPid),
-	    set_opt(socket, Sock, Config);
+	    Server = ?config(server_host, Config),
+	    Port = ?config(server_port, Config),
+	    ReceiverPid = spawn(fun() ->
+					start_receiver(NS, Owner, Server, Port)
+				end),
+	    set_opt(receiver, ReceiverPid, Config);
 	_ ->
 	    Config
     end.
@@ -259,7 +258,6 @@ process_stream_features(Config) ->
 
 disconnect(Config) ->
     ct:comment("Disconnecting"),
-    Socket = ?config(socket, Config),
     try
 	send_text(Config, ?STREAM_TRAILER)
     catch exit:normal ->
@@ -267,13 +265,12 @@ disconnect(Config) ->
     end,
     receive {xmlstreamend, <<"stream:stream">>} -> ok end,
     flush(Config),
-    ejabberd_socket:close(Socket),
+    ok = recv_call(Config, close),
     ct:comment("Disconnected"),
-    set_opt(socket, undefined, Config).
+    set_opt(receiver, undefined, Config).
 
 close_socket(Config) ->
-    Socket = ?config(socket, Config),
-    ejabberd_socket:close(Socket),
+    ok = recv_call(Config, close),
     Config.
 
 starttls(Config) ->
@@ -289,18 +286,15 @@ starttls(Config, ShouldFail) ->
 	#starttls_failure{} ->
 	    ct:fail(starttls_failed);
 	#starttls_proceed{} ->
-	    {ok, TLSSocket} = ejabberd_socket:starttls(
-				?config(socket, Config),
-				[{certfile, ?config(certfile, Config)},
-				 connect]),
-	    set_opt(socket, TLSSocket, Config)
+	    ok = recv_call(Config, {starttls, ?config(certfile, Config)}),
+	    Config
     end.
 
 zlib(Config) ->
     send(Config, #compress{methods = [<<"zlib">>]}),
     receive #compressed{} -> ok end,
-    {ok, ZlibSocket} = ejabberd_socket:compress(?config(socket, Config)),
-    process_stream_features(init_stream(set_opt(socket, ZlibSocket, Config))).
+    ok = recv_call(Config, compress),
+    process_stream_features(init_stream(Config)).
 
 auth(Config) ->
     auth(Config, false).
@@ -442,7 +436,7 @@ wait_auth_SASL_result(Config, ShouldFail) ->
 	#sasl_success{} when ShouldFail ->
 	    ct:fail(sasl_auth_should_have_failed);
         #sasl_success{} ->
-            ejabberd_socket:reset_stream(?config(socket, Config)),
+	    ok = recv_call(Config, reset_stream),
             send(Config, stream_header(Config)),
 	    Type = ?config(type, Config),
 	    NS = if Type == client -> ?NS_CLIENT;
@@ -528,7 +522,7 @@ decode(El, NS, Opts) ->
     end.
 
 send_text(Config, Text) ->
-    ejabberd_socket:send(?config(socket, Config), Text).
+    ok = recv_call(Config, {send_text, Text}).
 
 send(State, Pkt) ->
     {NewID, NewPkt} = case Pkt of
@@ -774,23 +768,71 @@ get_roster(Config) ->
     {LUser, LServer, _} = jid:tolower(my_jid(Config)),
     mod_roster:get_roster(LUser, LServer).
 
-receiver(NS, Owner) ->
+recv_call(Config, Msg) ->
+    Receiver = ?config(receiver, Config),
+    Ref = make_ref(),
+    Receiver ! {Ref, Msg},
+    receive
+	{Ref, Reply} ->
+	    Reply
+    end.
+
+start_receiver(NS, Owner, Server, Port) ->
     MRef = erlang:monitor(process, Owner),
-    receiver(NS, Owner, MRef).
+    {ok, Socket} = xmpp_socket:connect(
+		     Server, Port,
+		     [binary, {packet, 0}, {active, false}], infinity),
+    receiver(NS, Owner, Socket, MRef).
 
-receiver(NS, Owner, MRef) ->
+receiver(NS, Owner, Socket, MRef) ->
     receive
+	{Ref, reset_stream} ->
+	    Socket1 = xmpp_socket:reset_stream(Socket),
+	    Owner ! {Ref, ok},
+	    receiver(NS, Owner, Socket1, MRef);
+	{Ref, {starttls, Certfile}} ->
+	    {ok, TLSSocket} = xmpp_socket:starttls(
+				Socket,
+				[{certfile, Certfile}, connect]),
+	    Owner ! {Ref, ok},
+	    receiver(NS, Owner, TLSSocket, MRef);
+	{Ref, compress} ->
+	    {ok, ZlibSocket} = xmpp_socket:compress(Socket),
+	    Owner ! {Ref, ok},
+	    receiver(NS, Owner, ZlibSocket, MRef);
+	{Ref, {send_text, Text}} ->
+	    ok = xmpp_socket:send(Socket, Text),
+	    Owner ! {Ref, ok},
+	    receiver(NS, Owner, Socket, MRef);
+	{Ref, close} ->
+	    xmpp_socket:close(Socket),
+	    Owner ! {Ref, ok},
+	    receiver(NS, Owner, Socket, MRef);
         {'$gen_event', {xmlstreamelement, El}} ->
 	    Owner ! decode_stream_element(NS, El),
-	    receiver(NS, Owner, MRef);
+	    receiver(NS, Owner, Socket, MRef);
 	{'$gen_event', {xmlstreamstart, Name, Attrs}} ->
 	    Owner ! decode(#xmlel{name = Name, attrs = Attrs}, <<>>, []),
-	    receiver(NS, Owner, MRef);
+	    receiver(NS, Owner, Socket, MRef);
 	{'$gen_event', Event} ->
             Owner ! Event,
-	    receiver(NS, Owner, MRef);
+	    receiver(NS, Owner, Socket, MRef);
 	{'DOWN', MRef, process, Owner, _} ->
-	    ok
+	    ok;
+	{tcp, _, Data} ->
+	    case xmpp_socket:recv(Socket, Data) of
+		{ok, Socket1} ->
+		    receiver(NS, Owner, Socket1, MRef);
+		{error, _} ->
+		    Owner ! closed,
+		    receiver(NS, Owner, Socket, MRef)
+	    end;
+	{tcp_error, _, _} ->
+	    Owner ! closed,
+	    receiver(NS, Owner, Socket, MRef);
+	{tcp_closed, _} ->
+	    Owner ! closed,
+	    receiver(NS, Owner, Socket, MRef)
     end.
 
 %%%===================================================================