2 files changed, 57 insertions, 1 deletions
diff --git a/app/services/ci/ensure_stage_service.rb b/app/services/ci/ensure_stage_service.rb
index dc2f49e8db1..9a48f7096cf 100644
--- a/app/services/ci/ensure_stage_service.rb
+++ b/app/services/ci/ensure_stage_service.rb
@@ -7,6 +7,8 @@ module Ci
   # stage.
   #
   class EnsureStageService < BaseService
+    PipelineStageError = Class.new(StandardError)
+
     def execute(build)
       @build = build
 
@@ -22,8 +24,11 @@ module Ci
 
     private
 
-    def ensure_stage
+    def ensure_stage(attempts: 2)
       find_stage || create_stage
+    rescue ActiveRecord::RecordNotUnique
+      retry if (attempts -= 1) > 0
+      raise PipelineStageError, 'Fix me!'
     end
 
     def find_stage
diff --git a/spec/services/ci/ensure_stage_service_spec.rb b/spec/services/ci/ensure_stage_service_spec.rb
new file mode 100644
index 00000000000..47ba1e74c15
--- /dev/null
+++ b/spec/services/ci/ensure_stage_service_spec.rb
@@ -0,0 +1,51 @@
+require 'spec_helper'
+
+describe Ci::EnsureStageService, '#execute' do
+  set(:project) { create(:project) }
+  set(:user) { create(:user) }
+
+  let(:stage) { create(:ci_stage_entity) }
+  let(:job) { build(:ci_build) }
+
+  let(:service) { described_class.new(project, user) }
+
+  context 'when build has a stage assigned' do
+   it 'does not create a new stage' do
+     job.assign_attributes(stage_id: stage.id)
+
+     expect { service.execute(job) }.not_to change { Ci::Stage.count }
+   end
+  end
+
+  context 'when build does not have a stage assigned' do
+   it 'creates a new stage' do
+     job.assign_attributes(stage_id: nil, stage: 'test')
+
+     expect { service.execute(job) }.to change { Ci::Stage.count }.by(1)
+   end
+  end
+
+  context 'when build is invalid' do
+   it 'does not create a new stage' do
+     job.assign_attributes(stage_id: nil, ref: nil)
+
+     expect { service.execute(job) }.not_to change { Ci::Stage.count }
+   end
+  end
+
+  context 'when new stage can not be created because of an exception' do
+    before do
+      allow(Ci::Stage).to receive(:create!)
+        .and_raise(ActiveRecord::RecordNotUnique.new('Duplicates!'))
+    end
+
+    it 'retries up to two times' do
+      job.assign_attributes(stage_id: nil)
+
+      expect(service).to receive(:find_stage).exactly(2).times
+
+      expect { service.execute(job) }
+        .to raise_error(Ci::EnsureStageService::PipelineStageError)
+    end
+  end
+end