Add latest changes from gitlab-org/gitlab@14-9-stable-eev14.9.0-rc42

7 files changed, 616 insertions, 50 deletions

diff --git a/spec/tasks/dev_rake_spec.rb b/spec/tasks/dev_rake_spec.rb
new file mode 100644
index 00000000000..7bc27d2732c
--- /dev/null
+++ b/spec/tasks/dev_rake_spec.rb
@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+require 'rake_helper'
+
+RSpec.describe 'dev rake tasks' do
+  before do
+    Rake.application.rake_require 'tasks/gitlab/setup'
+    Rake.application.rake_require 'tasks/gitlab/shell'
+    Rake.application.rake_require 'tasks/dev'
+  end
+
+  describe 'setup' do
+    subject(:setup_task) { run_rake_task('dev:setup') }
+
+    let(:connections) { Gitlab::Database.database_base_models.values.map(&:connection) }
+
+    it 'sets up the development environment', :aggregate_failures do
+      expect(Rake::Task['gitlab:setup']).to receive(:invoke)
+
+      expect(connections).to all(receive(:execute).with('ANALYZE'))
+
+      expect(Rake::Task['gitlab:shell:setup']).to receive(:invoke)
+
+      setup_task
+    end
+  end
+
+  describe 'load' do
+    subject(:load_task) { run_rake_task('dev:load') }
+
+    it 'eager loads the application', :aggregate_failures do
+      expect(Rails.configuration).to receive(:eager_load=).with(true)
+      expect(Rails.application).to receive(:eager_load!)
+
+      load_task
+    end
+  end
+end
diff --git a/spec/tasks/gitlab/background_migrations_rake_spec.rb b/spec/tasks/gitlab/background_migrations_rake_spec.rb
index 079b4d3aea8..98920df71ee 100644
--- a/spec/tasks/gitlab/background_migrations_rake_spec.rb
+++ b/spec/tasks/gitlab/background_migrations_rake_spec.rb
@@ -10,6 +10,16 @@ RSpec.describe 'gitlab:background_migrations namespace rake tasks' do
   describe 'finalize' do
     subject(:finalize_task) { run_rake_task('gitlab:background_migrations:finalize', *arguments) }
 
+    let(:connection) { double(:connection) }
+    let(:main_model) { double(:model, connection: connection) }
+    let(:base_models) { { main: main_model } }
+    let(:databases) { [Gitlab::Database::MAIN_DATABASE_NAME] }
+
+    before do
+      allow(Gitlab::Database).to receive(:database_base_models).and_return(base_models)
+      allow(Gitlab::Database).to receive(:db_config_names).and_return(databases)
+    end
+
     context 'without the proper arguments' do
       let(:arguments) { %w[CopyColumnUsingBackgroundMigrationJob events id] }
 
@@ -26,24 +36,135 @@ RSpec.describe 'gitlab:background_migrations namespace rake tasks' do
 
       it 'finalizes the matching migration' do
         expect(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner).to receive(:finalize)
-          .with('CopyColumnUsingBackgroundMigrationJob', 'events', 'id', [%w[id1 id2]])
+          .with('CopyColumnUsingBackgroundMigrationJob', 'events', 'id', [%w[id1 id2]], connection: connection)
 
         expect { finalize_task }.to output(/Done/).to_stdout
       end
     end
