diff --git a/lib/gitlab/background_migration/migrate_build_stage_id_reference.rb b/lib/gitlab/background_migration/migrate_build_stage_id_reference.rb
index 711126ea0d3f3b5646bcbf237afcc975c687fae2..c8669ca32720a4cf88f1475ab35998a2905a76ec 100644
--- a/lib/gitlab/background_migration/migrate_build_stage_id_reference.rb
+++ b/lib/gitlab/background_migration/migrate_build_stage_id_reference.rb
@@ -5,12 +5,13 @@ module Gitlab
         raise ArgumentError unless id.present?
 
         sql = <<-SQL.strip_heredoc
-          UPDATE "ci_builds" SET "stage_id" = (
-            SELECT id FROM ci_stages
-              WHERE ci_stages.pipeline_id = ci_builds.commit_id
-              AND ci_stages.name = ci_builds.stage
-            )
-          WHERE "ci_builds"."id" = #{id} AND "ci_builds"."stage_id" IS NULL
+          UPDATE "ci_builds"
+            SET "stage_id" =
+              (SELECT id FROM ci_stages
+                WHERE ci_stages.pipeline_id = ci_builds.commit_id
+                AND ci_stages.name = ci_builds.stage)
+          WHERE "ci_builds"."id" = #{id}
+            AND "ci_builds"."stage_id" IS NULL
         SQL
 
         ActiveRecord::Base.connection.execute(sql)
diff --git a/spec/migrations/migrate_stage_id_reference_in_background_spec.rb b/spec/migrations/migrate_stage_id_reference_in_background_spec.rb
index 1bd2c14b61ccfb88c5afd3d21d34283ecb8ce36d..63787d712331287f808f4995a0bfb691878b1833 100644
--- a/spec/migrations/migrate_stage_id_reference_in_background_spec.rb
+++ b/spec/migrations/migrate_stage_id_reference_in_background_spec.rb
@@ -58,11 +58,11 @@ describe MigrateStageIdReferenceInBackground, :migration, :sidekiq do
 
   it 'schedules background migrations' do
     Sidekiq::Testing.inline! do
-      expect(jobs.where(stage_id: nil)).to be_present
+      expect(jobs.where(stage_id: nil).count).to eq 5
 
       migrate!
 
-      expect(jobs.where(stage_id: nil)).to be_one
+      expect(jobs.where(stage_id: nil).count).to eq 1
     end
   end
 end