Implement proper associations with a persisted stage

2 files changed, 23 insertions, 2 deletions

diff --git a/app/models/ci/stage.rb b/app/models/ci/stage.rb
index da1c3753924..ca89caf4782 100644
--- a/app/models/ci/stage.rb
+++ b/app/models/ci/stage.rb
@@ -9,8 +9,8 @@ module Ci
     belongs_to :project
     belongs_to :pipeline
 
-    has_many :statuses, class_name: 'CommitStatus', foreign_key: :commit_id
-    has_many :builds, foreign_key: :commit_id
+    has_many :commit_statuses, foreign_key: :stage_id
+    has_many :builds, foreign_key: :stage_id
 
     validates :project, presence: true, unless: :importing?
     validates :pipeline, presence: true, unless: :importing?
diff --git a/spec/models/ci/stage_spec.rb b/spec/models/ci/stage_spec.rb
index 911c468ff1a..49573175266 100644
--- a/spec/models/ci/stage_spec.rb
+++ b/spec/models/ci/stage_spec.rb
@@ -1,6 +1,27 @@
 require 'spec_helper'
 
 describe Ci::Stage, :models do
+  describe 'associations' do
+    let(:stage) { create(:ci_stage_entity) }
+
+    before do
+      create(:ci_build, stage_id: stage.id)
+      create(:commit_status, stage_id: stage.id)
+    end
+
+    describe '#commit_statuses' do
+      it 'returns all commit statuses' do
+        expect(stage.commit_statuses.count).to be 2
+      end
+    end
+
+    describe '#builds' do
+      it 'returns only builds' do
+        expect(stage.builds).to be_one
+      end
+    end
+  end
+
   describe '#status' do
     context 'when stage is pending' do
       let(:stage) { create(:ci_stage_entity, status: 'pending') }