Use util to simplify connection establishment in suites

3 files changed, 56 insertions, 139 deletions

diff --git a/lib/diameter/test/diameter_relay_SUITE.erl b/lib/diameter/test/diameter_relay_SUITE.erl
index f7f65197ed..03f1115496 100644
--- a/lib/diameter/test/diameter_relay_SUITE.erl
+++ b/lib/diameter/test/diameter_relay_SUITE.erl
@@ -110,19 +110,6 @@
                         {module, ?MODULE},
                         {answer_errors, callback}]}]).
 
-%% Config for diameter:add_transport/2. In the listening case, listen
-%% on a free port that we then lookup using the implementation detail
-%% that diameter_tcp registers the port with diameter_reg.
--define(CONNECT(PortNr),
-        {connect, [{transport_module, diameter_tcp},
-                   {transport_config, [{raddr, ?ADDR},
-                                       {rport, PortNr},
-                                       {ip, ?ADDR},
-                                       {port, 0}]}]}).
--define(LISTEN,
-        {listen, [{transport_module, diameter_tcp},
-                  {transport_config, [{ip, ?ADDR}, {port, 0}]}]}).
-
 -define(SUCCESS, 2001).
 -define(LOOP_DETECTED, 3005).
 -define(UNABLE_TO_DELIVER, 3002).
@@ -130,9 +117,6 @@
 -define(LOGOUT, ?'DIAMETER_BASE_TERMINATION-CAUSE_DIAMETER_LOGOUT').
 -define(AUTHORIZE_ONLY, ?'DIAMETER_BASE_RE-AUTH-REQUEST-TYPE_AUTHORIZE_ONLY').
 
--define(A, list_to_atom).
--define(L, atom_to_list).
-
 %% ===========================================================================
 
 suite() ->
@@ -173,47 +157,29 @@ start(_Config) ->
     ok = diameter:start().
 
 start_services(_Config) ->
-    S = [server(N, ?DICT_COMMON) || N <- [?SERVER1,
-                                          ?SERVER2,
-                                          ?SERVER3,
-                                          ?SERVER4]],
-    R = [server(N, ?DICT_RELAY) || N <- [?RELAY1, ?RELAY2]],
+    [S1,S2,S3,S4] = [server(N, ?DICT_COMMON) || N <- [?SERVER1,
+                                                      ?SERVER2,
+                                                      ?SERVER3,
+                                                      ?SERVER4]],
+    [R1,R2] = [server(N, ?DICT_RELAY) || N <- [?RELAY1, ?RELAY2]],
 
     ok = diameter:start_service(?CLIENT, ?SERVICE(?CLIENT, ?DICT_COMMON)),
 
-    {save_config, S ++ R}.
+    {save_config, [{?RELAY1, [S1,S2,R2]},
+                   {?RELAY2, [S3,S4]},
+                   {?CLIENT, [R1,R2]}]}.
 
 connect(Config) ->
-    {_, [S1,S2,S3,S4,R1,R2] = SR} = proplists:get_value(saved_config, Config),
-
-    true = diameter:subscribe(?RELAY1),
-    true = diameter:subscribe(?RELAY2),
-    true = diameter:subscribe(?CLIENT),
+    {_, Conns} = proplists:get_value(saved_config, Config),
 
-    [R1S1,R1S2] = connect(?RELAY1, [S1,S2]),
-    [R2S3,R2S4] = connect(?RELAY2, [S3,S4]),
-    [CR1,CR2]   = connect(?CLIENT, [R1,R2]),
+    ?util:write_priv(Config,
+                     "cfg",
+                     lists:flatmap(fun({CN,Ss}) -> connect(CN, Ss) end,
+                                   Conns)).
 
-    R1R2 = connect(?RELAY1, R2),
-
-    ?util:write_priv(Config, "cfg", SR ++ [R1S1,R1S2,R2S3,R2S4,CR1,CR2,R1R2]).
-
-%% Remove the client transports and expect the corresponding server
-%% transport to go down.
 disconnect(Config) ->
