1 files changed, 64 insertions, 19 deletions
diff --git a/spec/lib/gitlab/background_migration_spec.rb b/spec/lib/gitlab/background_migration_spec.rb
index b28329792eb..9c6b07f02cc 100644
--- a/spec/lib/gitlab/background_migration_spec.rb
+++ b/spec/lib/gitlab/background_migration_spec.rb
@@ -24,7 +24,8 @@ describe Gitlab::BackgroundMigration do
         it 'steals jobs from a queue' do
           expect(queue[0]).to receive(:delete).and_return(true)
 
-          expect(described_class).to receive(:perform).with('Foo', [10, 20])
+          expect(described_class).to receive(:perform)
+            .with('Foo', [10, 20], anything)
 
           described_class.steal('Foo')
         end
@@ -32,7 +33,7 @@ describe Gitlab::BackgroundMigration do
         it 'does not steal job that has already been taken' do
           expect(queue[0]).to receive(:delete).and_return(false)
 
-          expect(described_class).not_to receive(:perform).with('Foo', [10, 20])
+          expect(described_class).not_to receive(:perform)
 
           described_class.steal('Foo')
         end
@@ -57,17 +58,40 @@ describe Gitlab::BackgroundMigration do
         before do
           stub_const("#{described_class}::Foo", migration)
 
-          allow(migration).to receive(:perform).with(10, 20)
-            .and_raise(StandardError, 'Migration error')
-
           allow(queue[0]).to receive(:delete).and_return(true)
           allow(queue[1]).to receive(:delete).and_return(true)
         end
 
-        it 'recovers from an exceptions and continues' do
-          expect(migration).to receive(:perform).twice
-          expect { described_class.steal('Foo') }
-            .to output(/Migration error/).to_stdout
+        context 'when standard error is being raised' do
+          before do
+            allow(migration).to receive(:perform).with(10, 20)
+              .and_raise(StandardError, 'Migration error')
+          end
+
+          it 'recovers from an exception and retries the migration' do
+            expect(migration).to receive(:perform).with(10, 20)
+              .exactly(3).times.ordered
+            expect(migration).to receive(:perform).with(20, 30)
+              .once.ordered
+            expect(Rails.logger).to receive(:warn)
+              .with(/Retrying background migration/).exactly(3).times
+
+            described_class.steal('Foo')
+          end
+        end
+
+        context 'when top level exception is being raised' do
+          it 'enqueues the migration again and reraises the error' do
+            allow(migration).to receive(:perform).with(10, 20)
+              .and_raise(Exception, 'Migration error').once
+
+            expect(BackgroundMigrationWorker).to receive(:perform_async)
+              .with('Foo', [10, 20]).once
+
+            expect(Rails.logger).not_to receive(:warn)
+            expect { described_class.steal('Foo') }
+              .to raise_error(Exception)
+          end
         end
       end
     end
@@ -91,9 +115,9 @@ describe Gitlab::BackgroundMigration do
       it 'steals from the scheduled sets queue first' do
         Sidekiq::Testing.disable! do
           expect(described_class).to receive(:perform)
-            .with('Object', [1]).ordered
+            .with('Object', [1], anything).ordered
           expect(described_class).to receive(:perform)
-            .with('Object', [2]).ordered
+            .with('Object', [2], anything).ordered
 
           BackgroundMigrationWorker.perform_async('Object', [2])
           BackgroundMigrationWorker.perform_in(10.minutes, 'Object', [1])
@@ -105,17 +129,38 @@ describe Gitlab::BackgroundMigration do
   end
 
   describe '.perform' do
-    it 'performs a background migration' do
-      instance = double(:instance)
-      klass = double(:klass, new: instance)
+    let(:migration) { spy(:migration) }
 
-      expect(described_class).to receive(:const_get)
-        .with('Foo')
-        .and_return(klass)
+    before do
+      stub_const("#{described_class.name}::Foo", migration)
+    end
 
-      expect(instance).to receive(:perform).with(10, 20)
+    context 'when retries count is not specified' do
+      it 'performs a background migration' do
+        expect(migration).to receive(:perform).with(10, 20).once
 
-      described_class.perform('Foo', [10, 20])
+        described_class.perform('Foo', [10, 20])
+      end
+    end
+
+    context 'when retries count is zero' do
+      it 'perform a background migration only once' do
+        expect(migration).to receive(:perform).with(10, 20)
+          .and_raise(StandardError).once
+
+        expect {  described_class.perform('Foo', [10, 20], retries: 0) }
+          .to raise_error(StandardError)
+      end
+    end
+
+    context 'when retries count is larger than zero' do
+      it 'retries a background migration when needed' do
+        expect(migration).to receive(:perform).with(10, 20)
+          .and_raise(StandardError).exactly(3).times
+
+        expect {  described_class.perform('Foo', [10, 20], retries: 3) }
+          .to raise_error(StandardError)
+      end
     end
   end
 end