fix exclusive lease specs fro hashed storage migration

2 files changed, 15 insertions, 9 deletions

diff --git a/app/workers/project_migrate_hashed_storage_worker.rb b/app/workers/project_migrate_hashed_storage_worker.rb
index e3ecd6bc950..127aa6b9d7d 100644
--- a/app/workers/project_migrate_hashed_storage_worker.rb
+++ b/app/workers/project_migrate_hashed_storage_worker.rb
@@ -8,23 +8,23 @@ class ProjectMigrateHashedStorageWorker
     project = Project.find_by(id: project_id)
     return if project.nil? || project.pending_delete?
 
-    uuid = try_obtain_lease_for(project_id)
+    uuid = lease_for(project_id).try_obtain
     if uuid
       ::Projects::HashedStorageMigrationService.new(project, logger).execute
     else
       false
     end
   rescue => ex
-    cancel_lease_for(project_id, uuid)
+    cancel_lease_for(project_id, uuid) if uuid
     raise ex
   end
 
-  private
-
-  def try_obtain_lease_for(project_id)
-    Gitlab::ExclusiveLease.new(lease_key(project_id), timeout: LEASE_TIMEOUT).try_obtain
+  def lease_for(project_id)
+    Gitlab::ExclusiveLease.new(lease_key(project_id), timeout: LEASE_TIMEOUT)
   end
 
+  private
+
   def lease_key(project_id)
     "project_migrate_hashed_storage_worker:#{project_id}"
   end
diff --git a/spec/workers/project_migrate_hashed_storage_worker_spec.rb b/spec/workers/project_migrate_hashed_storage_worker_spec.rb
index 8cacdfa7173..2e3951e7afc 100644
--- a/spec/workers/project_migrate_hashed_storage_worker_spec.rb
+++ b/spec/workers/project_migrate_hashed_storage_worker_spec.rb
@@ -1,13 +1,16 @@
 require 'spec_helper'
 
-describe ProjectMigrateHashedStorageWorker do
+describe ProjectMigrateHashedStorageWorker, :clean_gitlab_redis_shared_state do
   describe '#perform' do
     let(:project) { create(:project, :empty_repo) }
     let(:pending_delete_project) { create(:project, :empty_repo, pending_delete: true) }
 
     context 'when have exclusive lease' do
       before do
-        allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(true)
+        lease = subject.lease_for(project.id)
+
+        allow(Gitlab::ExclusiveLease).to receive(:new).and_return(lease)
+        allow(lease).to receive(:try_obtain).and_return(true)
       end
 
       it 'skips when project no longer exists' do
@@ -34,7 +37,10 @@ describe ProjectMigrateHashedStorageWorker do
 
     context 'when dont have exclusive lease' do
       before do
-        allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(false)
+        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' do