1 files changed, 15 insertions, 6 deletions
diff --git a/spec/lib/gitlab/background_migration/set_correct_vulnerability_state_spec.rb b/spec/lib/gitlab/background_migration/set_correct_vulnerability_state_spec.rb
index d5b98e49a31..2372ce21c4c 100644
--- a/spec/lib/gitlab/background_migration/set_correct_vulnerability_state_spec.rb
+++ b/spec/lib/gitlab/background_migration/set_correct_vulnerability_state_spec.rb
@@ -34,7 +34,7 @@ RSpec.describe Gitlab::BackgroundMigration::SetCorrectVulnerabilityState do
   let(:detected_state) { 1 }
   let(:dismissed_state) { 2 }
 
-  subject(:perform_migration) do
+  let(:migration_job) do
     described_class.new(start_id: vulnerability_with_dismissed_at.id,
                         end_id: vulnerability_without_dismissed_at.id,
                         batch_table: :vulnerabilities,
@@ -42,15 +42,24 @@ RSpec.describe Gitlab::BackgroundMigration::SetCorrectVulnerabilityState do
                         sub_batch_size: 1,
                         pause_ms: 0,
                         connection: ActiveRecord::Base.connection)
-                   .perform
   end
 
-  it 'changes vulnerability state to `dismissed` when dismissed_at is not nil' do
-    expect { perform_migration }.to change { vulnerability_with_dismissed_at.reload.state }.to(dismissed_state)
+  describe '#filter_batch' do
+    it 'filters out vulnerabilities where dismissed_at is null' do
+      expect(migration_job.filter_batch(vulnerabilities)).to contain_exactly(vulnerability_with_dismissed_at)
+    end
   end
 
-  it 'does not change the state when dismissed_at is nil' do
-    expect { perform_migration }.not_to change { vulnerability_without_dismissed_at.reload.state }
+  describe '#perform' do
+    subject(:perform_migration) { migration_job.perform }
+
+    it 'changes vulnerability state to `dismissed` when dismissed_at is not nil' do
+      expect { perform_migration }.to change { vulnerability_with_dismissed_at.reload.state }.to(dismissed_state)
+    end
+
+    it 'does not change the state when dismissed_at is nil' do
+      expect { perform_migration }.not_to change { vulnerability_without_dismissed_at.reload.state }
+    end
   end
 
   private