+
+    context 'when multiple database feature is enabled' do
+      subject(:finalize_task) { run_rake_task("gitlab:background_migrations:finalize:#{ci_database_name}", *arguments) }
+
+      let(:ci_database_name) { Gitlab::Database::CI_DATABASE_NAME }
+      let(:ci_model) { double(:model, connection: connection) }
+      let(:base_models) { { 'main' => main_model, 'ci' => ci_model } }
+      let(:databases) { [Gitlab::Database::MAIN_DATABASE_NAME, ci_database_name] }
+
+      before do
+        skip_if_multiple_databases_not_setup
+
+        allow(Gitlab::Database).to receive(:database_base_models).and_return(base_models)
+      end
+
+      it 'ignores geo' do
+        expect { run_rake_task('gitlab:background_migrations:finalize:geo}') }
+          .to raise_error(RuntimeError).with_message(/Don't know how to build task/)
+      end
+
+      context 'without the proper arguments' do
+        let(:arguments) { %w[CopyColumnUsingBackgroundMigrationJob events id] }
+
+        it 'exits without finalizing the migration' do
+          expect(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner).not_to receive(:finalize)
+
+          expect { finalize_task }.to output(/Must specify job_arguments as an argument/).to_stdout
+            .and raise_error(SystemExit) { |error| expect(error.status).to eq(1) }
+        end
+      end
+
+      context 'with the proper arguments' do
+        let(:arguments) { %w[CopyColumnUsingBackgroundMigrationJob events id [["id1"\,"id2"]]] }
+
+        it 'finalizes the matching migration' do
+          expect(Gitlab::Database::BackgroundMigration::BatchedMigrationRunner).to receive(:finalize)
+            .with('CopyColumnUsingBackgroundMigrationJob', 'events', 'id', [%w[id1 id2]], connection: connection)
+
+          expect { finalize_task }.to output(/Done/).to_stdout
+        end
+      end
+
+      context 'when database name is not passed' do
+        it 'aborts the rake task' do
+          expect { run_rake_task('gitlab:background_migrations:finalize') }.to output(/Please specify the database/).to_stdout
+            .and raise_error(SystemExit) { |error| expect(error.status).to eq(1) }
+        end
+      end
+    end
   end
 
   describe 'status' do
     subject(:status_task) { run_rake_task('gitlab:background_migrations:status') }
 
+    let(:migration1) { create(:batched_background_migration, :finished, job_arguments: [%w[id1 id2]]) }
+    let(:migration2) { create(:batched_background_migration, :failed, job_arguments: []) }
+
+    let(:main_database_name) { Gitlab::Database::MAIN_DATABASE_NAME }
+    let(:model) { Gitlab::Database.database_base_models[main_database_name] }
+    let(:connection) { double(:connection) }
+    let(:base_models) { { 'main' => model } }
+
+    around do |example|
+      Gitlab::Database::SharedModel.using_connection(model.connection) do
+        example.run
+      end
+    end
+
     it 'outputs the status of background migrations' do
-      migration1 = create(:batched_background_migration, :finished, job_arguments: [%w[id1 id2]])
-      migration2 = create(:batched_background_migration, :failed, job_arguments: [])
+      allow(Gitlab::Database).to receive(:database_base_models).and_return(base_models)
 
       expect { status_task }.to output(<<~OUTPUT).to_stdout
+        Database: #{main_database_name}
         finished   | #{migration1.job_class_name},#{migration1.table_name},#{migration1.column_name},[["id1","id2"]]
         failed     | #{migration2.job_class_name},#{migration2.table_name},#{migration2.column_name},[]
       OUTPUT
     end
+
+    context 'when multiple database feature is enabled' do
+      before do
+        skip_if_multiple_databases_not_setup
+      end
+
+      context 'with a single database' do
+        subject(:status_task) { run_rake_task("gitlab:background_migrations:status:#{main_database_name}") }
+
+        it 'outputs the status of background migrations' do
+          expect { status_task }.to output(<<~OUTPUT).to_stdout
+            Database: #{main_database_name}
+            finished   | #{migration1.job_class_name},#{migration1.table_name},#{migration1.column_name},[["id1","id2"]]
+            failed     | #{migration2.job_class_name},#{migration2.table_name},#{migration2.column_name},[]
+          OUTPUT
+        end
+
+        it 'ignores geo' do
+          expect { run_rake_task('gitlab:background_migrations:status:geo') }
+            .to raise_error(RuntimeError).with_message(/Don't know how to build task/)
+        end
+      end
+
+      context 'with multiple databases' do
+        subject(:status_task) { run_rake_task('gitlab:background_migrations:status') }
+
+        let(:base_models) { { 'main' => main_model, 'ci' => ci_model } }
+        let(:main_model) { double(:model, connection: connection) }
+        let(:ci_model) { double(:model, connection: connection) }
+
+        it 'outputs the status for each database' do
+          allow(Gitlab::Database).to receive(:database_base_models).and_return(base_models)
+
+          expect(Gitlab::Database::SharedModel).to receive(:using_connection).with(main_model.connection).and_yield
+          expect(Gitlab::Database::BackgroundMigration::BatchedMigration).to receive(:find_each).and_yield(migration1)
+
+          expect(Gitlab::Database::SharedModel).to receive(:using_connection).with(ci_model.connection).and_yield
+          expect(Gitlab::Database::BackgroundMigration::BatchedMigration).to receive(:find_each).and_yield(migration2)
+
+          expect { status_task }.to output(<<~OUTPUT).to_stdout
+            Database: main
+            finished   | #{migration1.job_class_name},#{migration1.table_name},#{migration1.column_name},[["id1","id2"]]
+            Database: ci
+            failed     | #{migration2.job_class_name},#{migration2.table_name},#{migration2.column_name},[]
+          OUTPUT
+        end
+      end
+    end
   end
 end
diff --git a/spec/tasks/gitlab/backup_rake_spec.rb b/spec/tasks/gitlab/backup_rake_spec.rb
index e9aa8cbb991..df9f2a0d3bb 100644
--- a/spec/tasks/gitlab/backup_rake_spec.rb
+++ b/spec/tasks/gitlab/backup_rake_spec.rb
@@ -72,7 +72,6 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
         before do
           allow(YAML).to receive(:load_file)
             .and_return({ gitlab_version: gitlab_version })
-          expect(Rake::Task['gitlab:db:drop_tables']).to receive(:invoke)
           expect_next_instance_of(::Backup::Manager) do |instance|
             backup_types.each do |subtask|
               expect(instance).to receive(:run_restore_task).with(subtask).ordered
@@ -85,10 +84,6 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
         it 'invokes restoration on match' do
           expect { run_rake_task('gitlab:backup:restore') }.to output.to_stdout_from_any_process
         end
-
-        it 'prints timestamps on messages' do
-          expect { run_rake_task('gitlab:backup:restore') }.to output(/.*\d{4}-\d{2}-\d{2}\s\d{2}:\d{2}:\d{2}\s[-+]\d{4}\s--\s.*/).to_stdout_from_any_process
-        end
       end
     end
 
@@ -131,8 +126,6 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
         allow(YAML).to receive(:load_file)
           .and_return({ gitlab_version: Gitlab::VERSION })
 
-        expect(Rake::Task['gitlab:db:drop_tables']).to receive(:invoke)
-
         expect_next_instance_of(::Backup::Manager) do |instance|
           backup_types.each do |subtask|
             expect(instance).to receive(:run_restore_task).with(subtask).ordered
@@ -183,8 +176,8 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
 
           expect(exit_status).to eq(0)
           expect(tar_contents).to match(user_backup_path)
-          expect(tar_contents).to match("#{user_backup_path}/custom_hooks.tar")
-          expect(tar_contents).to match("#{user_backup_path}.bundle")
+          expect(tar_contents).to match("#{user_backup_path}/.+/001.custom_hooks.tar")
+          expect(tar_contents).to match("#{user_backup_path}/.+/001.bundle")
         end
 
         it 'restores files correctly' do
@@ -367,14 +360,14 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
           expect(exit_status).to eq(0)
 
           [
-            "#{project_a.disk_path}.bundle",
-            "#{project_a.disk_path}.wiki.bundle",
-            "#{project_a.disk_path}.design.bundle",
-            "#{project_b.disk_path}.bundle",
-            "#{project_snippet_a.disk_path}.bundle",
-            "#{project_snippet_b.disk_path}.bundle"
+            "#{project_a.disk_path}/.+/001.bundle",
+            "#{project_a.disk_path}.wiki/.+/001.bundle",
+            "#{project_a.disk_path}.design/.+/001.bundle",
+            "#{project_b.disk_path}/.+/001.bundle",
+            "#{project_snippet_a.disk_path}/.+/001.bundle",
+            "#{project_snippet_b.disk_path}/.+/001.bundle"
           ].each do |repo_name|
-            expect(tar_lines.grep(/#{repo_name}/).size).to eq 1
+            expect(tar_lines).to include(a_string_matching(repo_name))
           end
         end
 
@@ -435,7 +428,7 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
         expect(::Backup::Repositories).to receive(:new)
           .with(anything, strategy: anything, max_concurrency: 5, max_storage_concurrency: 2)
           .and_call_original
-        expect(::Backup::GitalyBackup).to receive(:new).with(anything, max_parallelism: 5, storage_parallelism: 2).and_call_original
+        expect(::Backup::GitalyBackup).to receive(:new).with(anything, max_parallelism: 5, storage_parallelism: 2, incremental: false).and_call_original
 
         expect { run_rake_task('gitlab:backup:create') }.to output.to_stdout_from_any_process
       end
@@ -486,7 +479,6 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
       allow(Rake::Task['gitlab:shell:setup'])
         .to receive(:invoke).and_return(true)
 
-      expect(Rake::Task['gitlab:db:drop_tables']).to receive :invoke
       expect_next_instance_of(::Backup::Manager) do |instance|
         (backup_types - %w{repositories uploads}).each do |subtask|
           expect(instance).to receive(:run_restore_task).with(subtask).ordered
@@ -531,7 +523,6 @@ RSpec.describe 'gitlab:app namespace rake task', :delete do
       allow(Rake::Task['gitlab:shell:setup'])
         .to receive(:invoke).and_return(true)
 
-      expect(Rake::Task['gitlab:db:drop_tables']).to receive :invoke
       expect_next_instance_of(::Backup::Manager) do |instance|
         backup_types.each do |subtask|
           expect(instance).to receive(:run_restore_task).with(subtask).ordered
diff --git a/spec/tasks/gitlab/db_rake_spec.rb b/spec/tasks/gitlab/db_rake_spec.rb
index c3fd8135ae0..8d3ec7b1ee2 100644
--- a/spec/tasks/gitlab/db_rake_spec.rb
+++ b/spec/tasks/gitlab/db_rake_spec.rb
@@ -20,6 +20,14 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
     allow(Rake::Task['db:seed_fu']).to receive(:invoke).and_return(true)
   end
 
+  describe 'clear_all_connections' do
+    it 'calls clear_all_connections!' do
+      expect(ActiveRecord::Base).to receive(:clear_all_connections!)
+
+      run_rake_task('gitlab:db:clear_all_connections')
+    end
+  end
+
   describe 'mark_migration_complete' do
     context 'with a single database' do
       let(:main_model) { ActiveRecord::Base }
@@ -253,45 +261,78 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
   end
 
   describe 'drop_tables' do
-    subject { run_rake_task('gitlab:db:drop_tables') }
-
-    let(:tables) { %w(one two) }
+    let(:tables) { %w(one two schema_migrations) }
     let(:views) { %w(three four) }
-    let(:connection) { ActiveRecord::Base.connection }
+    let(:schemas) { Gitlab::Database::EXTRA_SCHEMAS }
 
-    before do
-      allow(connection).to receive(:execute).and_return(nil)
+    context 'with a single database' do
+      let(:connection) { ActiveRecord::Base.connection }
+
+      before do
+        skip_if_multiple_databases_are_setup
+
+        allow(connection).to receive(:execute).and_return(nil)
+
+        allow(connection).to receive(:tables).and_return(tables)
+        allow(connection).to receive(:views).and_return(views)
+      end
+
+      it 'drops all objects for the database', :aggregate_failures do
+        expect_objects_to_be_dropped(connection)
 
-      allow(connection).to receive(:tables).and_return(tables)
-      allow(connection).to receive(:views).and_return(views)
+        run_rake_task('gitlab:db:drop_tables')
+      end
     end
 
-    it 'drops all tables, except schema_migrations' do
-      expect(connection).to receive(:execute).with('DROP TABLE IF EXISTS "one" CASCADE')
-      expect(connection).to receive(:execute).with('DROP TABLE IF EXISTS "two" CASCADE')
+    context 'with multiple databases', :aggregate_failures do
+      let(:main_model) { double(:model, connection: double(:connection, tables: tables, views: views)) }
+      let(:ci_model) { double(:model, connection: double(:connection, tables: tables, views: views)) }
+      let(:base_models) { { 'main' => main_model, 'ci' => ci_model } }
 
-      subject
+      before do
+        skip_if_multiple_databases_not_setup
+
+        allow(Gitlab::Database).to receive(:database_base_models).and_return(base_models)
+
+        allow(main_model.connection).to receive(:table_exists?).with('schema_migrations').and_return(true)
+        allow(ci_model.connection).to receive(:table_exists?).with('schema_migrations').and_return(true)
+
+        (tables + views + schemas).each do |name|
+          allow(main_model.connection).to receive(:quote_table_name).with(name).and_return("\"#{name}\"")
+          allow(ci_model.connection).to receive(:quote_table_name).with(name).and_return("\"#{name}\"")
+        end
+      end
+
+      it 'drops all objects for all databases', :aggregate_failures do
+        expect_objects_to_be_dropped(main_model.connection)
+        expect_objects_to_be_dropped(ci_model.connection)
+
+        run_rake_task('gitlab:db:drop_tables')
+      end
+
+      context 'when the single database task is used' do
+        it 'drops all objects for the given database', :aggregate_failures do
+          expect_objects_to_be_dropped(main_model.connection)
+
+          expect(ci_model.connection).not_to receive(:execute)
+
+          run_rake_task('gitlab:db:drop_tables:main')
+        end
+      end
     end
 
-    it 'drops all views' do
+    def expect_objects_to_be_dropped(connection)
+      expect(connection).to receive(:execute).with('DROP TABLE IF EXISTS "one" CASCADE')
+      expect(connection).to receive(:execute).with('DROP TABLE IF EXISTS "two" CASCADE')
+
       expect(connection).to receive(:execute).with('DROP VIEW IF EXISTS "three" CASCADE')
       expect(connection).to receive(:execute).with('DROP VIEW IF EXISTS "four" CASCADE')
 
-      subject
-    end
-
-    it 'truncates schema_migrations table' do
       expect(connection).to receive(:execute).with('TRUNCATE schema_migrations')
 
-      subject
-    end
-
-    it 'drops extra schemas' do
       Gitlab::Database::EXTRA_SCHEMAS.each do |schema|
         expect(connection).to receive(:execute).with("DROP SCHEMA IF EXISTS \"#{schema}\" CASCADE")
       end
-
-      subject
     end
   end
 
@@ -422,13 +463,11 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
 
   context 'with multiple databases', :reestablished_active_record_base do
     before do
-      allow(ActiveRecord::Tasks::DatabaseTasks).to receive(:setup_initial_database_yaml).and_return([:main, :geo])
+      skip_if_multiple_databases_not_setup
     end
 
     describe 'db:structure:dump' do
       it 'invokes gitlab:db:clean_structure_sql' do
-        skip unless Gitlab.ee?
-
         expect(Rake::Task['gitlab:db:clean_structure_sql']).to receive(:invoke).twice.and_return(true)
 
         expect { run_rake_task('db:structure:dump:main') }.not_to raise_error
@@ -437,13 +476,33 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
 
     describe 'db:schema:dump' do
       it 'invokes gitlab:db:clean_structure_sql' do
-        skip unless Gitlab.ee?
-
         expect(Rake::Task['gitlab:db:clean_structure_sql']).to receive(:invoke).once.and_return(true)
 
         expect { run_rake_task('db:schema:dump:main') }.not_to raise_error
       end
     end
+
+    describe 'db:migrate' do
+      it 'invokes gitlab:db:create_dynamic_partitions' do
+        expect(Rake::Task['gitlab:db:create_dynamic_partitions']).to receive(:invoke).once.and_return(true)
+
+        expect { run_rake_task('db:migrate:main') }.not_to raise_error
+      end
+    end
+
+    describe 'db:migrate:geo' do
+      it 'does not invoke gitlab:db:create_dynamic_partitions' do
+        skip 'Skipping because geo database is not setup' unless geo_configured?
+
+        expect(Rake::Task['gitlab:db:create_dynamic_partitions']).not_to receive(:invoke)
+
+        expect { run_rake_task('db:migrate:geo') }.not_to raise_error
+      end
+
+      def geo_configured?
+        !!ActiveRecord::Base.configurations.configs_for(env_name: Rails.env, name: 'geo')
+      end
+    end
   end
 
   describe 'gitlab:db:reset_as_non_superuser' do
diff --git a/spec/tasks/gitlab/refresh_project_statistics_build_artifacts_size_rake_spec.rb b/spec/tasks/gitlab/refresh_project_statistics_build_artifacts_size_rake_spec.rb
new file mode 100644
index 00000000000..e57704d0ebe
--- /dev/null
+++ b/spec/tasks/gitlab/refresh_project_statistics_build_artifacts_size_rake_spec.rb
@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+require 'rake_helper'
+
+RSpec.describe 'gitlab:refresh_project_statistics_build_artifacts_size rake task', :silence_stdout do
+  let(:rake_task) { 'gitlab:refresh_project_statistics_build_artifacts_size' }
+
+  describe 'enqueuing build artifacts size statistics refresh for given list of project IDs' do
+    let_it_be(:project_1) { create(:project) }
+    let_it_be(:project_2) { create(:project) }
+    let_it_be(:project_3) { create(:project) }
+
+    let(:string_of_ids) { "#{project_1.id} #{project_2.id} #{project_3.id} 999999" }
+
+    before do
+      Rake.application.rake_require('tasks/gitlab/refresh_project_statistics_build_artifacts_size')
+
+      stub_const("BUILD_ARTIFACTS_SIZE_REFRESH_ENQUEUE_BATCH_SIZE", 2)
+    end
+
+    context 'when given a list of space-separated IDs through STDIN' do
+      before do
+        allow($stdin).to receive(:tty?).and_return(false)
+        allow($stdin).to receive(:read).and_return(string_of_ids)
+      end
+
+      it 'enqueues the projects for refresh' do
+        expect { run_rake_task(rake_task) }.to output(/Done/).to_stdout
+
+        expect(Projects::BuildArtifactsSizeRefresh.all.map(&:project)).to match_array([project_1, project_2, project_3])
+      end
+    end
+
+    context 'when given a list of space-separated IDs through rake argument' do
+      it 'enqueues the projects for refresh' do
+        expect { run_rake_task(rake_task, string_of_ids) }.to output(/Done/).to_stdout
+
+        expect(Projects::BuildArtifactsSizeRefresh.all.map(&:project)).to match_array([project_1, project_2, project_3])
+      end
+    end
+
+    context 'when not given any IDs' do
+      it 'returns an error message' do
+        expect { run_rake_task(rake_task) }.to output(/Please provide a string of space-separated project IDs/).to_stdout
+      end
+    end
+  end
+end
diff --git a/spec/tasks/gitlab/setup_rake_spec.rb b/spec/tasks/gitlab/setup_rake_spec.rb
new file mode 100644
index 00000000000..6e4d5087517
--- /dev/null
+++ b/spec/tasks/gitlab/setup_rake_spec.rb
@@ -0,0 +1,141 @@
+# frozen_string_literal: true
+
+require 'rake_helper'
+
+RSpec.describe 'gitlab:setup namespace rake tasks', :silence_stdout do
+  before do
+    Rake.application.rake_require 'active_record/railties/databases'
+    Rake.application.rake_require 'tasks/seed_fu'
+    Rake.application.rake_require 'tasks/gitlab/setup'
+  end
+
+  describe 'setup' do
+    subject(:setup_task) { run_rake_task('gitlab:setup') }
+
+    let(:storages) do
+      {
+        'name1' => 'some details',
+        'name2' => 'other details'
+      }
+    end
+
+    let(:server_service1) { double(:server_service) }
+    let(:server_service2) { double(:server_service) }
+
+    let(:connections) { Gitlab::Database.database_base_models.values.map(&:connection) }
+
+    before do
+      allow(Gitlab).to receive_message_chain('config.repositories.storages').and_return(storages)
+
+      stub_warn_user_is_not_gitlab
+
+      allow(main_object).to receive(:ask_to_continue)
+    end
+
+    it 'sets up the application', :aggregate_failures do
+      expect_gitaly_connections_to_be_checked
+
+      expect_connections_to_be_terminated
+      expect_database_to_be_setup
+
+      setup_task
+    end
+
+    context 'when an environment variable is set to force execution' do
+      before do
+        stub_env('force', 'yes')
+      end
+
+      it 'sets up the application without prompting the user', :aggregate_failures do
+        expect_gitaly_connections_to_be_checked
+
+        expect(main_object).not_to receive(:ask_to_continue)
+
+        expect_connections_to_be_terminated
+        expect_database_to_be_setup
+
+        setup_task
+      end
+    end
+
+    context 'when the gitaly connection check raises an error' do
+      it 'exits the task without setting up the database', :aggregate_failures do
+        expect(Gitlab::GitalyClient::ServerService).to receive(:new).with('name1').and_return(server_service1)
+        expect(server_service1).to receive(:info).and_raise(GRPC::Unavailable)
+
+        expect_connections_not_to_be_terminated
+        expect_database_not_to_be_setup
+
+        expect { setup_task }.to output(/Failed to connect to Gitaly/).to_stdout
+          .and raise_error(SystemExit) { |error| expect(error.status).to eq(1) }
+      end
+    end
+
+    context 'when the task is aborted' do
+      it 'exits without setting up the database', :aggregate_failures do
+        expect_gitaly_connections_to_be_checked
+
+        expect(main_object).to receive(:ask_to_continue).and_raise(Gitlab::TaskAbortedByUserError)
+
+        expect_connections_not_to_be_terminated
+        expect_database_not_to_be_setup
+
+        expect { setup_task }.to output(/Quitting/).to_stdout
+          .and raise_error(SystemExit) { |error| expect(error.status).to eq(1) }
+      end
+    end
+
+    context 'when in the production environment' do
+      it 'sets up the database without terminating connections', :aggregate_failures do
+        expect_gitaly_connections_to_be_checked
+
+        expect(Rails.env).to receive(:production?).and_return(true)
+
+        expect_connections_not_to_be_terminated
+        expect_database_to_be_setup
+
+        setup_task
+      end
+    end
+
+    context 'when the database is not found when terminating connections' do
+      it 'continues setting up the database', :aggregate_failures do
+        expect_gitaly_connections_to_be_checked
+
+        expect(connections).to all(receive(:execute).and_raise(ActiveRecord::NoDatabaseError))
+
+        expect_database_to_be_setup
+
+        setup_task
+      end
+    end
+
+    def expect_gitaly_connections_to_be_checked
+      expect(Gitlab::GitalyClient::ServerService).to receive(:new).with('name1').and_return(server_service1)
+      expect(server_service1).to receive(:info)
+
+      expect(Gitlab::GitalyClient::ServerService).to receive(:new).with('name2').and_return(server_service2)
+      expect(server_service2).to receive(:info)
+    end
+
+    def expect_connections_to_be_terminated
+      expect(connections).to all(receive(:execute).with(/SELECT pg_terminate_backend/))
+    end
+
+    def expect_connections_not_to_be_terminated
+      connections.each do |connection|
+        expect(connection).not_to receive(:execute)
+      end
+    end
+
+    def expect_database_to_be_setup
+      expect(Rake::Task['db:reset']).to receive(:invoke)
+      expect(Rake::Task['db:seed_fu']).to receive(:invoke)
+    end
+
+    def expect_database_not_to_be_setup
+      expect(Rake::Task['db:reset']).not_to receive(:invoke)
+      expect(Rake::Task['db:seed_fu']).not_to receive(:invoke)
+    end
+  end
+end
diff --git a/spec/tasks/rubocop_rake_spec.rb b/spec/tasks/rubocop_rake_spec.rb
new file mode 100644
index 00000000000..cf7e45aae28
--- /dev/null
+++ b/spec/tasks/rubocop_rake_spec.rb
@@ -0,0 +1,168 @@
+# frozen_string_literal: true
+# rubocop:disable RSpec/VerifiedDoubles
+
+require 'fast_spec_helper'
+require 'rake'
+require 'fileutils'
+
+require_relative '../support/silence_stdout'
+require_relative '../support/helpers/next_instance_of'
+require_relative '../support/helpers/rake_helpers'
+require_relative '../../rubocop/todo_dir'
+
+RSpec.describe 'rubocop rake tasks', :silence_stdout do
+  include RakeHelpers
+
+  before do
+    stub_const('Rails', double(:rails_env))
+    allow(Rails).to receive(:env).and_return(double(production?: false))
+
+    stub_const('ENV', ENV.to_hash.dup)
+
+    Rake.application.rake_require 'tasks/rubocop'
+  end
+
+  describe 'todo:generate', :aggregate_failures do
+    let(:tmp_dir) { Dir.mktmpdir }
+    let(:rubocop_todo_dir) { File.join(tmp_dir, '.rubocop_todo') }
+    let(:todo_dir) { RuboCop::TodoDir.new(rubocop_todo_dir) }
+
+    around do |example|
+      Dir.chdir(tmp_dir) do
+        with_inflections do
+          example.run
+        end
+      end
+    end
+
+    before do
+      allow(RuboCop::TodoDir).to receive(:new).and_return(todo_dir)
+
+      # This Ruby file will trigger the following 3 offenses.
+      File.write('a.rb', <<~RUBY)
+        a+b
+
+      RUBY
+
+      # Mimic GitLab's .rubocop_todo.yml avoids relying on RuboCop's
+      # default.yml configuration.
+      File.write('.rubocop.yml', <<~YAML)
+        <% unless ENV['REVEAL_RUBOCOP_TODO'] == '1' %>
+          <% Dir.glob('.rubocop_todo/**/*.yml').each do |rubocop_todo_yaml| %>
+        - '<%= rubocop_todo_yaml %>'
+          <% end %>
+        - '.rubocop_todo.yml'
+        <% end %>
+
+        AllCops:
+          NewCops: enable # Avoiding RuboCop warnings
+
+        Layout/SpaceAroundOperators:
+          Enabled: true
+
+        Layout/TrailingEmptyLines:
+          Enabled: true
+
+        Lint/Syntax:
+          Enabled: true
+
+        Style/FrozenStringLiteralComment:
+          Enabled: true
+      YAML
+
+      # Required to verify that we are revealing all TODOs via
+      # ENV['REVEAL_RUBOCOP_TODO'] = '1'.
+      # This file can be removed from specs after we've moved all offenses from
+      # .rubocop_todo.yml to .rubocop_todo/**/*.yml.
+      File.write('.rubocop_todo.yml', <<~YAML)
+        # Too many offenses
+        Layout/SpaceAroundOperators:
+          Enabled: false
+      YAML
+
+      # Previous offense now fixed.
+      todo_dir.write('Lint/Syntax', '')
+    end
+
+    after do
+      FileUtils.remove_entry(tmp_dir)
+    end
+
+    context 'without arguments' do
+      let(:run_task) { run_rake_task('rubocop:todo:generate') }
+
+      it 'generates TODOs for all RuboCop rules' do
+        expect { run_task }.to output(<<~OUTPUT).to_stdout
+          Generating RuboCop TODOs with:
+            rubocop --parallel --format RuboCop::Formatter::TodoFormatter
+
+          This might take a while...
+          Written to .rubocop_todo/layout/space_around_operators.yml
+          Written to .rubocop_todo/layout/trailing_empty_lines.yml
+          Written to .rubocop_todo/style/frozen_string_literal_comment.yml
+        OUTPUT
+
+        expect(rubocop_todo_dir_listing).to contain_exactly(
+          'layout/space_around_operators.yml',
+          'layout/trailing_empty_lines.yml',
+          'style/frozen_string_literal_comment.yml'
+        )
+      end
+
+      it 'sets acronyms for inflections' do
+        run_task
+
+        expect(ActiveSupport::Inflector.inflections.acronyms).to include(
+          'rspec' => 'RSpec',
+          'graphql' => 'GraphQL'
+        )
+      end
+    end
+
+    context 'with cop names as arguments' do
+      let(:run_task) do
+        cop_names = %w[
+          Style/FrozenStringLiteralComment Layout/TrailingEmptyLines
+          Lint/Syntax
+        ]
+
+        run_rake_task('rubocop:todo:generate', cop_names)
+      end
+
+      it 'generates TODOs for given RuboCop cops' do
+        expect { run_task }.to output(<<~OUTPUT).to_stdout
+          Generating RuboCop TODOs with:
+            rubocop --parallel --format RuboCop::Formatter::TodoFormatter --only Layout/TrailingEmptyLines,Lint/Syntax,Style/FrozenStringLiteralComment
+
+          This might take a while...
+          Written to .rubocop_todo/layout/trailing_empty_lines.yml
+          Written to .rubocop_todo/style/frozen_string_literal_comment.yml
+        OUTPUT
+
+        expect(rubocop_todo_dir_listing).to contain_exactly(
+          'layout/trailing_empty_lines.yml',
+          'style/frozen_string_literal_comment.yml'
+        )
+      end
+    end
+
+    private
+
+    def rubocop_todo_dir_listing
+      Dir.glob("#{rubocop_todo_dir}/**/*")
+        .select { |path| File.file?(path) }
+        .map { |path| path.delete_prefix("#{rubocop_todo_dir}/") }
+    end
+
+    def with_inflections
+      original = ActiveSupport::Inflector::Inflections.instance_variable_get(:@__instance__)[:en]
+      ActiveSupport::Inflector::Inflections.instance_variable_set(:@__instance__, en: original.dup)
+
+      yield
+    ensure
+      ActiveSupport::Inflector::Inflections.instance_variable_set(:@__instance__, en: original)
+    end
+  end
+end
+
+# rubocop:enable RSpec/VerifiedDoubles