1 files changed, 25 insertions, 1 deletions
diff --git a/spec/factories/ci/processable.rb b/spec/factories/ci/processable.rb
index 76c7376d24a..49e66368f94 100644
--- a/spec/factories/ci/processable.rb
+++ b/spec/factories/ci/processable.rb
@@ -3,13 +3,37 @@
 FactoryBot.define do
   factory :ci_processable, class: 'Ci::Processable' do
     name { 'processable' }
-    stage { 'test' }
     stage_idx { ci_stage.try(:position) || 0 }
     ref { 'master' }
     tag { false }
     pipeline factory: :ci_pipeline
     project { pipeline.project }
     scheduling_type { 'stage' }
+    partition_id { pipeline.partition_id }
+
+    # This factory was updated to help with the efforts of the removal of `ci_builds.stage`:
+    # https://gitlab.com/gitlab-org/gitlab/-/issues/364377
+    # These additions can be removed once the specs that use the stage attribute have been updated
+
+    transient do
+      stage { 'test' }
+    end
+
+    after(:build) do |processable, evaluator|
+      processable.stage = evaluator.stage
+    end
+
+    before(:create) do |processable, evaluator|
+      next if processable.ci_stage
+
+      if ci_stage = processable.pipeline.stages.find_by(name: evaluator.stage)
+        processable.ci_stage = ci_stage
+      else
+        processable.ci_stage = create(:ci_stage, pipeline: processable.pipeline,
+                                 project: processable.project || evaluator.project,
+                                 name: evaluator.stage, position: evaluator.stage_idx, status: 'created')
+      end
+    end
 
     trait :waiting_for_resource do
       status { 'waiting_for_resource' }