From 1e0d14ecd8d3639ee9b53cc0baba7d859e612b1e Mon Sep 17 00:00:00 2001
From: Kendall Gifford <zettabyte@gmail.com>
Date: Sat, 15 Jun 2013 11:39:32 -0600
Subject: [PATCH] Allow us to shutdown pool within a wrapper

---
 lib/connection_pool.rb       |  8 +++++---
 test/test_connection_pool.rb | 15 ++++++++++++---
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/lib/connection_pool.rb b/lib/connection_pool.rb
index 4740973..c23dbee 100644
--- a/lib/connection_pool.rb
+++ b/lib/connection_pool.rb
@@ -79,9 +79,7 @@ class ConnectionPool
   end
 
   class Wrapper < ::BasicObject
-    METHODS = [:with]
-
-    attr_reader :pool
+    METHODS = [:with, :pool_shutdown]
 
     def initialize(options = {}, &block)
       @pool = ::ConnectionPool.new(options, &block)
@@ -93,6 +91,10 @@ class ConnectionPool
       @pool.checkin
     end
 
+    def pool_shutdown(&block)
+      @pool.shutdown(&block)
+    end
+
     def respond_to?(id, *args)
       METHODS.include?(id) || @pool.with { |c| c.respond_to?(id, *args) }
     end
diff --git a/test/test_connection_pool.rb b/test/test_connection_pool.rb
index eef432c..a89c252 100644
--- a/test/test_connection_pool.rb
+++ b/test/test_connection_pool.rb
@@ -200,8 +200,17 @@ class TestConnectionPool < Minitest::Test
     end
   end
 
-  def test_wrapper_provides_access_to_pool
-    wrapper = ConnectionPool::Wrapper.new(:size => 1) { true }
-    assert_equal ConnectionPool, wrapper.pool.class
+  def test_shutdown_is_executed_for_all_connections_in_wrapped_pool
+    recorders = []
+
+    wrapper = ConnectionPool::Wrapper.new(:size => 3) do
+      Recorder.new.tap { |r| recorders << r }
+    end
+
+    wrapper.pool_shutdown do |recorder|
+      recorder.do_work("shutdown")
+    end
+
+    assert_equal [["shutdown"]] * 3, recorders.map { |r| r.calls }
   end
 end