Use exclusive lease helper methods in the EE spec files

ee/spec/lib/gitlab/geo/log_cursor/daemon_spec.rb

@@ -2,6 +2,7 @@ require 'spec_helper'
 
 describe Gitlab::Geo::LogCursor::Daemon, :postgresql, :clean_gitlab_redis_shared_state do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   set(:primary) { create(:geo_node, :primary) }
   set(:secondary) { create(:geo_node) }
@@ -37,9 +38,14 @@ describe Gitlab::Geo::LogCursor::Daemon, :postgresql, :clean_gitlab_redis_shared
     end
 
     it 'skips execution if cannot achieve a lease' do
+      lease = stub_exclusive_lease_taken('geo_log_cursor_processed')
+
+      allow(lease).to receive(:try_obtain_with_ttl).and_return({ ttl: 1, uuid: false })
+      allow(lease).to receive(:same_uuid?).and_return(false)
+      allow(Gitlab::Geo::LogCursor::Lease).to receive(:exclusive_lease).and_return(lease)
+
       is_expected.to receive(:exit?).and_return(false, true)
       is_expected.not_to receive(:run_once!)
-      expect_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain_with_ttl).and_return({ ttl: 1, uuid: false })
 
       daemon.run!
     end

ee/spec/lib/gitlab/geo/log_cursor/lease_spec.rb

 require 'spec_helper'
 
 describe Gitlab::Geo::LogCursor::Lease, :clean_gitlab_redis_shared_state do
+  include ExclusiveLeaseHelpers
+
   describe '.exclusive_lease' do
     it 'returns an exclusive lease instance' do
       expect(described_class.send(:exclusive_lease)).to be_an_instance_of(Gitlab::ExclusiveLease)
@@ -8,15 +10,20 @@ describe Gitlab::Geo::LogCursor::Lease, :clean_gitlab_redis_shared_state do
   end
 
   describe '.renew!' do
+    let(:lease) { stub_exclusive_lease(described_class::LEASE_KEY, renew: true) }
+
+    before do
+      allow(described_class).to receive(:exclusive_lease).and_return(lease)
+    end
+
     it 'returns an exclusive lease instance' do
