1 files changed, 44 insertions, 13 deletions
diff --git a/spec/models/ci/build_metadata_spec.rb b/spec/models/ci/build_metadata_spec.rb
index e904463a5ca..16cff72db64 100644
--- a/spec/models/ci/build_metadata_spec.rb
+++ b/spec/models/ci/build_metadata_spec.rb
@@ -14,8 +14,8 @@ RSpec.describe Ci::BuildMetadata do
                          status: 'success')
   end
 
-  let(:build) { create(:ci_build, pipeline: pipeline) }
-  let(:metadata) { build.metadata }
+  let(:job) { create(:ci_build, pipeline: pipeline) }
+  let(:metadata) { job.metadata }
 
   it_behaves_like 'having unique enum values'
 
@@ -35,7 +35,7 @@ RSpec.describe Ci::BuildMetadata do
     context 'when project timeout is set' do
       context 'when runner is assigned to the job' do
         before do
-          build.update!(runner: runner)
+          job.update!(runner: runner)
         end
 
         context 'when runner timeout is not set' do
@@ -59,13 +59,13 @@ RSpec.describe Ci::BuildMetadata do
 
       context 'when job timeout is set' do
         context 'when job timeout is higher than project timeout' do
-          let(:build) { create(:ci_build, pipeline: pipeline, options: { job_timeout: 3000 }) }
+          let(:job) { create(:ci_build, pipeline: pipeline, options: { job_timeout: 3000 }) }
 
           it_behaves_like 'sets timeout', 'job_timeout_source', 3000
         end
 
         context 'when job timeout is lower than project timeout' do
-          let(:build) { create(:ci_build, pipeline: pipeline, options: { job_timeout: 1000 }) }
+          let(:job) { create(:ci_build, pipeline: pipeline, options: { job_timeout: 1000 }) }
 
           it_behaves_like 'sets timeout', 'job_timeout_source', 1000
         end
@@ -73,18 +73,18 @@ RSpec.describe Ci::BuildMetadata do
 
       context 'when both runner and job timeouts are set' do
         before do
-          build.update!(runner: runner)
+          job.update!(runner: runner)
         end
 
         context 'when job timeout is higher than runner timeout' do
-          let(:build) { create(:ci_build, pipeline: pipeline, options: { job_timeout: 3000 }) }
+          let(:job) { create(:ci_build, pipeline: pipeline, options: { job_timeout: 3000 }) }
           let(:runner) { create(:ci_runner, maximum_timeout: 2100) }
 
           it_behaves_like 'sets timeout', 'runner_timeout_source', 2100
         end
 
         context 'when job timeout is lower than runner timeout' do
-          let(:build) { create(:ci_build, pipeline: pipeline, options: { job_timeout: 1900 }) }
+          let(:job) { create(:ci_build, pipeline: pipeline, options: { job_timeout: 1900 }) }
           let(:runner) { create(:ci_runner, maximum_timeout: 2100) }
 
           it_behaves_like 'sets timeout', 'job_timeout_source', 1900
@@ -135,20 +135,51 @@ RSpec.describe Ci::BuildMetadata do
 
   describe 'set_cancel_gracefully' do
     it 'sets cancel_gracefully' do
-      build.set_cancel_gracefully
+      job.set_cancel_gracefully
 
-      expect(build.cancel_gracefully?).to be true
+      expect(job.cancel_gracefully?).to be true
     end
 
     it 'returns false' do
-      expect(build.cancel_gracefully?).to be false
+      expect(job.cancel_gracefully?).to be false
     end
   end
 
   context 'loose foreign key on ci_builds_metadata.project_id' do
     it_behaves_like 'cleanup by a loose foreign key' do
-      let!(:parent) { create(:project) }
-      let!(:model) { create(:ci_build_metadata, project: parent) }
+      let!(:parent) { project }
+      let!(:model) { metadata }
+    end
+  end
+
+  describe 'partitioning' do
+    context 'with job' do
+      let(:status) { build(:commit_status, partition_id: 123) }
+      let(:metadata) { build(:ci_build_metadata, build: status) }
+
+      it 'copies the partition_id from job' do
+        expect { metadata.valid? }.to change(metadata, :partition_id).to(123)
+      end
+
+      context 'when it is already set' do
+        let(:metadata) { build(:ci_build_metadata, build: status, partition_id: 125) }
+
+        it 'does not change the partition_id value' do
+          expect { metadata.valid? }.not_to change(metadata, :partition_id)
+        end
+      end
+    end
+
+    context 'without job' do
+      subject(:metadata) do
+        build(:ci_build_metadata, build: nil)
+      end
+
+      it { is_expected.to validate_presence_of(:partition_id) }
+
+      it 'does not change the partition_id value' do
+        expect { metadata.valid? }.not_to change(metadata, :partition_id)
+      end
     end
   end
 end