1 files changed, 5 insertions, 5 deletions
diff --git a/spec/serializers/stage_entity_spec.rb b/spec/serializers/stage_entity_spec.rb
index b977d5d33aa..95d3fd254d4 100644
--- a/spec/serializers/stage_entity_spec.rb
+++ b/spec/serializers/stage_entity_spec.rb
@@ -12,12 +12,12 @@ RSpec.describe StageEntity do
   end
 
   let(:stage) do
-    build(:ci_stage, pipeline: pipeline, name: 'test')
+    create(:ci_stage, pipeline: pipeline, status: :success)
   end
 
   before do
     allow(request).to receive(:current_user).and_return(user)
-    create(:ci_build, :success, pipeline: pipeline)
+    create(:ci_build, :success, pipeline: pipeline, stage_id: stage.id)
   end
 
   describe '#as_json' do
@@ -74,7 +74,7 @@ RSpec.describe StageEntity do
     end
 
     context 'with a skipped stage ' do
-      let(:stage) { create(:ci_stage_entity, status: 'skipped') }
+      let(:stage) { create(:ci_stage, status: 'skipped') }
 
       it 'contains play_all_manual' do
         expect(subject[:status][:action]).to be_present
@@ -82,7 +82,7 @@ RSpec.describe StageEntity do
     end
 
     context 'with a scheduled stage ' do
-      let(:stage) { create(:ci_stage_entity, status: 'scheduled') }
+      let(:stage) { create(:ci_stage, status: 'scheduled') }
 
       it 'contains play_all_manual' do
         expect(subject[:status][:action]).to be_present
@@ -90,7 +90,7 @@ RSpec.describe StageEntity do
     end
 
     context 'with a manual stage ' do
-      let(:stage) { create(:ci_stage_entity, status: 'manual') }
+      let(:stage) { create(:ci_stage, status: 'manual') }
 
       it 'contains play_all_manual' do
         expect(subject[:status][:action]).to be_present