-      expect_any_instance_of(Gitlab::ExclusiveLease).to receive(:renew)
+      expect(lease).to receive(:renew)
 
       described_class.renew!
     end
 
     it 'logs with the correct caller class' do
       stub_const("Gitlab::Geo::LogCursor::Logger::PID", 111)
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:renew).and_return(true)
 
       expect(::Gitlab::Logger).to receive(:debug).with(pid: 111,
                                                        class: 'Gitlab::Geo::LogCursor::Lease',
@@ -44,7 +51,12 @@ describe Gitlab::Geo::LogCursor::Lease, :clean_gitlab_redis_shared_state do
     end
 
     it 'returns > 0 if there was an error' do
-      expect(Gitlab::ExclusiveLease).to receive(:cancel)
+      lease = stub_exclusive_lease(described_class::LEASE_KEY, 'uuid')
+
+      allow(lease).to receive(:try_obtain_with_ttl).and_return({ ttl: 0, uuid: 'uuid' })
+      allow(described_class).to receive(:exclusive_lease).and_return(lease)
+
+      expect_to_cancel_exclusive_lease(described_class::LEASE_KEY, 'uuid')
 
       result = described_class.try_obtain_with_ttl { raise StandardError }
 

ee/spec/services/geo/file_download_service_spec.rb

@@ -2,21 +2,26 @@ require 'spec_helper'
 
 describe Geo::FileDownloadService do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   set(:primary)  { create(:geo_node, :primary) }
   set(:secondary) { create(:geo_node) }
 
   before do
     stub_current_geo_node(secondary)
-
-    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(true)
   end
 
   shared_examples_for 'a service that downloads the file and registers the sync result' do |file_type|
     let(:download_service) { described_class.new(file_type, file.id) }
     let(:registry) { file_type == 'job_artifact' ? Geo::JobArtifactRegistry : Geo::FileRegistry }
+
     subject(:execute!) { download_service.execute }
 
+    before do
+      stub_exclusive_lease("file_download_service:#{file_type}:#{file.id}",
+        timeout: Geo::FileDownloadService::LEASE_TIMEOUT)
+    end
+
     context 'for a new file' do
       context 'when the downloader fails before attempting a transfer' do
         it 'logs that the download failed before attempting a transfer' do

ee/spec/services/geo/file_registry_removal_service_spec.rb

@@ -2,18 +2,18 @@ require 'spec_helper'
 
 describe Geo::FileRegistryRemovalService do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   set(:secondary) { create(:geo_node) }
 
   before do
     stub_current_geo_node(secondary)
-
-    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(true)
   end
 
   describe '#execute' do
     it 'delegates log_error to the Geo logger' do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(false)
+      stub_exclusive_lease_taken("file_registry_removal_service:lfs:99")
+
       expect(Gitlab::Geo::Logger).to receive(:error)
 
       described_class.new(:lfs, 99).execute
@@ -22,6 +22,11 @@ describe Geo::FileRegistryRemovalService do
     shared_examples 'removes' do
       subject(:service) { described_class.new(file_registry.file_type, file_registry.file_id) }
 
+      before do
+        stub_exclusive_lease("file_registry_removal_service:#{file_registry.file_type}:#{file_registry.file_id}",
+          timeout: Geo::FileRegistryRemovalService::LEASE_TIMEOUT)
+      end
+
       it 'file from disk' do
         expect do
           service.execute
@@ -38,6 +43,11 @@ describe Geo::FileRegistryRemovalService do
     shared_examples 'removes artifact' do
       subject(:service) { described_class.new('job_artifact', registry.artifact_id) }
 
+      before do
+        stub_exclusive_lease("file_registry_removal_service:job_artifact:#{registry.artifact_id}",
+          timeout: Geo::FileRegistryRemovalService::LEASE_TIMEOUT)
+      end
+
       it 'file from disk' do
         expect do
           service.execute

ee/spec/services/geo/repository_sync_service_spec.rb

@@ -2,13 +2,15 @@ require 'spec_helper'
 
 describe Geo::RepositorySyncService do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   set(:primary) { create(:geo_node, :primary) }
   set(:secondary) { create(:geo_node) }
-
-  let(:lease) { double(try_obtain: true) }
   set(:project) { create(:project_empty_repo) }
+
   let(:repository) { project.repository }
+  let(:lease_key) { "geo_sync_service:repository:#{project.id}" }
+  let(:lease_uuid) { 'uuid'}
 
   subject { described_class.new(project) }
 
@@ -23,7 +25,8 @@ describe Geo::RepositorySyncService do
     let(:url_to_repo) { "#{primary.url}#{project.full_path}.git" }
 
     before do
-      allow(subject).to receive(:exclusive_lease).and_return(lease)
+      stub_exclusive_lease(lease_key, lease_uuid)
+      stub_exclusive_lease("geo_project_housekeeping:#{project.id}")
 
       allow_any_instance_of(Repository).to receive(:fetch_as_mirror)
         .and_return(true)
@@ -47,8 +50,7 @@ describe Geo::RepositorySyncService do
     end
 
     it 'returns the lease when succeed' do
-      expect(Gitlab::ExclusiveLease).to receive(:cancel).once.with(
-        subject.__send__(:lease_key), anything).and_call_original
+      expect_to_cancel_exclusive_lease(lease_key, lease_uuid)
 
       subject.execute
     end
@@ -64,14 +66,13 @@ describe Geo::RepositorySyncService do
         .with(url_to_repo, remote_name: 'geo', forced: true)
         .and_raise(Gitlab::Shell::Error)
 
-      expect(Gitlab::ExclusiveLease).to receive(:cancel).once.with(
-        subject.__send__(:lease_key), anything).and_call_original
+      expect_to_cancel_exclusive_lease(lease_key, lease_uuid)
 
       subject.execute
     end
 
     it 'does not fetch project repository if cannot obtain a lease' do
-      allow(lease).to receive(:try_obtain) { false }
+      stub_exclusive_lease_taken(lease_key)
 
       expect(repository).not_to receive(:fetch_as_mirror)
 

ee/spec/services/geo/wiki_sync_service_spec.rb

@@ -2,13 +2,15 @@ require 'spec_helper'
 
 RSpec.describe Geo::WikiSyncService do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   set(:primary) { create(:geo_node, :primary) }
   set(:secondary) { create(:geo_node) }
+  set(:project) { create(:project_empty_repo) }
 
-  let(:lease) { double(try_obtain: true, lease_key: anything ) }
-  let(:project) { create(:project_empty_repo) }
   let(:repository) { project.wiki.repository }
+  let(:lease_key) { "geo_sync_service:wiki:#{project.id}" }
+  let(:lease_uuid) { 'uuid'}
 
   subject { described_class.new(project) }
 
@@ -23,7 +25,7 @@ RSpec.describe Geo::WikiSyncService do
     let(:url_to_repo) { "#{primary.url}#{project.full_path}.wiki.git" }
 
     before do
-      allow(subject).to receive(:exclusive_lease).and_return(lease)
+      stub_exclusive_lease(lease_key, lease_uuid)
 
       allow_any_instance_of(Repository).to receive(:fetch_as_mirror)
         .and_return(true)
@@ -39,8 +41,7 @@ RSpec.describe Geo::WikiSyncService do
     end
 
     it 'releases lease' do
-      expect(Gitlab::ExclusiveLease).to receive(:cancel).at_least(:once).with(
-        subject.__send__(:lease_key), anything).and_call_original
+      expect_to_cancel_exclusive_lease(lease_key, lease_uuid)
 
       subject.execute
     end
@@ -52,7 +53,7 @@ RSpec.describe Geo::WikiSyncService do
     end
 
     it 'does not fetch wiki repository if cannot obtain a lease' do
-      allow(lease).to receive(:try_obtain) { false }
+      stub_exclusive_lease_taken(lease_key)
 
       expect(repository).not_to receive(:fetch_as_mirror)
 

ee/spec/workers/geo/prune_event_log_worker_spec.rb

@@ -2,13 +2,18 @@ require 'spec_helper'
 
 describe Geo::PruneEventLogWorker, :geo do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   subject(:worker) { described_class.new }
+
   set(:primary) { create(:geo_node, :primary) }
   set(:secondary) { create(:geo_node) }
 
+  let(:lease_key) { 'geo/prune_event_log_worker' }
+  let(:lease_timeout) { Geo::PruneEventLogWorker::LEASE_TIMEOUT }
+
   before do
-    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(true)
+    stub_exclusive_lease(lease_key, timeout: lease_timeout)
   end
 
   describe '#perform' do
@@ -30,7 +35,7 @@ describe Geo::PruneEventLogWorker, :geo do
       end
 
       it 'logs error when it cannot obtain lease' do
-        allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { nil }
+        stub_exclusive_lease_taken(lease_key, timeout: lease_timeout)
 
         expect(worker).to receive(:log_error).with(/^Cannot obtain an exclusive lease/)
 

ee/spec/workers/geo/repositories_clean_up_worker_spec.rb

@@ -2,10 +2,12 @@ require 'spec_helper'
 
 describe Geo::RepositoriesCleanUpWorker do
   describe '#perform' do
+    include ExclusiveLeaseHelpers
+
     let(:geo_node) { create(:geo_node) }
 
     before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { true }
+      stub_exclusive_lease
     end
 
     context 'when node has namespace restrictions' do
@@ -68,7 +70,7 @@ describe Geo::RepositoriesCleanUpWorker do
     end
 
     it 'does not perform GeoRepositoryDestroyWorker when cannnot obtain a lease' do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { false }
+      stub_exclusive_lease_taken
 
       expect(GeoRepositoryDestroyWorker).not_to receive(:perform_async)
 

ee/spec/workers/geo/repository_shard_sync_worker_spec.rb

@@ -4,6 +4,7 @@ require 'spec_helper'
 # can't see changes inside a transaction of a different connection.
 describe Geo::RepositoryShardSyncWorker, :geo, :delete, :clean_gitlab_redis_cache do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   let!(:primary) { create(:geo_node, :primary) }
   let!(:secondary) { create(:geo_node) }
@@ -25,8 +26,7 @@ describe Geo::RepositoryShardSyncWorker, :geo, :delete, :clean_gitlab_redis_cach
     end
 
     before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { true }
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:renew) { true }
+      stub_exclusive_lease(renew: true)
 
       Gitlab::ShardHealthCache.update([shard_name])
     end
