1 files changed, 150 insertions, 1 deletions
diff --git a/spec/lib/gitlab/background_migration/batched_migration_job_spec.rb b/spec/lib/gitlab/background_migration/batched_migration_job_spec.rb
index 98866bb765f..f03f90ddbbb 100644
--- a/spec/lib/gitlab/background_migration/batched_migration_job_spec.rb
+++ b/spec/lib/gitlab/background_migration/batched_migration_job_spec.rb
@@ -3,6 +3,113 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::BackgroundMigration::BatchedMigrationJob do
+  let(:connection) { Gitlab::Database.database_base_models[:main].connection }
+
+  describe '.generic_instance' do
+    it 'defines generic instance with only some of the attributes set' do
+      generic_instance = described_class.generic_instance(
+        batch_table: 'projects', batch_column: 'id',
+        job_arguments: %w(x y), connection: connection
+      )
+
+      expect(generic_instance.send(:batch_table)).to eq('projects')
+      expect(generic_instance.send(:batch_column)).to eq('id')
+      expect(generic_instance.instance_variable_get('@job_arguments')).to eq(%w(x y))
+      expect(generic_instance.send(:connection)).to eq(connection)
+
+      %i(start_id end_id sub_batch_size pause_ms).each do |attr|
+        expect(generic_instance.send(attr)).to eq(0)
+      end
+    end
+  end
+
+  describe '.job_arguments' do
+    let(:job_class) do
+      Class.new(described_class) do
+        job_arguments :value_a, :value_b
+      end
+    end
+
+    subject(:job_instance) do
+      job_class.new(start_id: 1, end_id: 10,
+                    batch_table: '_test_table',
+                    batch_column: 'id',
+                    sub_batch_size: 2,
+                    pause_ms: 1000,
+                    job_arguments: %w(a b),
+                    connection: connection)
+    end
+
+    it 'defines methods' do
+      expect(job_instance.value_a).to eq('a')
+      expect(job_instance.value_b).to eq('b')
+      expect(job_class.job_arguments_count).to eq(2)
+    end
+
+    context 'when no job arguments are defined' do
+      let(:job_class) do
+        Class.new(described_class)
+      end
+
+      it 'job_arguments_count is 0' do
+        expect(job_class.job_arguments_count).to eq(0)
+      end
+    end
+  end
+
+  describe '.scope_to' do
+    subject(:job_instance) do
+      job_class.new(start_id: 1, end_id: 10,
+                    batch_table: '_test_table',
+                    batch_column: 'id',
+                    sub_batch_size: 2,
+                    pause_ms: 1000,
+                    job_arguments: %w(a b),
+                    connection: connection)
+    end
+
+    context 'when additional scoping is defined' do
+      let(:job_class) do
+        Class.new(described_class) do
+          job_arguments :value_a, :value_b
+          scope_to ->(r) { "#{r}-#{value_a}-#{value_b}".upcase }
+        end
+      end
+
+      it 'applies additional scope to the provided relation' do
+        expect(job_instance.filter_batch('relation')).to eq('RELATION-A-B')
+      end
+    end
+
+    context 'when there is no additional scoping defined' do
+      let(:job_class) do
+        Class.new(described_class) do
+        end
+      end
+
+      it 'returns provided relation as is' do
+        expect(job_instance.filter_batch('relation')).to eq('relation')
+      end
+    end
+  end
+
+  describe 'descendants', :eager_load do
+    it 'have the same method signature for #perform' do
+      expected_arity = described_class.instance_method(:perform).arity
+      offences = described_class.descendants.select { |klass| klass.instance_method(:perform).arity != expected_arity }
+
+      expect(offences).to be_empty, "expected no descendants of #{described_class} to accept arguments for " \
+        "'#perform', but some do: #{offences.join(", ")}"
+    end
+
+    it 'do not use .batching_scope' do
+      offences = described_class.descendants.select { |klass| klass.respond_to?(:batching_scope) }
+
+      expect(offences).to be_empty, "expected no descendants of #{described_class} to define '.batching_scope', " \
+        "but some do: #{offences.join(", ")}"
+    end
+  end
+
   describe '#perform' do
     let(:connection) { Gitlab::Database.database_base_models[:main].connection }
 
@@ -66,6 +173,30 @@ RSpec.describe Gitlab::BackgroundMigration::BatchedMigrationJob do
         expect(test_table.order(:id).pluck(:to_column)).to contain_exactly(5, 10, nil, 20)
       end
 
+      context 'with additional scoping' do
+        let(:job_class) do
+          Class.new(described_class) do
+            scope_to ->(r) { r.where('mod(id, 2) = 0') }
+
+            def perform(*job_arguments)
+              each_sub_batch(
+                operation_name: :update,
+                batching_arguments: { order_hint: :updated_at },
+                batching_scope: -> (relation) { relation.where.not(bar: nil) }
+              ) do |sub_batch|
+                sub_batch.update_all('to_column = from_column')
+              end
+            end
+          end
+        end
+
+        it 'respects #filter_batch' do
+          expect { perform_job }.to change { test_table.where(to_column: nil).count }.from(4).to(2)
+
+          expect(test_table.order(:id).pluck(:to_column)).to contain_exactly(nil, 10, nil, 20)
+        end
+      end
+
       it 'instruments the batch operation' do
         expect(job_instance.batch_metrics.affected_rows).to be_empty
 
@@ -84,7 +215,7 @@ RSpec.describe Gitlab::BackgroundMigration::BatchedMigrationJob do
 
       context 'when batching_arguments are given' do
         it 'forwards them for batching' do
-          expect(job_instance).to receive(:parent_batch_relation).and_return(test_table)
+          expect(job_instance).to receive(:base_relation).and_return(test_table)
 
           expect(test_table).to receive(:each_batch).with(column: 'id', of: 2, order_hint: :updated_at)
 
@@ -155,6 +286,24 @@ RSpec.describe Gitlab::BackgroundMigration::BatchedMigrationJob do
 
         expect(job_instance.batch_metrics.affected_rows[:insert]).to contain_exactly(2, 1)
       end
+
+      context 'when used in combination with scope_to' do
+        let(:job_class) do
+          Class.new(described_class) do
+            scope_to ->(r) { r.where.not(from_column: 10) }
+
+            def perform(*job_arguments)
+              distinct_each_batch(operation_name: :insert) do |sub_batch|
+              end
+            end
+          end
+        end
+
+        it 'raises an error' do
+          expect { perform_job }.to raise_error RuntimeError,
+            /distinct_each_batch can not be used when additional filters are defined with scope_to/
+        end
+      end
     end
   end
 end