-    [S1,S2,S3,S4,R1,R2,R1S1,R1S2,R2S3,R2S4,CR1,CR2,R1R2]
-        = ?util:read_priv(Config, "cfg"),
-
-    [?CLIENT | Svcs] = ?SERVICES,
-    [] = [{S,T} || S <- Svcs, T <- [diameter:subscribe(S)], T /= true],
-
-    disconnect(?RELAY1, S1, R1S1),
-    disconnect(?RELAY1, S2, R1S2),
-    disconnect(?RELAY2, S3, R2S3),
-    disconnect(?RELAY2, S4, R2S4),
-    disconnect(?CLIENT, R1, CR1),
-    disconnect(?CLIENT, R2, CR2),
-    disconnect(?RELAY1, R2, R1R2).
+    lists:foreach(fun({{CN,CR},{SN,SR}}) -> ?util:disconnect(CN,CR,SN,SR) end,
+                  ?util:read_priv(Config, "cfg")).
 
 stop_services(_Config) ->
     [] = [{H,T} || H <- ?SERVICES,
@@ -223,6 +189,15 @@ stop_services(_Config) ->
 stop(_Config) ->
     ok = diameter:stop().
 
+%% ----------------------------------------
+
+server(Name, Dict) ->
+    ok = diameter:start_service(Name, ?SERVICE(Name, Dict)),
+    {Name, ?util:listen(Name, tcp)}.
+
+connect(Name, Refs) ->
+    [{{Name, ?util:connect(Name, tcp, LRef)}, T} || {_, LRef} = T <- Refs].
+
 %% ===========================================================================
 %% traffic testcases
 
@@ -266,30 +241,6 @@ send_timeout(Tmo) ->
 realm(Host) ->
     tl(lists:dropwhile(fun(C) -> C /= $. end, Host)).
 
-server(Host, Dict) ->
-    ok = diameter:start_service(Host, ?SERVICE(Host, Dict)),
-    {ok, LRef} = diameter:add_transport(Host, ?LISTEN),
-    {LRef, hd([_] = ?util:lport(tcp, LRef, 20))}.
-
-connect(Host, {_LRef, PortNr}) ->
-    {ok, Ref} = diameter:add_transport(Host, ?CONNECT(PortNr)),
-    ok = receive
-             #diameter_event{service = Host,
-                             info = {up, Ref, _, _, #diameter_packet{}}} ->
-                 ok
-         after 2000 ->
-                 false
-         end,
-    Ref;
-connect(Host, Ports) ->
-    [connect(Host, P) || P <- Ports].
-
-disconnect(Client, {LRef, _PortNr}, CRef) ->
-    ok = diameter:remove_transport(Client, CRef),
-    ok = receive #diameter_event{info = {down, LRef, _, _}} -> ok
-         after 2000 -> false
-         end.
-
 call(Server) ->
     Realm = realm(Server),
     Req = ['STR', {'Destination-Realm', Realm},
@@ -325,7 +276,7 @@ peer_down(_SvcName, _Peer, State) ->
 pick_peer([Peer | _], _, Svc, _State)
   when Svc == ?RELAY1;
        Svc == ?RELAY2;
-       Svc == ?CLIENT->
+       Svc == ?CLIENT ->
     {ok, Peer}.
 
 %% prepare_request/3
diff --git a/lib/diameter/test/diameter_traffic_SUITE.erl b/lib/diameter/test/diameter_traffic_SUITE.erl
index d85539edce..6704f24532 100644
--- a/lib/diameter/test/diameter_traffic_SUITE.erl
+++ b/lib/diameter/test/diameter_traffic_SUITE.erl
@@ -127,19 +127,6 @@
                         {module, ?MODULE},
                         {answer_errors, callback}]}]).
 
