Extract DB load balancing setup in specs

Adds an Rspec tag that enables DB load balancing for the specified
context

spec/lib/gitlab/checks/matching_merge_request_spec.rb

@@ -49,13 +49,11 @@ RSpec.describe Gitlab::Checks::MatchingMergeRequest do
       end
     end
 
-    context 'with load balancing enabled', :request_store, :redis do
+    context 'with load balancing enabled', :db_load_balancing do
       let(:session) { ::Gitlab::Database::LoadBalancing::Session.current }
       let(:all_caught_up) { true }
 
       before do
-        allow(ActiveRecord::Base).to receive(:load_balancing_proxy)
-        expect(::Gitlab::Database::LoadBalancing).to receive(:enable?).at_least(:once).and_return(true)
         allow(::Gitlab::Database::LoadBalancing::Sticking).to receive(:all_caught_up?).and_return(all_caught_up)
 
         expect(::Gitlab::Database::LoadBalancing::Sticking).to receive(:select_valid_host).with(:project, project.id).and_call_original

spec/lib/gitlab/database/load_balancing_spec.rb

@@ -364,7 +364,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
   # - In each test, we listen to the SQL queries (via sql.active_record
   # instrumentation) while triggering real queries from the defined model.
   # - We assert the desinations (replica/primary) of the queries in order.
-  describe 'LoadBalancing integration tests', :delete do
+  describe 'LoadBalancing integration tests', :db_load_balancing, :delete do
     before(:all) do
       ActiveRecord::Schema.define do
         create_table :load_balancing_test, force: true do |t|
@@ -379,30 +379,14 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       end
     end
 
-    shared_context 'LoadBalancing setup' do
-      let(:development_db_config) { ActiveRecord::Base.configurations.configs_for(env_name: 'development').first.configuration_hash }
-      let(:hosts) { [development_db_config[:host]] }
-      let(:model) do
-        Class.new(ApplicationRecord) do
-          self.table_name = "load_balancing_test"
-        end
+    let(:model) do
+      Class.new(ApplicationRecord) do
+        self.table_name = "load_balancing_test"
       end
+    end
 
-      before do
-        # Preloading testing class
-        model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy
-
-        # Setup load balancing
-        allow(ActiveRecord::Base).to receive(:load_balancing_proxy).and_return(
-          ::Gitlab::Database::LoadBalancing::ConnectionProxy.new(hosts)
-        )
-
-        original_db_config = Gitlab::Database.main.config
-        modified_db_config = original_db_config.merge(load_balancing: { hosts: hosts })
-        allow(Gitlab::Database.main).to receive(:config).and_return(modified_db_config)
-
-        ::Gitlab::Database::LoadBalancing::Session.clear_session
-      end
+    before do
+      model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy
     end
 
     where(:queries, :include_transaction, :expected_results) do
@@ -713,8 +697,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     end
 
     with_them do
-      include_context 'LoadBalancing setup'
-
       it 'redirects queries to the right roles' do
         roles = []
 
@@ -783,8 +765,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       end
 
       with_them do
-        include_context 'LoadBalancing setup'
-
         it 'redirects queries to the right roles' do
           roles = []
 
@@ -803,8 +783,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     end
 
     context 'a write inside a transaction inside fallback_to_replicas_for_ambiguous_queries block' do
-      include_context 'LoadBalancing setup'
-
       it 'raises an exception' do
         expect do
           ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do

spec/lib/gitlab/sidekiq_logging/structured_logger_spec.rb

@@ -291,12 +291,7 @@ RSpec.describe Gitlab::SidekiqLogging::StructuredLogger do
         include_examples 'performs database queries'
       end
 
-      context 'when load balancing is enabled' do
-        before do
-          allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-          allow(ActiveRecord::Base).to receive(:load_balancing_proxy)
-        end
-
+      context 'when load balancing is enabled', :db_load_balancing do
         let(:db_config_name) { ::Gitlab::Database.db_config_name(ApplicationRecord.connection) }
 
         let(:expected_end_payload_with_db) do

spec/models/ci/build_spec.rb

@@ -344,13 +344,9 @@ RSpec.describe Ci::Build do
   end
 
   describe '#stick_build_if_status_changed' do
-    it 'sticks the build if the status changed' do
+    it 'sticks the build if the status changed', :db_load_balancing do
       job = create(:ci_build, :pending)
 
