Add helper methods to stub Gitlab::ExclusiveLease

5 files changed, 114 insertions, 94 deletions

diff --git a/spec/workers/project_cache_worker_spec.rb b/spec/workers/project_cache_worker_spec.rb
index 6b1f2ff3227..b9b5445562f 100644
--- a/spec/workers/project_cache_worker_spec.rb
+++ b/spec/workers/project_cache_worker_spec.rb
@@ -1,14 +1,17 @@
 require 'spec_helper'
 
 describe ProjectCacheWorker do
+  include ExclusiveLeaseHelpers
+
   let(:worker) { described_class.new }
   let(:project) { create(:project, :repository) }
   let(:statistics) { project.statistics }
+  let(:lease_key) { "project_cache_worker:#{project.id}:update_statistics" }
+  let(:lease_timeout) { ProjectCacheWorker::LEASE_TIMEOUT }
 
   describe '#perform' do
     before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain)
-        .and_return(true)
+      stub_exclusive_lease(lease_key, timeout: lease_timeout)
     end
 
     context 'with a non-existing project' do
@@ -63,9 +66,7 @@ describe ProjectCacheWorker do
   describe '#update_statistics' do
     context 'when a lease could not be obtained' do
       it 'does not update the repository size' do
-        allow(worker).to receive(:try_obtain_lease_for)
-          .with(project.id, :update_statistics)
-          .and_return(false)
+        stub_exclusive_lease_taken(lease_key, timeout: lease_timeout)
 
         expect(statistics).not_to receive(:refresh!)
 
@@ -75,9 +76,7 @@ describe ProjectCacheWorker do
 
     context 'when a lease could be obtained' do
       it 'updates the project statistics' do
-        allow(worker).to receive(:try_obtain_lease_for)
-          .with(project.id, :update_statistics)
-          .and_return(true)
+        stub_exclusive_lease(lease_key, timeout: lease_timeout)
 
         expect(statistics).to receive(:refresh!)
           .with(only: %i(repository_size))
diff --git a/spec/workers/project_migrate_hashed_storage_worker_spec.rb b/spec/workers/project_migrate_hashed_storage_worker_spec.rb
index 2e3951e7afc..9551e358af1 100644
--- a/spec/workers/project_migrate_hashed_storage_worker_spec.rb
+++ b/spec/workers/project_migrate_hashed_storage_worker_spec.rb
@@ -1,53 +1,47 @@
 require 'spec_helper'
 
 describe ProjectMigrateHashedStorageWorker, :clean_gitlab_redis_shared_state do
+  include ExclusiveLeaseHelpers
+
   describe '#perform' do
     let(:project) { create(:project, :empty_repo) }
-    let(:pending_delete_project) { create(:project, :empty_repo, pending_delete: true) }
+    let(:lease_key) { "project_migrate_hashed_storage_worker:#{project.id}" }
+    let(:lease_timeout) { ProjectMigrateHashedStorageWorker::LEASE_TIMEOUT }
+
+    it 'skips when project no longer exists' do
+      expect(::Projects::HashedStorageMigrationService).not_to receive(:new)
+
+      subject.perform(-1)
+    end
 
-    context 'when have exclusive lease' do
-      before do
-        lease = subject.lease_for(project.id)
+    it 'skips when project is pending delete' do
+      pending_delete_project = create(:project, :empty_repo, pending_delete: true)
 
-        allow(Gitlab::ExclusiveLease).to receive(:new).and_return(lease)
-        allow(lease).to receive(:try_obtain).and_return(true)
-      end
+      expect(::Projects::HashedStorageMigrationService).not_to receive(:new)
 
-      it 'skips when project no longer exists' do
-        nonexistent_id = 999999999999
+      subject.perform(pending_delete_project.id)
+    end
 
-        expect(::Projects::HashedStorageMigrationService).not_to receive(:new)
-        subject.perform(nonexistent_id)
-      end
+    it 'delegates removal to service class when have exclusive lease' do
+      stub_exclusive_lease(lease_key, 'uuid', timeout: lease_timeout)
 
-      it 'skips when project is pending delete' do
-        expect(::Projects::HashedStorageMigrationService).not_to receive(:new)
+      migration_service = spy
 
-        subject.perform(pending_delete_project.id)
-      end
+      allow(::Projects::HashedStorageMigrationService)
+        .to receive(:new).with(project, subject.logger)
+        .and_return(migration_service)
 
-      it 'delegates removal to service class' do
-        service = double('service')
-        expect(::Projects::HashedStorageMigrationService).to receive(:new).with(project, subject.logger).and_return(service)
-        expect(service).to receive(:execute)
+      subject.perform(project.id)
 
-        subject.perform(project.id)
-      end
+      expect(migration_service).to have_received(:execute)
     end
 