-%% Config for diameter:add_transport/2. In the listening case, listen
-%% on a free port that we then lookup using the implementation detail
-%% that diameter_tcp registers the port with diameter_reg.
--define(CONNECT(PortNr),
-        {connect, [{transport_module, diameter_tcp},
-                   {transport_config, [{raddr, ?ADDR},
-                                       {rport, PortNr},
-                                       {ip, ?ADDR},
-                                       {port, 0}]}]}).
--define(LISTEN,
-        {listen, [{transport_module, diameter_tcp},
-                  {transport_config, [{ip, ?ADDR}, {port, 0}]}]}).
-
 -define(SUCCESS,
         ?'DIAMETER_BASE_RESULT-CODE_DIAMETER_SUCCESS').
 -define(COMMAND_UNSUPPORTED,
@@ -264,27 +251,13 @@ start_services(_Config) ->
     ok = diameter:start_service(?CLIENT, ?SERVICE(?CLIENT)).
 
 add_transports(Config) ->
-    {ok, LRef} = diameter:add_transport(?SERVER, ?LISTEN),
-    true = diameter:subscribe(?CLIENT),
-    [PortNr] = ?util:lport(tcp, LRef, 20),
-    {ok, CRef} = diameter:add_transport(?CLIENT, ?CONNECT(PortNr)),
-    {up, CRef, _Peer, _Cfg, #diameter_packet{}}
-        = receive #diameter_event{service = ?CLIENT, info = I} -> I
-          after 2000 -> false
-          end,
-    true = diameter:unsubscribe(?CLIENT),
+    LRef = ?util:listen(?SERVER, tcp),
+    CRef = ?util:connect(?CLIENT, tcp, LRef),
     ?util:write_priv(Config, "transport", {LRef, CRef}).
 
-%% Remove the client transport and expect the server transport to
-%% go down.
 remove_transports(Config) ->
     {LRef, CRef} = ?util:read_priv(Config, "transport"),
-    true = diameter:subscribe(?SERVER),
-    ok = diameter:remove_transport(?CLIENT, CRef),
-    {down, LRef, _, _}
-        = receive #diameter_event{service = ?SERVER, info = I} -> I
-          after 2000 -> false
-          end.
+    ?util:disconnect(?CLIENT, CRef, ?SERVER, LRef).
 
 stop_services(_Config) ->
     ok = diameter:stop_service(?CLIENT),
diff --git a/lib/diameter/test/diameter_util.erl b/lib/diameter/test/diameter_util.erl
index 3574094417..3fe8ea5363 100644
--- a/lib/diameter/test/diameter_util.erl
+++ b/lib/diameter/test/diameter_util.erl
@@ -33,9 +33,9 @@
 %% diameter-specific
 -export([lport/2,
          lport/3,
-         service/1,
          listen/2,
-         connect/3]).
+         connect/3,
+         disconnect/4]).
 
 %% common_test-specific
 -export([write_priv/3,
@@ -254,27 +254,12 @@ lp(M, Ref, T) ->
     end.
 
 %% ---------------------------------------------------------------------------
-%% service/1
-%%
-%% Start a new diameter service and return its generated name.
-
-service(Opts) ->
-    Name = make_ref(),
-
-    case diameter:start_service(Name, Opts) of
-        ok ->
-            {ok, Name};
-        {error, _} = No ->
-            No
-    end.
-
-%% ---------------------------------------------------------------------------
 %% listen/2
 %%
 %% Add a listening transport on the loopback address and a free port.
 
 listen(SvcName, Prot) ->
-    add_transport(SvcName, {listen, transport(Prot, listen)}).
+    add_transport(SvcName, {listen, opts(Prot, listen)}).
 
 %% ---------------------------------------------------------------------------
 %% connect/3
@@ -282,15 +267,31 @@ listen(SvcName, Prot) ->
 %% Add a connecting transport on and connect to a listening transport
 %% with the specified reference.
 
-connect(SvcName, Prot, LRef) ->
-    [PortNr] = lport(Prot, LRef),
-    Ref = add_transport(SvcName, {connect, transport(Prot, PortNr)}),
-    true = diameter:subscribe(SvcName),
-    receive
-        {diameter_event, SvcName, {up, Ref, _, _, _}} -> Ref
-    after 2000 ->
-            error({up, SvcName, Prot, PortNr})
-    end.
+connect(Client, Prot, LRef) ->
+    [PortNr] = lport(Prot, LRef, 20),
+    Ref = add_transport(Client, {connect, opts(Prot, PortNr)}),
+    true = diameter:subscribe(Client),
+    ok = receive
+             {diameter_event, Client, {up, Ref, _, _, _}} -> ok
+         after 2000 ->
+                 {Client, Prot, PortNr, process_info(self(), messages)}
+         end,
+    Ref.
+
+%% ---------------------------------------------------------------------------
+%% disconnect/4
+%%
+%% Remove the client transport and expect the server transport to go
+%% down.
+
+disconnect(Client, Ref, Server, LRef) ->
+    true = diameter:subscribe(Server),
+    ok = diameter:remove_transport(Client, Ref),
+    ok = receive
+             {diameter_event, Server, {down, LRef, _, _}} -> ok
+         after 2000 ->
+                 {Client, Ref, Server, LRef, process_info(self(), messages)}
+         end.
 
 %% ---------------------------------------------------------------------------
 
@@ -305,14 +306,6 @@ tmod(tcp) ->
 tmod(sctp) ->
     diameter_sctp.
 
-transport(Prot, T) ->
-    {tag(T), opts(Prot, T)}.
-
-tag(listen = T) ->
-    T;
-tag(_PortNr) ->
-    connect.
-
 opts(Prot, T) ->
     [{transport_module, tmod(Prot)},
      {transport_config, [{ip, ?ADDR}, {port, 0} | opts(T)]}].