1 files changed, 25 insertions, 6 deletions
diff --git a/spec/lib/gitlab/exclusive_lease_helpers_spec.rb b/spec/lib/gitlab/exclusive_lease_helpers_spec.rb
index 747fe369c78..9914518cda5 100644
--- a/spec/lib/gitlab/exclusive_lease_helpers_spec.rb
+++ b/spec/lib/gitlab/exclusive_lease_helpers_spec.rb
@@ -22,9 +22,7 @@ describe Gitlab::ExclusiveLeaseHelpers, :clean_gitlab_redis_shared_state do
     end
 
     context 'when the lease is not obtained yet' do
-      before do
-        stub_exclusive_lease(unique_key, 'uuid')
-      end
+      let!(:lease) { stub_exclusive_lease(unique_key, 'uuid') }
 
       it 'calls the given block' do
         expect { |b| class_instance.in_lock(unique_key, &b) }.to yield_with_args(false)
@@ -37,7 +35,7 @@ describe Gitlab::ExclusiveLeaseHelpers, :clean_gitlab_redis_shared_state do
       end
 
       it 'cancels the exclusive lease after the block' do
-        expect_to_cancel_exclusive_lease(unique_key, 'uuid')
+        expect(lease).to receive(:cancel).once
 
         subject
       end
@@ -81,11 +79,32 @@ describe Gitlab::ExclusiveLeaseHelpers, :clean_gitlab_redis_shared_state do
         end
       end
 
+      context 'when we specify no retries' do
+        let(:options) { { retries: 0 } }
+
+        it 'never sleeps' do
+          expect(class_instance).not_to receive(:sleep)
+
+          expect { subject }.to raise_error('Failed to obtain a lock')
+        end
+      end
+
       context 'when sleep second is specified' do
-        let(:options) { { retries: 0, sleep_sec: 0.05.seconds } }
+        let(:options) { { retries: 1, sleep_sec: 0.05.seconds } }
+
+        it 'receives the specified argument' do
+          expect_any_instance_of(Object).to receive(:sleep).with(0.05.seconds).once
+
+          expect { subject }.to raise_error('Failed to obtain a lock')
+        end
+      end
+
+      context 'when sleep second is specified as a lambda' do
+        let(:options) { { retries: 2, sleep_sec: ->(num) { 0.1 + num } } }
 
         it 'receives the specified argument' do
-          expect(class_instance).to receive(:sleep).with(0.05.seconds).once
+          expect_any_instance_of(Object).to receive(:sleep).with(1.1.seconds).once
+          expect_any_instance_of(Object).to receive(:sleep).with(2.1.seconds).once
 
           expect { subject }.to raise_error('Failed to obtain a lock')
         end