Restart stream until it is on expected node in test

Author: acogoluegnes

deps/rabbitmq_stream/test/rabbit_stream_partitions_SUITE.erl

@@ -276,14 +276,15 @@ super_stream_sac_consumer_should_get_disconnected_on_network_partition(Config) -
     CL = coordinator_leader(Config),
 
     Ss = rabbit_data_coercion:to_binary(?FUNCTION_NAME),
-    [Partition, _, _] = init_super_stream(Config, CL, Ss),
+
+    Partition = init_super_stream(Config, CL, Ss, 1, CL),
     [L, F1, F2] = topology(Config, Partition),
 
     wait_for_coordinator_ready(Config),
 
     %% we expect the stream leader and the coordinator leader to be on the same node
     %% another node will be isolated
-    ?assertEqual(L#node.name, coordinator_leader(Config)),
+    ?assertEqual(L#node.name, CL),
 
     {ok, So0, C0_00} = stream_test_utils:connect(L#node.stream_port),
     {ok, So1, C1_00} = stream_test_utils:connect(F1#node.stream_port),
@@ -362,15 +363,16 @@ super_stream_sac_consumer_should_get_disconnected_on_network_partition(Config) -
 super_stream_sac_consumer_should_get_disconnected_on_coord_leader_network_partition(Config) ->
     init_coordinator(Config),
     CL = coordinator_leader(Config),
+    [CF1, _] = all_nodes(Config) -- [CL],
     Ss = rabbit_data_coercion:to_binary(?FUNCTION_NAME),
-    [_, Partition, _] = init_super_stream(Config, CL, Ss),
+    Partition = init_super_stream(Config, CL, Ss, 2, CF1),
     [L, F1, F2] = topology(Config, Partition),
 
     wait_for_coordinator_ready(Config),
 
     %% check stream leader and coordinator are not on the same node
     %% the coordinator leader node will be isolated
-    ?assertNotEqual(L#node.name, coordinator_leader(Config)),
+    ?assertNotEqual(L#node.name, CL),
 
     {ok, So0, C0_00} = stream_test_utils:connect(L#node.stream_port),
     {ok, So1, C1_00} = stream_test_utils:connect(F1#node.stream_port),
@@ -533,7 +535,7 @@ delete_stream(Port, St) ->
     {ok, C1} = stream_test_utils:delete_stream(S, C0, St),
     {ok, _} = stream_test_utils:close(S, C1).
 
-init_super_stream(Config, Node, Ss) ->
+init_super_stream(Config, Node, Ss, PartitionIndex, ExpectedNode) ->
     {ok, S, C0} = stream_test_utils:connect(Config, Node),
     NC = node_count(Config),
     Partitions = [unicode:characters_to_binary([Ss, <<"-">>, integer_to_binary(N)])
@@ -544,10 +546,46 @@ init_super_stream(Config, Node, Ss) ->
     {Cmd1, C1} = receive_commands(S, C0),
     ?assertMatch({response, ?CORR_ID, {create_super_stream, ?RESPONSE_CODE_OK}},
                  Cmd1),
-    Partition = lists:nth(2, Partitions),
-    wait_for_members(S, C1, Partition, NC),
+    [wait_for_members(S, C1, P, NC) || P <- Partitions],
+    Partition = lists:nth(PartitionIndex, Partitions),
+    [#node{name = LN} | _] = topology(Config, Partition),
+    P = case LN of
+            ExpectedNode ->
+                Partition;
+            _ ->
+                enforce_stream_leader_on_node(Config, S, C1,
+                                              Partitions, Partition,
+                                              ExpectedNode, 10)
+        end,
     {ok, _} = stream_test_utils:close(S, C1),
-    Partitions.
+    P.
+
+
+enforce_stream_leader_on_node(_, _, _, _, _, _, 0) ->
+    ct:fail("could not create super stream partition on chosen node");
+enforce_stream_leader_on_node(Config, S, C,
+                              Partitions, Partition, Node, Count) ->
+    CL = coordinator_leader(Config),
+    NC = node_count(Config),
+    [begin
+         case P of
+             Partition ->
+                 restart_stream(Config, CL, P, Node);
+             _ ->
+                 restart_stream(Config, CL, P, undefined)
+         end,
+         wait_for_members(S, C, P, NC)
+     end || P <- Partitions],
+    [#node{name = LN} | _] = topology(Config, Partition),
+    case LN of
+        Node ->
+            Partition;
+        _ ->
+            timer:sleep(500),
+            enforce_stream_leader_on_node(Config, S, C,
+                                          Partitions, Partition, Node,
+                                          Count - 1)
+    end.
 
 delete_super_stream(Port, Ss) ->
     {ok, S, C0} = stream_test_utils:connect(Port),
@@ -631,6 +669,13 @@ coordinator_leader(Config) ->
             undefined
     end.
 
+restart_stream(Config, Node, S, undefined) ->
+    rpc(Config, Node, rabbit_stream_queue, restart_stream, [<<"/">>, S, #{}]);
+restart_stream(Config, Node, S, Leader) ->
+    Opts = #{preferred_leader_node => Leader},
+    rpc(Config, Node, rabbit_stream_queue, restart_stream, [<<"/">>, S, Opts]).
+
+
 rpc(Config, M, F, A) ->
     rpc(Config, 0, M, F, A).