1 files changed, 43 insertions, 7 deletions
diff --git a/spec/lib/gitlab/database_spec.rb b/spec/lib/gitlab/database_spec.rb
index c67b5af5e3c..a9a8d5e6314 100644
--- a/spec/lib/gitlab/database_spec.rb
+++ b/spec/lib/gitlab/database_spec.rb
@@ -15,6 +15,22 @@ RSpec.describe Gitlab::Database do
     end
   end
 
+  describe '.default_pool_size' do
+    before do
+      allow(Gitlab::Runtime).to receive(:max_threads).and_return(7)
+    end
+
+    it 'returns the max thread size plus a fixed headroom of 10' do
+      expect(described_class.default_pool_size).to eq(17)
+    end
+
+    it 'returns the max thread size plus a DB_POOL_HEADROOM if this env var is present' do
+      stub_env('DB_POOL_HEADROOM', '7')
+
+      expect(described_class.default_pool_size).to eq(14)
+    end
+  end
+
   describe '.has_config?' do
     context 'two tier database config' do
       before do
@@ -139,23 +155,43 @@ RSpec.describe Gitlab::Database do
     it { expect(described_class.nulls_first_order('column', 'DESC')).to eq 'column DESC NULLS FIRST'}
   end
 
-  describe '.db_config_name' do
-    it 'returns the db_config name for the connection' do
-      connection = ActiveRecord::Base.connection
+  describe '.db_config_for_connection' do
+    context 'when the regular connection is used' do
+      it 'returns db_config' do
+        connection = ActiveRecord::Base.retrieve_connection
 
-      expect(described_class.db_config_name(connection)).to be_a(String)
-      expect(described_class.db_config_name(connection)).to eq(connection.pool.db_config.name)
+        expect(described_class.db_config_for_connection(connection)).to eq(connection.pool.db_config)
+      end
+    end
+
+    context 'when the connection is LoadBalancing::ConnectionProxy' do
+      it 'returns nil' do
+        lb_config = ::Gitlab::Database::LoadBalancing::Configuration.new(ActiveRecord::Base)
+        lb = ::Gitlab::Database::LoadBalancing::LoadBalancer.new(lb_config)
+        proxy = ::Gitlab::Database::LoadBalancing::ConnectionProxy.new(lb)
+
+        expect(described_class.db_config_for_connection(proxy)).to be_nil
+      end
     end
 
     context 'when the pool is a NullPool' do
-      it 'returns unknown' do
+      it 'returns nil' do
         connection = double(:active_record_connection, pool: ActiveRecord::ConnectionAdapters::NullPool.new)
 
-        expect(described_class.db_config_name(connection)).to eq('unknown')
+        expect(described_class.db_config_for_connection(connection)).to be_nil
       end
     end
   end
 
+  describe '.db_config_name' do
+    it 'returns the db_config name for the connection' do
+      connection = ActiveRecord::Base.connection
+
+      expect(described_class.db_config_name(connection)).to be_a(String)
+      expect(described_class.db_config_name(connection)).to eq(connection.pool.db_config.name)
+    end
+  end
+
   describe '#true_value' do
     it 'returns correct value' do
       expect(described_class.true_value).to eq "'t'"