-      allow(ActiveRecord::Base).to receive(:load_balancing_proxy)
-      allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-        .and_return(true)
-
       expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:stick)
         .with(:build, job.id)
 

spec/models/project_feature_usage_spec.rb

@@ -133,12 +133,9 @@ RSpec.describe ProjectFeatureUsage, type: :model do
 
       subject { project.feature_usage }
 
-      context 'database load balancing is configured' do
+      context 'database load balancing is configured', :db_load_balancing do
         before do
-          # Do not pollute AR for other tests, but rather simulate effect of configure_proxy.
-          allow(ActiveRecord::Base).to receive(:load_balancing_proxy=)
-          proxy = ::Gitlab::Database::LoadBalancing.configure_proxy
-          allow(ActiveRecord::Base).to receive(:connection).and_return(proxy)
+          allow(ActiveRecord::Base).to receive(:connection).and_return(::Gitlab::Database::LoadBalancing.proxy)
 
           ::Gitlab::Database::LoadBalancing::Session.clear_session
         end

spec/services/ci/register_job_service_spec.rb

@@ -14,14 +14,11 @@ module Ci
     let!(:pending_job) { create(:ci_build, :pending, :queued, pipeline: pipeline) }
 
     describe '#execute' do
-      context 'checks database loadbalancing stickiness' do
+      context 'checks database loadbalancing stickiness', :db_load_balancing do
         subject { described_class.new(shared_runner).execute }
 
         before do
           project.update!(shared_runners_enabled: false)
-
-          allow(ActiveRecord::Base).to receive(:load_balancing_proxy)
-          allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
         end
 
         it 'result is valid if replica did caught-up' do

spec/services/users/activity_service_spec.rb

@@ -85,18 +85,15 @@ RSpec.describe Users::ActivityService do
     end
   end
 
-  context 'with DB Load Balancing', :request_store, :redis, :clean_gitlab_redis_shared_state do
+  context 'with DB Load Balancing' do
     let(:user) { create(:user, last_activity_on: last_activity_on) }
 
     context 'when last activity is in the past' do
       let(:user) { create(:user, last_activity_on: Date.today - 1.week) }
 
-      context 'database load balancing is configured' do
+      context 'database load balancing is configured', :db_load_balancing do
         before do
-          # Do not pollute AR for other tests, but rather simulate effect of configure_proxy.
-          allow(ActiveRecord::Base).to receive(:load_balancing_proxy=)
-          proxy = ::Gitlab::Database::LoadBalancing.configure_proxy
-          allow(ActiveRecord::Base).to receive(:connection).and_return(proxy)
+          allow(ActiveRecord::Base).to receive(:connection).and_return(::Gitlab::Database::LoadBalancing.proxy)
         end
 
         let(:service) do

spec/support/database_load_balancing.rb

+# frozen_string_literal: true
+
+RSpec.configure do |config|
+  config.before(:each, :db_load_balancing) do
+    allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+
+    proxy = ::Gitlab::Database::LoadBalancing::ConnectionProxy.new([Gitlab::Database.main.config['host']])
+
+    allow(ActiveRecord::Base).to receive(:load_balancing_proxy).and_return(proxy)
+
+    ::Gitlab::Database::LoadBalancing::Session.clear_session
+    redis_shared_state_cleanup!
+  end
+
+  config.after(:each, :db_load_balancing) do
+    ::Gitlab::Database::LoadBalancing::Session.clear_session
+    redis_shared_state_cleanup!
+  end
+end

spec/workers/authorized_project_update/user_refresh_from_replica_worker_spec.rb

@@ -44,12 +44,7 @@ RSpec.describe AuthorizedProjectUpdate::UserRefreshFromReplicaWorker do
       end
     end
 
-    context 'with load balancing enabled' do
-      before do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-        allow(ActiveRecord::Base).to receive(:load_balancing_proxy)
-      end
-
+    context 'with load balancing enabled', :db_load_balancing do
       it 'reads from the replica database' do
         expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original
 

spec/workers/container_expiration_policy_worker_spec.rb

@@ -156,12 +156,7 @@ RSpec.describe ContainerExpirationPolicyWorker do
         subject
       end
 
-      context 'with load balancing enabled' do
-        before do
-          allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-          allow(ActiveRecord::Base).to receive(:load_balancing_proxy)
-        end
-
+      context 'with load balancing enabled', :db_load_balancing do
         it 'reads the counts from the replica' do
           expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original