Merge pull request #295 from pope1ni/unix-socket-fixes

3 files changed, 88 insertions, 30 deletions

diff --git a/pymemcache/client/base.py b/pymemcache/client/base.py
index 1bf006f..7b3ed6d 100644
--- a/pymemcache/client/base.py
+++ b/pymemcache/client/base.py
@@ -37,7 +37,6 @@ VALID_STORE_RESULTS = {
     b'prepend': (b'STORED', b'NOT_STORED'),
     b'cas':     (b'STORED', b'EXISTS', b'NOT_FOUND'),
 }
-VALID_STRING_TYPES = (six.text_type, six.string_types)
 
 
 # Some of the values returned by the "stats" command
@@ -87,9 +86,9 @@ def check_key_helper(key, allow_unicode_keys, key_prefix=b''):
     if allow_unicode_keys:
         if isinstance(key, six.text_type):
             key = key.encode('utf8')
-    elif isinstance(key, VALID_STRING_TYPES):
+    elif isinstance(key, six.string_types):
         try:
-            if isinstance(key, bytes):
+            if isinstance(key, six.binary_type):
                 key = key.decode().encode('ascii')
             else:
                 key = key.encode('ascii')
@@ -264,7 +263,7 @@ class Client(object):
         self.sock = None
         if isinstance(key_prefix, six.text_type):
             key_prefix = key_prefix.encode('ascii')
-        if not isinstance(key_prefix, bytes):
+        if not isinstance(key_prefix, six.binary_type):
             raise TypeError("key_prefix should be bytes.")
         self.key_prefix = key_prefix
         self.default_noreply = default_noreply
@@ -794,7 +793,7 @@ class Client(object):
 
     def _check_integer(self, value, name):
         """Check that a value is an integer and encode it as a binary string"""
-        if not isinstance(value, six.integer_types):  # includes "long" on py2
+        if not isinstance(value, six.integer_types):
             raise MemcacheIllegalInputError(
                 '%s must be integer, got bad value: %r' % (name, value)
             )
