1 files changed, 81 insertions, 0 deletions
diff --git a/spec/models/ci/pipeline_spec.rb b/spec/models/ci/pipeline_spec.rb
index e7845b693a1..f3725e24470 100644
--- a/spec/models/ci/pipeline_spec.rb
+++ b/spec/models/ci/pipeline_spec.rb
@@ -500,6 +500,87 @@ describe Ci::Pipeline, :mailer do
         end
       end
     end
+
+    describe '#stages' do
+      before do
+        create(:ci_stage_entity, project: project,
+                                 pipeline: pipeline,
+                                 name: 'build')
+      end
+
+      it 'returns persisted stages' do
+        expect(pipeline.stages).not_to be_empty
+        expect(pipeline.stages).to all(be_persisted)
+      end
+    end
+
+    describe '#ordered_stages' do
+      before do
+        create(:ci_stage_entity, project: project,
+                                 pipeline: pipeline,
+                                 position: 4,
+                                 name: 'deploy')
+
+        create(:ci_build, project: project,
+                          pipeline: pipeline,
+                          stage: 'test',
+                          stage_idx: 3,
+                          name: 'test')
+
+        create(:ci_build, project: project,
+                          pipeline: pipeline,
+                          stage: 'build',
+                          stage_idx: 2,
+                          name: 'build')
+
+        create(:ci_stage_entity, project: project,
+                                 pipeline: pipeline,
+                                 position: 1,
+                                 name: 'sanity')
+
+        create(:ci_stage_entity, project: project,
+                                 pipeline: pipeline,
+                                 position: 5,
+                                 name: 'cleanup')
+      end
+
+      subject { pipeline.ordered_stages }
+
+      context 'when using legacy stages' do
+        before do
+          stub_feature_flags(ci_pipeline_persisted_stages: false)
+        end
+
+        it 'returns legacy stages in valid order' do
+          expect(subject.map(&:name)).to eq %w[build test]
+        end
+      end
+
+      context 'when using persisted stages' do
+        before do
+          stub_feature_flags(ci_pipeline_persisted_stages: true)
+        end
+
+        context 'when pipelines is not complete' do
+          it 'still returns legacy stages' do
+            expect(subject).to all(be_a Ci::LegacyStage)
+            expect(subject.map(&:name)).to eq %w[build test]
+          end
+        end
+
+        context 'when pipeline is complete' do
+          before do
+            pipeline.succeed!
+          end
+
+          it 'returns stages in valid order' do
+            expect(subject).to all(be_a Ci::Stage)
+            expect(subject.map(&:name))
+              .to eq %w[sanity build test deploy cleanup]
+          end
+        end
+      end
+    end
   end
 
   describe 'state machine' do