-    context 'when dont have exclusive lease' do
-      before do
-        lease = subject.lease_for(project.id)
-
-        allow(Gitlab::ExclusiveLease).to receive(:new).and_return(lease)
-        allow(lease).to receive(:try_obtain).and_return(false)
-      end
+    it 'skips when dont have lease when dont have exclusive lease' do
+      stub_exclusive_lease_taken(lease_key, timeout: lease_timeout)
 
-      it 'skips when dont have lease' do
-        expect(::Projects::HashedStorageMigrationService).not_to receive(:new)
+      expect(::Projects::HashedStorageMigrationService).not_to receive(:new)
 
-        subject.perform(project.id)
-      end
+      subject.perform(project.id)
     end
   end
 end
diff --git a/spec/workers/propagate_service_template_worker_spec.rb b/spec/workers/propagate_service_template_worker_spec.rb
index b8b65ead9b3..af1fb80a51d 100644
--- a/spec/workers/propagate_service_template_worker_spec.rb
+++ b/spec/workers/propagate_service_template_worker_spec.rb
@@ -1,29 +1,29 @@
 require 'spec_helper'
 
 describe PropagateServiceTemplateWorker do
-  let!(:service_template) do
-    PushoverService.create(
-      template: true,
-      active: true,
-      properties: {
-        device: 'MyDevice',
-        sound: 'mic',
-        priority: 4,
-        user_key: 'asdf',
-        api_key: '123456789'
-      })
-  end
-
-  before do
-    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain)
-      .and_return(true)
-  end
+  include ExclusiveLeaseHelpers
 
   describe '#perform' do
     it 'calls the propagate service with the template' do
-      expect(Projects::PropagateServiceTemplate).to receive(:propagate).with(service_template)
+      template = PushoverService.create(
+        template: true,
+        active: true,
+        properties: {
+          device: 'MyDevice',
+          sound: 'mic',
+          priority: 4,
+          user_key: 'asdf',
+          api_key: '123456789'
+        })
+
+      stub_exclusive_lease("propagate_service_template_worker:#{template.id}",
+        timeout: PropagateServiceTemplateWorker::LEASE_TIMEOUT)
+
+      expect(Projects::PropagateServiceTemplate)
+        .to receive(:propagate)
+        .with(template)
 
-      subject.perform(service_template.id)
+      subject.perform(template.id)
     end
   end
 end
diff --git a/spec/workers/repository_remove_remote_worker_spec.rb b/spec/workers/repository_remove_remote_worker_spec.rb
index 5968c5da3c9..a653f6f926c 100644
--- a/spec/workers/repository_remove_remote_worker_spec.rb
+++ b/spec/workers/repository_remove_remote_worker_spec.rb
@@ -1,44 +1,50 @@
 require 'rails_helper'
 
 describe RepositoryRemoveRemoteWorker do
-  subject(:worker) { described_class.new }
+  include ExclusiveLeaseHelpers
 
   describe '#perform' do
-    let(:remote_name) { 'joe'}
     let!(:project) { create(:project, :repository) }
+    let(:remote_name) { 'joe'}
+    let(:lease_key) { "remove_remote_#{project.id}_#{remote_name}" }
+    let(:lease_timeout) { RepositoryRemoveRemoteWorker::LEASE_TIMEOUT }
 
-    context 'when it cannot obtain lease' do
-      it 'logs error' do
-        allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { nil }
-
-        expect_any_instance_of(Repository).not_to receive(:remove_remote)
-        expect(worker).to receive(:log_error).with('Cannot obtain an exclusive lease. There must be another instance already in execution.')
-
-        worker.perform(project.id, remote_name)
-      end
+    it 'returns nil when project does not exist' do
+      expect(subject.perform(-1, 'remote_name')).to be_nil
     end
 
-    context 'when it gets the lease' do
+    context 'when project exists' do
       before do
-        allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(true)
+        allow(Project)
+          .to receive(:find_by)
+          .with(id: project.id)
+          .and_return(project)
       end
 
-      context 'when project does not exist' do
-        it 'returns nil' do
-          expect(worker.perform(-1, 'remote_name')).to be_nil
-        end
-      end
+      it 'does not remove remote when cannot obtain lease' do
+        stub_exclusive_lease_taken(lease_key, timeout: lease_timeout)
+
+        expect(project.repository)
+          .not_to receive(:remove_remote)
 
-      context 'when project exists' do
-        it 'removes remote from repository' do
-          masterrev = project.repository.find_branch('master').dereferenced_target
+        expect(subject)
+          .to receive(:log_error)
+          .with('Cannot obtain an exclusive lease. There must be another instance already in execution.')
 
-          create_remote_branch(remote_name, 'remote_branch', masterrev)
+        subject.perform(project.id, remote_name)
+      end
 
-          expect_any_instance_of(Repository).to receive(:remove_remote).with(remote_name).and_call_original
+      it 'removes remote from repository when obtain a lease' do
+        stub_exclusive_lease(lease_key, timeout: lease_timeout)
+        masterrev = project.repository.find_branch('master').dereferenced_target
+        create_remote_branch(remote_name, 'remote_branch', masterrev)
 
