Add e2e for praefect set-replication-factor

qa/qa/service/praefect_manager.rb

@@ -418,6 +418,30 @@ module QA
         shell "docker exec #{@praefect} bash -c '#{cmd}'"
       end
 
+      # set_replication_factor assigns or unassigns random storage nodes as necessary to reach the desired replication factor for a repository
+      def set_replication_factor(relative_path, virtual_storage, factor)
+        cmd = "/opt/gitlab/embedded/bin/praefect -config /var/opt/gitlab/praefect/config.toml set-replication-factor -repository #{relative_path} -virtual-storage #{virtual_storage} -replication-factor #{factor}"
+        shell "docker exec #{@praefect} bash -c '#{cmd}'"
+      end
+
+      # get_replication_storages retrieves a list of currently assigned storages for a repository
+      def get_replication_storages(relative_path, virtual_storage)
+        storage_repositories = []
+        query = "SELECT storage FROM repository_assignments WHERE relative_path='#{relative_path}' AND virtual_storage='#{virtual_storage}';"
+        shell(sql_to_docker_exec_cmd(query)) { |line| storage_repositories << line.strip }
+        # Returned data from query will be in format
+        #    storage
+        #    --------
+        #    gitaly1
+        #    gitaly3
+        #    gitaly2
+        #   (3 rows)
+        #
+
+        # remove 2 header rows and last 2 rows from query response (including blank line)
+        storage_repositories[2..-3]
+      end
+
       def add_repo_to_disk(node, repo_path)
         cmd = "GIT_DIR=. git init --initial-branch=main /var/opt/gitlab/git-data/repositories/#{repo_path}"
         shell "docker exec --user git #{node} bash -c '#{cmd}'"

qa/qa/specs/features/api/3_create/gitaly/praefect_repo_sync_spec.rb

@@ -2,7 +2,7 @@
 
 module QA
   RSpec.describe 'Create' do
-    context 'Praefect repository commands', :orchestrated, :gitaly_cluster, quarantine: { issue: 'https://gitlab.com/gitlab-org/gitlab/-/issues/347415', type: :investigating } do
+    context 'Praefect repository commands', :orchestrated, :gitaly_cluster do
       let(:praefect_manager) { Service::PraefectManager.new }
 
       let(:repo1) { { "relative_path" => "@hashed/repo1.git", "storage" => "gitaly1", "virtual_storage" => "default" } }
@@ -22,7 +22,7 @@ module QA
         praefect_manager.remove_repository_from_praefect_database(repo2["relative_path"])
       end
 
-      it 'allows admin to manage difference between praefect database and disk state', testcase: 'https://gitlab.com/gitlab-org/gitlab/-/quality/test_cases/347606' do
+      it 'allows admin to manage difference between praefect database and disk state', testcase: 'https://gitlab.com/gitlab-org/gitlab/-/quality/test_cases/347606', quarantine: { issue: 'https://gitlab.com/gitlab-org/gitlab/-/issues/347415', type: :investigating } do
         # Some repos are on disk that praefect is not aware of
         untracked_repositories = praefect_manager.list_untracked_repositories
         expect(untracked_repositories).to include(repo1)
@@ -59,6 +59,18 @@ module QA
         untracked_repositories = praefect_manager.list_untracked_repositories
         expect(untracked_repositories).not_to include(repo1)
       end
+
+      it 'allows admin to control the number of replicas of data', testcase: 'https://gitlab.com/gitlab-org/quality/testcases/-/quality/test_cases/2434' do
+        praefect_manager.track_repository_in_praefect(repo1['relative_path'], repo1['storage'], repo1['virtual_storage'])
+
+        praefect_manager.set_replication_factor(repo1['relative_path'], repo1['virtual_storage'], 2)
+        replication_storages = praefect_manager.get_replication_storages(repo1['relative_path'], repo1['virtual_storage'])
+        expect(replication_storages).to have_attributes(size: 2)
+
+        praefect_manager.set_replication_factor(repo1['relative_path'], repo1['virtual_storage'], 3)
+        replication_storages = praefect_manager.get_replication_storages(repo1['relative_path'], repo1['virtual_storage'])
+        expect(replication_storages).to eq(%w(gitaly1 gitaly2 gitaly3))
+      end
     end
   end
 end