Fixed custom client for HashClient using pooling.

2 files changed, 16 insertions, 6 deletions

diff --git a/pymemcache/client/hash.py b/pymemcache/client/hash.py
index 4716549..9c61cb6 100644
--- a/pymemcache/client/hash.py
+++ b/pymemcache/client/hash.py
@@ -120,13 +120,10 @@ class HashClient(object):
         else:
             key = '%s:%s' % server
 
+        _class = PooledClient if self.use_pooling else self.client_class
+        client = _class(server, **self.default_kwargs)
         if self.use_pooling:
-            client = PooledClient(
-                server,
-                **self.default_kwargs
-            )
-        else:
-            client = self.client_class(server, **self.default_kwargs)
+            client.client_class = self.client_class
 
         self.clients[key] = client
         self.hasher.add_node(key)
diff --git a/pymemcache/test/test_client_hash.py b/pymemcache/test/test_client_hash.py
index f972f34..4ffad5c 100644
--- a/pymemcache/test/test_client_hash.py
+++ b/pymemcache/test/test_client_hash.py
@@ -358,6 +358,19 @@ class TestHashClient(ClientTestMixin, unittest.TestCase):
         client.add_server(('host', 11211))
         assert isinstance(client.clients['host:11211'], MyClient)
 
+    def test_custom_client_with_pooling(self):
+        class MyClient(Client):
+            pass
+
+        client = HashClient([], use_pooling=True)
+        client.client_class = MyClient
+        client.add_server(('host', 11211))
+        assert isinstance(client.clients['host:11211'], PooledClient)
+
+        pool = client.clients['host:11211'].client_pool
+        with pool.get_and_release(destroy_on_fail=True) as c:
+            assert isinstance(c, MyClient)
+
     def test_mixed_inet_and_unix_sockets(self):
         servers = [
             '/tmp/pymemcache.{pid}'.format(pid=os.getpid()),