Merge branch 'qmnguyen0711/fix-a-flakky-database-test' into 'master'

Flaky test: Release DB proxy's host when clearing configuration

See merge request gitlab-org/gitlab!57320

ee/lib/gitlab/database/load_balancing.rb

@@ -137,12 +137,6 @@ module Gitlab
         ActiveRecord::Base.singleton_class.prepend(ActiveRecordProxy)
       end
 
-      # Clear configuration
-      def self.clear_configuration
-        @proxy = nil
-        remove_instance_variable(:@feature_available) if defined?(@feature_available)
-      end
-
       def self.active_record_models
         ActiveRecord::Base.descendants
       end

ee/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb

@@ -87,7 +87,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::ConnectionProxy do
       end
     end
 
-    it 'inserts data in bulk', quarantine: 'https://gitlab.com/gitlab-org/gitlab/-/issues/325786' do
+    it 'inserts data in bulk' do
       expect(model_class).to receive(:connection)
         .at_least(:once)
         .and_return(proxy)

ee/spec/lib/gitlab/database/load_balancing_spec.rb

@@ -3,6 +3,8 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::LoadBalancing do
+  include_context 'clear DB Load Balancing configuration'
+
   describe '.proxy' do
     context 'when configured' do
       before do
@@ -10,10 +12,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
         subject.configure_proxy
       end
 
-      after do
-        subject.clear_configuration
-      end
-
       it 'returns the connection proxy' do
         expect(subject.proxy).to be_an_instance_of(subject::ConnectionProxy)
       end
@@ -133,7 +131,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     let!(:license) { create(:license, plan: ::License::PREMIUM_PLAN) }
 
     before do
-      subject.clear_configuration
+      clear_load_balancing_configuration
       allow(described_class).to receive(:hosts).and_return(%w(foo))
     end
 
@@ -176,7 +174,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     context 'without a license' do
       before do
         License.destroy_all # rubocop: disable Cop/DestroyAll
-        subject.clear_configuration
+        clear_load_balancing_configuration
       end
 
       it 'is disabled' do
@@ -207,10 +205,13 @@ RSpec.describe Gitlab::Database::LoadBalancing do
   describe '.configured?' do
     let!(:license) { create(:license, plan: ::License::PREMIUM_PLAN) }
 
+    before do
+      clear_load_balancing_configuration
+    end
+
     it 'returns true when Sidekiq is being used' do
       allow(described_class).to receive(:hosts).and_return(%w(foo))
       allow(Gitlab::Runtime).to receive(:sidekiq?).and_return(true)
-
       expect(described_class.configured?).to eq(true)
     end
 
@@ -238,7 +239,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
     context 'without a license' do
       before do
         License.destroy_all # rubocop: disable Cop/DestroyAll
-        subject.clear_configuration
+        clear_load_balancing_configuration
       end
 
       it 'is not configured' do
@@ -248,10 +249,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
   end
 
   describe '.configure_proxy' do
-    after do
-      described_class.clear_configuration
-    end
-
     it 'configures the connection proxy' do
       allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
 
@@ -370,10 +367,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
         subject.configure_proxy(proxy)
       end
 
-      after do
-        subject.clear_configuration
-      end
-
       context 'when the load balancer returns :replica' do
         it 'returns :replica' do
           allow(load_balancer).to receive(:db_role_for_connection).and_return(:replica)
@@ -445,7 +438,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
         model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy
 
         # Setup load balancing
-        subject.clear_configuration
+        clear_load_balancing_configuration
         allow(ActiveRecord::Base.singleton_class).to receive(:prepend)
         subject.configure_proxy(::Gitlab::Database::LoadBalancing::ConnectionProxy.new(hosts))
         allow(ActiveRecord::Base.configurations[Rails.env])
@@ -454,10 +447,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
           .and_return('hosts' => hosts)
         ::Gitlab::Database::LoadBalancing::Session.clear_session
       end
-
-      after do
-        subject.clear_configuration
-      end
     end
 
     where(:queries, :include_transaction, :expected_results) do

ee/spec/models/project_feature_usage_spec.rb

@@ -3,6 +3,8 @@
 require 'spec_helper'
 
 RSpec.describe ProjectFeatureUsage, :request_store do
+  include_context 'clear DB Load Balancing configuration'
+
   describe '#log_jira_dvcs_integration_usage' do
     let!(:project) { create(:project) }
 
@@ -17,10 +19,6 @@ RSpec.describe ProjectFeatureUsage, :request_store do
         ::Gitlab::Database::LoadBalancing::Session.clear_session
       end
 
-      after do
-        ::Gitlab::Database::LoadBalancing.clear_configuration
-      end
-
       it 'logs Jira DVCS Cloud last sync' do
         freeze_time do
           subject.log_jira_dvcs_integration_usage

ee/spec/services/users/activity_service_spec.rb

@@ -3,6 +3,8 @@
 require 'spec_helper'
 
 RSpec.describe Users::ActivityService, '#execute', :request_store, :redis, :clean_gitlab_redis_shared_state do
+  include_context 'clear DB Load Balancing configuration'
+
   let(:user) { create(:user, last_activity_on: last_activity_on) }
 
   context 'when last activity is in the past' do
@@ -33,10 +35,6 @@ RSpec.describe Users::ActivityService, '#execute', :request_store, :redis, :clea
         expect(::Gitlab::Database::LoadBalancing::Session.current).to be_performed_write
         expect(::Gitlab::Database::LoadBalancing::Session.current).not_to be_using_primary
       end
-
-      after do
-        ::Gitlab::Database::LoadBalancing.clear_configuration
-      end
     end
 
     context 'database load balancing is not configured' do

spec/support/shared_contexts/load_balancing_configuration_shared_context.rb

+# frozen_string_literal: true
+
+RSpec.shared_context 'clear DB Load Balancing configuration' do
+  def clear_load_balancing_configuration
+    proxy = ::Gitlab::Database::LoadBalancing.instance_variable_get(:@proxy)
+    proxy.load_balancer.release_host if proxy
+    ::Gitlab::Database::LoadBalancing.instance_variable_set(:@proxy, nil)
+
+    ::Gitlab::Database::LoadBalancing.remove_instance_variable(:@feature_available) if ::Gitlab::Database::LoadBalancing.instance_variable_defined?(:@feature_available)
+
+    ::Gitlab::Database::LoadBalancing::Session.clear_session
+  end
+
+  around do |example|
+    clear_load_balancing_configuration
+
+    example.run
+
+    clear_load_balancing_configuration
+  end
+end