1 files changed, 4 insertions, 2 deletions
diff --git a/spec/factories/ci/stages.rb b/spec/factories/ci/stages.rb
index d37eabb3e8c..d3c8bf9d54f 100644
--- a/spec/factories/ci/stages.rb
+++ b/spec/factories/ci/stages.rb
@@ -1,5 +1,5 @@
 FactoryGirl.define do
-  factory :ci_stage, class: Ci::Stage do
+  factory :ci_stage, class: Ci::LegacyStage do
     skip_create
 
     transient do
@@ -10,7 +10,9 @@ FactoryGirl.define do
     end
 
     initialize_with do
-      Ci::Stage.new(pipeline, name: name, status: status, warnings: warnings)
+      Ci::LegacyStage.new(pipeline, name: name,
+                                    status: status,
+                                    warnings: warnings)
     end
   end
 end