@@ -48,6 +48,7 @@ describe Geo::RepositoryShardSyncWorker, :geo, :delete, :clean_gitlab_redis_cach
 
     it 'does not perform Geo::ProjectSyncWorker when shard becomes unhealthy' do
       Gitlab::ShardHealthCache.update([])
+
       expect(Geo::ProjectSyncWorker).not_to receive(:perform_async)
 
       subject.perform(shard_name)

ee/spec/workers/geo/repository_verification/primary/shard_worker_spec.rb

@@ -2,6 +2,7 @@ require 'spec_helper'
 
 describe Geo::RepositoryVerification::Primary::ShardWorker, :postgresql, :clean_gitlab_redis_cache do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   let!(:primary)   { create(:geo_node, :primary) }
   let(:shard_name) { Gitlab.config.repositories.storages.keys.first }
@@ -13,8 +14,7 @@ describe Geo::RepositoryVerification::Primary::ShardWorker, :postgresql, :clean_
 
   describe '#perform' do
     before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { true }
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:renew) { true }
+      stub_exclusive_lease(renew: true)
 
       Gitlab::ShardHealthCache.update([shard_name])
     end

ee/spec/workers/geo/repository_verification/primary/single_worker_spec.rb

@@ -2,6 +2,7 @@ require 'spec_helper'
 
 describe Geo::RepositoryVerification::Primary::SingleWorker, :postgresql, :clean_gitlab_redis_cache do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   set(:project) { create(:project) }
 
@@ -11,13 +12,10 @@ describe Geo::RepositoryVerification::Primary::SingleWorker, :postgresql, :clean
 
   before do
     stub_current_geo_node(primary)
+    stub_exclusive_lease
   end
 
   describe '#perform' do
