1 files changed, 48 insertions, 0 deletions
diff --git a/spec/lib/gitlab/database_spec.rb b/spec/lib/gitlab/database_spec.rb
index b142b3a2781..c9be832b1a0 100644
--- a/spec/lib/gitlab/database_spec.rb
+++ b/spec/lib/gitlab/database_spec.rb
@@ -71,6 +71,54 @@ describe Gitlab::Database, lib: true do
     end
   end
 
+  describe '.with_connection_pool' do
+    it 'creates a new connection pool and disconnect it after used' do
+      closed_pool = nil
+
+      described_class.with_connection_pool(1) do |pool|
+        pool.with_connection do |connection|
+          connection.execute('SELECT 1 AS value')
+        end
+
+        expect(pool).to be_connected
+
+        closed_pool = pool
+      end
+
+      expect(closed_pool).not_to be_connected
+    end
+
+    it 'disconnects the pool even an exception was raised' do
+      error = Class.new(RuntimeError)
+      closed_pool = nil
+
+      begin
+        described_class.with_connection_pool(1) do |pool|
+          pool.with_connection do |connection|
+            connection.execute('SELECT 1 AS value')
+          end
+
+          closed_pool = pool
+
+          raise error.new('boom')
+        end
+      rescue error
+      end
+
+      expect(closed_pool).not_to be_connected
+    end
+  end
+
+  describe '.create_connection_pool' do
+    it 'creates a new connection pool with specific pool size' do
+      pool = described_class.create_connection_pool(5)
+
+      expect(pool)
+        .to be_kind_of(ActiveRecord::ConnectionAdapters::ConnectionPool)
+      expect(pool.spec.config[:pool]).to eq(5)
+    end
+  end
+
   describe '#true_value' do
     it 'returns correct value for PostgreSQL' do
       expect(described_class).to receive(:postgresql?).and_return(true)