2 files changed, 130 insertions, 30 deletions
diff --git a/spec/lib/gitlab/hashed_storage/migrator_spec.rb b/spec/lib/gitlab/hashed_storage/migrator_spec.rb
index 6154b3e2f76..d03a74ac9eb 100644
--- a/spec/lib/gitlab/hashed_storage/migrator_spec.rb
+++ b/spec/lib/gitlab/hashed_storage/migrator_spec.rb
@@ -1,6 +1,8 @@
+# frozen_string_literal: true
+
 require 'spec_helper'
 
-describe Gitlab::HashedStorage::Migrator do
+describe Gitlab::HashedStorage::Migrator, :sidekiq do
   describe '#bulk_schedule_migration' do
     it 'schedules job to HashedStorage::MigratorWorker' do
       Sidekiq::Testing.fake! do
@@ -182,4 +184,52 @@ describe Gitlab::HashedStorage::Migrator do
       end
     end
   end
+
+  describe 'migration_pending?' do
+    set(:project) { create(:project, :empty_repo) }
+
+    it 'returns true when there are MigratorWorker jobs scheduled' do
+      Sidekiq::Testing.fake! do
+        ::HashedStorage::MigratorWorker.perform_async(1, 5)
+
+        expect(subject.migration_pending?).to be_truthy
+      end
+    end
+
+    it 'returns true when there are ProjectMigrateWorker jobs scheduled' do
+      Sidekiq::Testing.fake! do
+        ::HashedStorage::ProjectMigrateWorker.perform_async(1)
+
+        expect(subject.migration_pending?).to be_truthy
+      end
+    end
+
+    it 'returns false when queues are empty' do
+      expect(subject.migration_pending?).to be_falsey
+    end
+  end
+
+  describe 'rollback_pending?' do
+    set(:project) { create(:project, :empty_repo) }
+
+    it 'returns true when there are RollbackerWorker jobs scheduled' do
+      Sidekiq::Testing.fake! do
+        ::HashedStorage::RollbackerWorker.perform_async(1, 5)
+
+        expect(subject.rollback_pending?).to be_truthy
+      end
+    end
+
+    it 'returns true when there are jobs scheduled' do
+      Sidekiq::Testing.fake! do
+        ::HashedStorage::ProjectRollbackWorker.perform_async(1)
+
+        expect(subject.rollback_pending?).to be_truthy
+      end
+    end
+
+    it 'returns false when queues are empty' do
+      expect(subject.rollback_pending?).to be_falsey
+    end
+  end
 end
diff --git a/spec/tasks/gitlab/storage_rake_spec.rb b/spec/tasks/gitlab/storage_rake_spec.rb
index 6b50670c3c0..736809eee5b 100644
--- a/spec/tasks/gitlab/storage_rake_spec.rb
+++ b/spec/tasks/gitlab/storage_rake_spec.rb
@@ -1,6 +1,6 @@
 require 'rake_helper'
 
-describe 'rake gitlab:storage:*' do
+describe 'rake gitlab:storage:*', :sidekiq do
   before do
     Rake.application.rake_require 'tasks/gitlab/storage'
 
@@ -43,9 +43,7 @@ describe 'rake gitlab:storage:*' do
     end
   end
 
-  describe 'gitlab:storage:migrate_to_hashed' do
-    let(:task) { 'gitlab:storage:migrate_to_hashed' }
-
+  shared_examples "make sure database is writable" do
     context 'read-only database' do
       it 'does nothing' do
         expect(Gitlab::Database).to receive(:read_only?).and_return(true)
@@ -55,48 +53,68 @@ describe 'rake gitlab:storage:*' do
         expect { run_rake_task(task) }.to output(/This task requires database write access. Exiting./).to_stderr
       end
     end
+  end
 
