From 60ee35ea7e269dca3eecc84a68fd1a5feaa64ec2 Mon Sep 17 00:00:00 2001
From: Ayanda Dube <adube14@bloomberg.net>
Date: Tue, 1 Oct 2024 15:16:16 +0100
Subject: [PATCH] QQ tests for force-shrink to current member operations

---
 deps/rabbit/test/quorum_queue_SUITE.erl | 83 ++++++++++++++++++++++++-
 1 file changed, 82 insertions(+), 1 deletion(-)

diff --git a/deps/rabbit/test/quorum_queue_SUITE.erl b/deps/rabbit/test/quorum_queue_SUITE.erl
index 0643842bf511..5c625fa26eb1 100644
--- a/deps/rabbit/test/quorum_queue_SUITE.erl
+++ b/deps/rabbit/test/quorum_queue_SUITE.erl
@@ -92,7 +92,9 @@ groups() ->
                                             format,
                                             add_member_2,
                                             single_active_consumer_priority_take_over,
-                                            single_active_consumer_priority
+                                            single_active_consumer_priority,
+                                            force_shrink_member_to_current_member,
+                                            force_all_queues_shrink_member_to_current_member
                                            ]
                        ++ all_tests()},
                       {cluster_size_5, [], [start_queue,
@@ -1151,6 +1153,85 @@ single_active_consumer_priority(Config) ->
                 rpc:call(Server0, ra, local_query, [RaNameQ3, QueryFun])),
     ok.
 
+force_shrink_member_to_current_member(Config) ->
+    [Server0, Server1, Server2] =
+        rabbit_ct_broker_helpers:get_node_configs(Config, nodename),
+
+    Ch = rabbit_ct_client_helpers:open_channel(Config, Server0),
+    QQ = ?config(queue_name, Config),
+    ?assertEqual({'queue.declare_ok', QQ, 0, 0},
+                 declare(Ch, QQ, [{<<"x-queue-type">>, longstr, <<"quorum">>}])),
+
+    RaName = ra_name(QQ),
+    rabbit_ct_client_helpers:publish(Ch, QQ, 3),
+    wait_for_messages_ready([Server0], RaName, 3),
+
+    {ok, Q0} = rpc:call(Server0, rabbit_amqqueue, lookup, [QQ, <<"/">>]),
+    #{nodes := Nodes0} = amqqueue:get_type_state(Q0),
+    ?assertEqual(3, length(Nodes0)),
+
+    rabbit_ct_broker_helpers:rpc(Config, 0, rabbit_quorum_queue,
+        force_shrink_member_to_current_member, [<<"/">>, QQ]),
+
+    wait_for_messages_ready([Server0], RaName, 3),
+
+    {ok, Q1} = rpc:call(Server0, rabbit_amqqueue, lookup, [QQ, <<"/">>]),
+    #{nodes := Nodes1} = amqqueue:get_type_state(Q1),
+    ?assertEqual(1, length(Nodes1)),
+
+    %% grow queues back to all nodes
+    [rpc:call(Server0, rabbit_quorum_queue, grow, [S, <<"/">>, <<".*">>, all]) || S <- [Server1, Server2]],
+
+    wait_for_messages_ready([Server0], RaName, 3),
+    {ok, Q2} = rpc:call(Server0, rabbit_amqqueue, lookup, [QQ, <<"/">>]),
+    #{nodes := Nodes2} = amqqueue:get_type_state(Q2),
+    ?assertEqual(3, length(Nodes2)).
+
+force_all_queues_shrink_member_to_current_member(Config) ->
+    [Server0, Server1, Server2] =
+        rabbit_ct_broker_helpers:get_node_configs(Config, nodename),
+
+    Ch = rabbit_ct_client_helpers:open_channel(Config, Server0),
+    QQ = ?config(queue_name, Config),
+    AQ = ?config(alt_queue_name, Config),
+    ?assertEqual({'queue.declare_ok', QQ, 0, 0},
+                 declare(Ch, QQ, [{<<"x-queue-type">>, longstr, <<"quorum">>}])),
+    ?assertEqual({'queue.declare_ok', AQ, 0, 0},
+                 declare(Ch, AQ, [{<<"x-queue-type">>, longstr, <<"quorum">>}])),
+
+    QQs = [QQ, AQ],
+
+    [begin
+        RaName = ra_name(Q),
+        rabbit_ct_client_helpers:publish(Ch, Q, 3),
+        wait_for_messages_ready([Server0], RaName, 3),
+        {ok, Q0} = rpc:call(Server0, rabbit_amqqueue, lookup, [Q, <<"/">>]),
+        #{nodes := Nodes0} = amqqueue:get_type_state(Q0),
+        ?assertEqual(3, length(Nodes0))
+    end || Q <- QQs],
+
+    rabbit_ct_broker_helpers:rpc(Config, 0, rabbit_quorum_queue,
+        force_all_queues_shrink_member_to_current_member, []),
+
+    [begin
+        RaName = ra_name(Q),
+        wait_for_messages_ready([Server0], RaName, 3),
+        {ok, Q0} = rpc:call(Server0, rabbit_amqqueue, lookup, [Q, <<"/">>]),
+        #{nodes := Nodes0} = amqqueue:get_type_state(Q0),
+        ?assertEqual(1, length(Nodes0))
+    end || Q <- QQs],
+
+    %% grow queues back to all nodes
+    [rpc:call(Server0, rabbit_quorum_queue, grow, [S, <<"/">>, <<".*">>, all]) || S <- [Server1, Server2]],
+
+    [begin
+        RaName = ra_name(Q),
+        wait_for_messages_ready([Server0], RaName, 3),
+        {ok, Q0} = rpc:call(Server0, rabbit_amqqueue, lookup, [Q, <<"/">>]),
+        #{nodes := Nodes0} = amqqueue:get_type_state(Q0),
+        ?assertEqual(3, length(Nodes0))
+    end || Q <- QQs].
+
 priority_queue_fifo(Config) ->
     %% testing: if hi priority messages are published before lo priority
     %% messages they are always consumed first (fifo)