-          worker.perform(project.id, remote_name)
-        end
+        expect(project.repository)
+          .to receive(:remove_remote)
+          .with(remote_name)
+          .and_call_original
+
+        subject.perform(project.id, remote_name)
       end
     end
   end
@@ -47,6 +53,7 @@ describe RepositoryRemoveRemoteWorker do
     rugged = Gitlab::GitalyClient::StorageSettings.allow_disk_access do
       project.repository.rugged
     end
+
     rugged.references.create("refs/remotes/#{remote_name}/#{branch_name}", target.id)
   end
 end
diff --git a/spec/workers/stuck_ci_jobs_worker_spec.rb b/spec/workers/stuck_ci_jobs_worker_spec.rb
index 2605c14334f..856886e3df5 100644
--- a/spec/workers/stuck_ci_jobs_worker_spec.rb
+++ b/spec/workers/stuck_ci_jobs_worker_spec.rb
@@ -1,14 +1,21 @@
 require 'spec_helper'
 
 describe StuckCiJobsWorker do
+  include ExclusiveLeaseHelpers
+
   let!(:runner) { create :ci_runner }
   let!(:job) { create :ci_build, runner: runner }
-  let(:worker) { described_class.new }
-  let(:exclusive_lease_uuid) { SecureRandom.uuid }
+  let(:trace_lease_key) { "trace:archive:#{job.id}" }
+  let(:trace_lease_uuid) { SecureRandom.uuid }
+  let(:worker_lease_key) { StuckCiJobsWorker::EXCLUSIVE_LEASE_KEY }
+  let(:worker_lease_uuid) { SecureRandom.uuid }
+
+  subject(:worker) { described_class.new }
 
   before do
+    stub_exclusive_lease(worker_lease_key, worker_lease_uuid)
+    stub_exclusive_lease(trace_lease_key, trace_lease_uuid)
     job.update!(status: status, updated_at: updated_at)
-    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(exclusive_lease_uuid)
   end
 
   shared_examples 'job is dropped' do
@@ -44,16 +51,19 @@ describe StuckCiJobsWorker do
 
       context 'when job was not updated for more than 1 day ago' do
         let(:updated_at) { 2.days.ago }
+
         it_behaves_like 'job is dropped'
       end
 
       context 'when job was updated in less than 1 day ago' do
         let(:updated_at) { 6.hours.ago }
+
         it_behaves_like 'job is unchanged'
       end
 
       context 'when job was not updated for more than 1 hour ago' do
         let(:updated_at) { 2.hours.ago }
+
         it_behaves_like 'job is unchanged'
       end
     end
@@ -65,11 +75,14 @@ describe StuckCiJobsWorker do
 
       context 'when job was not updated for more than 1 hour ago' do
         let(:updated_at) { 2.hours.ago }
+
         it_behaves_like 'job is dropped'
       end
 
-      context 'when job was updated in less than 1 hour ago' do
+      context 'when job was updated in less than 1
+       hour ago' do
         let(:updated_at) { 30.minutes.ago }
+
         it_behaves_like 'job is unchanged'
       end
     end
@@ -80,11 +93,13 @@ describe StuckCiJobsWorker do
 
     context 'when job was not updated for more than 1 hour ago' do
       let(:updated_at) { 2.hours.ago }
+
       it_behaves_like 'job is dropped'
     end
 
     context 'when job was updated in less than 1 hour ago' do
       let(:updated_at) { 30.minutes.ago }
+
       it_behaves_like 'job is unchanged'
     end
   end
@@ -93,6 +108,7 @@ describe StuckCiJobsWorker do
     context "when job is #{status}" do
       let(:status) { status }
       let(:updated_at) { 2.days.ago }
+
       it_behaves_like 'job is unchanged'
     end
   end
@@ -119,23 +135,27 @@ describe StuckCiJobsWorker do
     it 'is guard by exclusive lease when executed concurrently' do
       expect(worker).to receive(:drop).at_least(:once).and_call_original
       expect(worker2).not_to receive(:drop)
+
       worker.perform
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(false)
+
+      stub_exclusive_lease_taken(worker_lease_key)
+
       worker2.perform
     end
 
     it 'can be executed in sequence' do
       expect(worker).to receive(:drop).at_least(:once).and_call_original
       expect(worker2).to receive(:drop).at_least(:once).and_call_original
+
       worker.perform
       worker2.perform
     end
 
-    it 'cancels exclusive lease after worker perform' do
-      worker.perform
+    it 'cancels exclusive leases after worker perform' do
+      expect_to_cancel_exclusive_lease(trace_lease_key, trace_lease_uuid)
+      expect_to_cancel_exclusive_lease(worker_lease_key, worker_lease_uuid)
 
-      expect(Gitlab::ExclusiveLease.new(described_class::EXCLUSIVE_LEASE_KEY, timeout: 1.hour))
-        .not_to be_exists
+      worker.perform
     end
   end
 end