-    context '0 legacy projects' do
-      it 'does nothing' do
-        expect(::HashedStorage::MigratorWorker).not_to receive(:perform_async)
+  shared_examples "handles custom BATCH env var" do |worker_klass|
+    context 'in batches of 1' do
+      before do
+        stub_env('BATCH' => 1)
+      end
+
+      it "enqueues one #{worker_klass} per project" do
+        projects.each do |project|
+          expect(worker_klass).to receive(:perform_async).with(project.id, project.id)
+        end
 
         run_rake_task(task)
       end
     end
 
-    context '3 legacy projects' do
-      let(:projects) { create_list(:project, 3, :legacy_storage) }
+    context 'in batches of 2' do
+      before do
+        stub_env('BATCH' => 2)
+      end
 
-      context 'in batches of 1' do
-        before do
-          stub_env('BATCH' => 1)
+      it "enqueues one #{worker_klass} per 2 projects" do
+        projects.map(&:id).sort.each_slice(2) do |first, last|
+          last ||= first
+          expect(worker_klass).to receive(:perform_async).with(first, last)
         end
 
-        it 'enqueues one HashedStorage::MigratorWorker per project' do
-          projects.each do |project|
-            expect(::HashedStorage::MigratorWorker).to receive(:perform_async).with(project.id, project.id)
-          end
-
-          run_rake_task(task)
-        end
+        run_rake_task(task)
       end
+    end
+  end
 
-      context 'in batches of 2' do
-        before do
-          stub_env('BATCH' => 2)
-        end
+  describe 'gitlab:storage:migrate_to_hashed' do
+    let(:task) { 'gitlab:storage:migrate_to_hashed' }
 
-        it 'enqueues one HashedStorage::MigratorWorker per 2 projects' do
-          projects.map(&:id).sort.each_slice(2) do |first, last|
-            last ||= first
-            expect(::HashedStorage::MigratorWorker).to receive(:perform_async).with(first, last)
-          end
+    context 'with rollback already scheduled' do
+      it 'does nothing' do
+        Sidekiq::Testing.fake! do
+          ::HashedStorage::RollbackerWorker.perform_async(1, 5)
+
+          expect(Project).not_to receive(:with_unmigrated_storage)
 
-          run_rake_task(task)
+          expect { run_rake_task(task) }.to output(/There is already a rollback operation in progress/).to_stderr
         end
       end
     end
 
+    context 'with 0 legacy projects' do
+      it 'does nothing' do
+        expect(::HashedStorage::MigratorWorker).not_to receive(:perform_async)
+
+        run_rake_task(task)
+      end
+    end
+
+    context 'with 3 legacy projects' do
+      let(:projects) { create_list(:project, 3, :legacy_storage) }
+
+      it_behaves_like "handles custom BATCH env var", ::HashedStorage::MigratorWorker
+    end
+
     context 'with same id in range' do
       it 'displays message when project cant be found' do
         stub_env('ID_FROM', 99999)
@@ -123,6 +141,38 @@ describe 'rake gitlab:storage:*' do
     end
   end
 
+  describe 'gitlab:storage:rollback_to_legacy' do
+    let(:task) { 'gitlab:storage:rollback_to_legacy' }
+
+    it_behaves_like 'make sure database is writable'
+
+    context 'with migration already scheduled' do
+      it 'does nothing' do
+        Sidekiq::Testing.fake! do
+          ::HashedStorage::MigratorWorker.perform_async(1, 5)
+
+          expect(Project).not_to receive(:with_unmigrated_storage)
+
+          expect { run_rake_task(task) }.to output(/There is already a migration operation in progress/).to_stderr
+        end
+      end
+    end
+
+    context 'with 0 hashed projects' do
+      it 'does nothing' do
+        expect(::HashedStorage::RollbackerWorker).not_to receive(:perform_async)
+
+        run_rake_task(task)
+      end
+    end
+
+    context 'with 3 hashed projects' do
+      let(:projects) { create_list(:project, 3) }
+
+      it_behaves_like "handles custom BATCH env var", ::HashedStorage::RollbackerWorker
+    end
+  end
+
   describe 'gitlab:storage:legacy_projects' do
     it_behaves_like 'rake entities summary', 'projects', 'Legacy' do
       let(:task) { 'gitlab:storage:legacy_projects' }