@@ -808,7 +807,7 @@ class Client(object):
         The value will be (re)encoded so that we can accept strings or bytes.
         """
         # convert non-binary values to binary
-        if isinstance(cas, (six.integer_types, VALID_STRING_TYPES)):
+        if isinstance(cas, (six.integer_types, six.string_types)):
             try:
                 cas = six.text_type(cas).encode(self.encoding)
             except UnicodeEncodeError:
@@ -1042,7 +1041,7 @@ class PooledClient(object):
         self.allow_unicode_keys = allow_unicode_keys
         if isinstance(key_prefix, six.text_type):
             key_prefix = key_prefix.encode('ascii')
-        if not isinstance(key_prefix, bytes):
+        if not isinstance(key_prefix, six.binary_type):
             raise TypeError("key_prefix should be bytes.")
         self.key_prefix = key_prefix
         self.client_pool = pool.ObjectPool(
diff --git a/pymemcache/client/hash.py b/pymemcache/client/hash.py
index 536b915..9c61cb6 100644
--- a/pymemcache/client/hash.py
+++ b/pymemcache/client/hash.py
@@ -46,7 +46,8 @@ class HashClient(object):
         Constructor.
 
         Args:
-          servers: list(tuple(hostname, port))
+          servers: list() of tuple(hostname, port) or string containing a UNIX
+                   socket path.
           hasher: optional class three functions ``get_node``, ``add_node``,
                   and ``remove_node``
                   defaults to Rendezvous (HRW) hash.
@@ -101,30 +102,48 @@ class HashClient(object):
                 'lock_generator': lock_generator
             })
 
-        for server, port in servers:
-            self.add_server(server, port)
+        for server in servers:
+            self.add_server(server)
         self.encoding = encoding
         self.tls_context = tls_context
 
-    def add_server(self, server, port):
-        key = '%s:%s' % (server, port)
+    def add_server(self, server, port=None):
+        # To maintain backward compatibility, if a port is provided, assume
+        # that server wasn't provided as a (host, port) tuple.
+        if port is not None:
+            if not isinstance(server, six.string_types):
+                raise TypeError('Server must be a string when passing port.')
+            server = (server, port)
 
-        if self.use_pooling:
-            client = PooledClient(
-                (server, port),
-                **self.default_kwargs
-            )
+        if isinstance(server, six.string_types):
+            key = server
         else:
-            client = self.client_class((server, port), **self.default_kwargs)
+            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.client_class = self.client_class
 
         self.clients[key] = client
         self.hasher.add_node(key)
 
-    def remove_server(self, server, port):
+    def remove_server(self, server, port=None):
+        # To maintain backward compatibility, if a port is provided, assume
+        # that server wasn't provided as a (host, port) tuple.
+        if port is not None:
+            if not isinstance(server, six.string_types):
+                raise TypeError('Server must be a string when passing port.')
+            server = (server, port)
+
+        if isinstance(server, six.string_types):
+            key = server
+        else:
+            key = '%s:%s' % server
+
         dead_time = time.time()
-        self._failed_clients.pop((server, port))
-        self._dead_clients[(server, port)] = dead_time
-        key = '%s:%s' % (server, port)
+        self._failed_clients.pop(server)
+        self._dead_clients[server] = dead_time
         self.hasher.remove_node(key)
 
     def _retry_dead(self):
@@ -141,7 +160,7 @@ class HashClient(object):
                     'bringing server back into rotation %s',
                     server
                 )
-                self.add_server(*server)
+                self.add_server(server)
                 del self._dead_clients[server]
             self._last_dead_check_time = current_time
 
@@ -185,7 +204,7 @@ class HashClient(object):
                     # We've reached our max retry attempts, we need to mark
                     # the sever as dead
                     logger.debug('marking server as dead: %s', client.server)
-                    self.remove_server(*client.server)
+                    self.remove_server(client.server)
 
             result = func(*args, **kwargs)
             return result
@@ -239,7 +258,7 @@ class HashClient(object):
                     # We've reached our max retry attempts, we need to mark
                     # the sever as dead
                     logger.debug('marking server as dead: %s', client.server)
-                    self.remove_server(*client.server)
+                    self.remove_server(client.server)
 
             succeeded, failed, err = self._set_many(
                 client, values, *args, **kwargs
@@ -289,7 +308,7 @@ class HashClient(object):
                 'attempts': 0,
             }
             logger.debug("marking server as dead %s", server)
-            self.remove_server(*server)
+            self.remove_server(server)
         # This client has failed previously, we need to update the metadata
         # to reflect that we have attempted it again
         else:
diff --git a/pymemcache/test/test_client_hash.py b/pymemcache/test/test_client_hash.py
index 322ce3d..4ffad5c 100644
--- a/pymemcache/test/test_client_hash.py
+++ b/pymemcache/test/test_client_hash.py
@@ -5,6 +5,7 @@ from pymemcache import pool
 
 from .test_client import ClientTestMixin, MockSocket
 import unittest
+import os
 import pytest
 import mock
 import socket
@@ -42,7 +43,7 @@ class TestHashClient(ClientTestMixin, unittest.TestCase):
         client_class = 'pymemcache.client.hash.HashClient.client_class'
         with mock.patch(client_class) as internal_client:
             client = HashClient([], timeout=999, key_prefix='foo_bar_baz')
-            client.add_server('127.0.0.1', '11211')
+            client.add_server(('127.0.0.1', '11211'))
 
         assert internal_client.call_args[0][0] == ('127.0.0.1', '11211')
         kwargs = internal_client.call_args[1]
@@ -313,7 +314,7 @@ class TestHashClient(ClientTestMixin, unittest.TestCase):
     @mock.patch("pymemcache.client.hash.HashClient.client_class")
     def test_dead_server_comes_back(self, client_patch):
         client = HashClient([], dead_timeout=0, retry_attempts=0)
-        client.add_server("127.0.0.1", 11211)
+        client.add_server(("127.0.0.1", 11211))
 
         test_client = client_patch.return_value
         test_client.server = ("127.0.0.1", 11211)
@@ -332,7 +333,7 @@ class TestHashClient(ClientTestMixin, unittest.TestCase):
     @mock.patch("pymemcache.client.hash.HashClient.client_class")
     def test_failed_is_retried(self, client_patch):
         client = HashClient([], retry_attempts=1, retry_timeout=0)
-        client.add_server("127.0.0.1", 11211)
+        client.add_server(("127.0.0.1", 11211))
 
         assert client_patch.call_count == 1
 
@@ -354,7 +355,46 @@ class TestHashClient(ClientTestMixin, unittest.TestCase):
 
         client = HashClient([])
         client.client_class = MyClient
-        client.add_server('host', 11211)
+        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()),
+            ('127.0.0.1', 11211),
+        ]
+        client = HashClient(servers)
+        assert set(servers) == {c.server for c in client.clients.values()}
+
+    def test_legacy_add_remove_server_signature(self):
+        server = ('127.0.0.1', 11211)
+        client = HashClient([])
+        assert client.clients == {}
+        client.add_server(*server)  # Unpack (host, port) tuple.
+        assert ('%s:%s' % server) in client.clients
+        client._mark_failed_server(server)
+        assert server in client._failed_clients
+        client.remove_server(*server)  # Unpack (host, port) tuple.
+        assert server in client._dead_clients
+        assert server not in client._failed_clients
+
+        # Ensure that server is a string if passing port argument:
+        with pytest.raises(TypeError):
+            client.add_server(server, server[-1])
+        with pytest.raises(TypeError):
+            client.remove_server(server, server[-1])
+
     # TODO: Test failover logic