-    before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { true }
-    end
-
     it 'does not calculate the checksum when not running on a primary' do
       stub_project_repository(project, repository)
       stub_wiki_repository(project.wiki, wiki)

ee/spec/workers/geo/repository_verification/secondary/shard_worker_spec.rb

@@ -2,6 +2,7 @@ require 'spec_helper'
 
 describe Geo::RepositoryVerification::Secondary::ShardWorker, :postgresql, :clean_gitlab_redis_cache do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   let!(:secondary) { create(:geo_node) }
   let(:shard_name) { Gitlab.config.repositories.storages.keys.first }
@@ -16,8 +17,7 @@ describe Geo::RepositoryVerification::Secondary::ShardWorker, :postgresql, :clea
 
   describe '#perform' do
     before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { true }
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:renew) { true }
+      stub_exclusive_lease(renew: true)
 
       Gitlab::ShardHealthCache.update([shard_name])
     end

ee/spec/workers/geo/repository_verification/secondary/single_worker_spec.rb

@@ -2,6 +2,7 @@ require 'spec_helper'
 
 describe Geo::RepositoryVerification::Secondary::SingleWorker, :postgresql, :clean_gitlab_redis_cache do
   include ::EE::GeoHelpers
+  include ExclusiveLeaseHelpers
 
   let!(:secondary) { create(:geo_node) }
   let(:project)    { create(:project, :repository, :wiki_repo) }
@@ -9,13 +10,10 @@ describe Geo::RepositoryVerification::Secondary::SingleWorker, :postgresql, :cle
 
   before do
     stub_current_geo_node(secondary)
+    stub_exclusive_lease
   end
 
   describe '#perform' do
-    before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain) { true }
-    end
-
     it 'does not calculate the checksum when not running on a secondary' do
       allow(Gitlab::Geo).to receive(:secondary?) { false }
 

ee/spec/workers/repository_update_mirror_worker_spec.rb

 require 'rails_helper'
 
 describe RepositoryUpdateMirrorWorker do
+  include ExclusiveLeaseHelpers
+
   describe '#perform' do
     let(:jid) { '12345678' }
     let!(:project) { create(:project) }
     let!(:import_state) { create(:import_state, :mirror, :scheduled, project: project) }
 
     before do
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(true)
       allow(subject).to receive(:jid).and_return(jid)
     end
 
@@ -56,6 +57,10 @@ describe RepositoryUpdateMirrorWorker do
       end
 
       context 'when we obtain the lease' do
+        before do
+          allow(stub_exclusive_lease).to receive(:exists?).and_return(false)
+        end
+
         it 'performs UpdateAllMirrorsWorker when reschedule_immediately? returns true' do
           allow(Gitlab::Mirror).to receive(:reschedule_immediately?).and_return(true)
 
@@ -74,7 +79,7 @@ describe RepositoryUpdateMirrorWorker do
       end
 
       it 'does not perform UpdateAllMirrorsWorker when we cannot obtain the lease' do
-        allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(false)
+        stub_exclusive_lease_taken
 
         expect(UpdateAllMirrorsWorker).not_to receive(:perform_async)
 
@@ -82,7 +87,7 @@ describe RepositoryUpdateMirrorWorker do
       end
 
       it 'does not perform UpdateAllMirrorsWorker when the lease already exists' do
-        allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:exists?).and_return(true)
+        stub_exclusive_lease_taken
 
         expect(UpdateAllMirrorsWorker).not_to receive(:perform_async)
 

ee/spec/workers/update_all_mirrors_worker_spec.rb

 require 'spec_helper'
 
 describe UpdateAllMirrorsWorker do
+  include ExclusiveLeaseHelpers
+
   subject(:worker) { described_class.new }
 
   before do
-    allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(true)
+    stub_exclusive_lease
   end
 
   describe '#perform' do
     it 'does not execute if cannot get the lease' do
-      create(:project, :mirror)
-
-      allow_any_instance_of(Gitlab::ExclusiveLease).to receive(:try_obtain).and_return(false)
+      stub_exclusive_lease_taken
 
       expect(worker).not_to receive(:schedule_mirrors!)
 

spec/support/helpers/exclusive_lease_helpers.rb

 module ExclusiveLeaseHelpers
-  def stub_exclusive_lease(key = nil, uuid = 'uuid', timeout: nil)
+  def stub_exclusive_lease(key = nil, uuid = 'uuid', renew: false, timeout: nil)
     key     ||= instance_of(String)
     timeout ||= instance_of(Integer)
-    lease   = instance_double(Gitlab::ExclusiveLease, try_obtain: uuid)
+
+    lease = instance_double(
+      Gitlab::ExclusiveLease,
+      try_obtain: uuid,
+      exists?: true,
+      renew: renew
+    )
 
     allow(Gitlab::ExclusiveLease)
       .to receive(:new)