Use a new stage_id reference to a persisted stage

2 files changed, 5 insertions, 6 deletions

diff --git a/app/models/commit_status.rb b/app/models/commit_status.rb
index afdc75f75fb..efb5cbd9d41 100644
--- a/app/models/commit_status.rb
+++ b/app/models/commit_status.rb
@@ -43,9 +43,7 @@ class CommitStatus < ActiveRecord::Base
   # `ci_builds.stage` attribute and migrate `ci_builds.stage_id` reference in
   # one of upcoming releases.
   #
-  def stage_entity
-    Ci::Stage.find_by(pipeline: pipeline, name: stage)
-  end
+  belongs_to :stage_entity, foreign_key: :stage_id, class_name: 'Ci::Stage'
 
   state_machine :status do
     event :enqueue do
diff --git a/spec/models/commit_status_spec.rb b/spec/models/commit_status_spec.rb
index c7651ce9b46..f54e1131813 100644
--- a/spec/models/commit_status_spec.rb
+++ b/spec/models/commit_status_spec.rb
@@ -409,9 +409,10 @@ describe CommitStatus, :models do
   end
 
   describe '#stage_entity' do
-    let!(:stage) do
-      create(:ci_stage_entity, pipeline: commit_status.pipeline,
-                               name: commit_status.stage)
+    let(:stage) { create(:ci_stage_entity) }
+
+    let(:commit_status) do
+      create(:commit_status, stage_id: stage.id)
     end
 
     it 'has a correct association with persisted stage' do