1 files changed, 29 insertions, 0 deletions
diff --git a/spec/models/ci/stage_spec.rb b/spec/models/ci/stage_spec.rb
index d55a8509a98..dd9af33a562 100644
--- a/spec/models/ci/stage_spec.rb
+++ b/spec/models/ci/stage_spec.rb
@@ -369,4 +369,33 @@ RSpec.describe Ci::Stage, :models do
       let!(:model) { create(:ci_stage, project: parent) }
     end
   end
+
+  describe 'partitioning' do
+    context 'with pipeline' do
+      let(:pipeline) { build(:ci_pipeline, partition_id: 123) }
+      let(:stage) { build(:ci_stage, pipeline: pipeline) }
+
+      it 'copies the partition_id from pipeline' do
+        expect { stage.valid? }.to change(stage, :partition_id).to(123)
+      end
+
+      context 'when it is already set' do
+        let(:stage) { build(:ci_stage, pipeline: pipeline, partition_id: 125) }
+
+        it 'does not change the partition_id value' do
+          expect { stage.valid? }.not_to change(stage, :partition_id)
+        end
+      end
+    end
+
+    context 'without pipeline' do
+      subject(:stage) { build(:ci_stage, pipeline: nil) }
+
+      it { is_expected.to validate_presence_of(:partition_id) }
+
+      it 'does not change the partition_id value' do
+        expect { stage.valid? }.not_to change(stage, :partition_id)
+      end
+    end
+  end
 end