Add client kill with filter

Signed-off-by: Theo Despoudis <thdespou@hotmail.com>

3 files changed, 81 insertions, 1 deletions

diff --git a/redis/client.py b/redis/client.py
index 957081c..6791d67 100755
--- a/redis/client.py
+++ b/redis/client.py
@@ -471,7 +471,7 @@ class Redis(object):
         {
             'CLIENT GETNAME': lambda r: r and nativestr(r),
             'CLIENT ID': int,
-            'CLIENT KILL': bool_ok,
+            'CLIENT KILL':  lambda r: int or nativestr(r) == 'OK',
             'CLIENT LIST': parse_client_list,
             'CLIENT SETNAME': bool_ok,
             'CLIENT UNBLOCK': lambda r: r and int(r) == 1 or False,
@@ -790,6 +790,40 @@ class Redis(object):
         "Disconnects the client at ``address`` (ip:port)"
         return self.execute_command('CLIENT KILL', address)
 
+    def client_kill_filter(self, *filter_options):
+        """
+        Disconnects the client using a variety of filter options"
+        :param filter_options: a tuple or list of filter options with the following format:
+        (filter, value, filter, value,...) or
+        [filter, value, filter, value,...]
+        """
+        if len(filter_options) == 1 and isinstance(filter_options[0], basestring):
+            return self.client_kill(filter_options[0])
+        if not isinstance(filter_options, (list, tuple)) or not filter_options:
+            raise DataError("CLIENT KILL <filter> <value> ... ... <filter> <value>" +
+                            "must be a non empty list or "
+                            "tuple to execute")
+        if len(filter_options) % 2 != 0:
+            raise DataError("CLIENT KILL <filter> <value> requires a filter and a value pair. Got %r" % (
+                filter_options,))
+        filter_types = ('addr', 'id', 'type', 'skipme')
+        client_types = ('normal', 'master', 'slave', 'pubsub')
+        yes_no = ('yes', 'no')
+        for index in range(0, len(filter_options), 2):
+            option = filter_options[index]
+            value = filter_options[index + 1]
+            key = str(option).lower()
+            if key not in filter_types:
+                raise DataError("CLIENT KILL <filter> must be one of %r" % (
+                    filter_types,))
+            if key == 'type' and option not in client_types:
+                raise DataError("CLIENT KILL TYPE <value> must be one of %r" % (
+                    client_types,))
+            if key == 'skipme' and value not in yes_no:
+                raise DataError("CLIENT KILL SKIPME <value> must be one of %r" % (
+                    yes_no,))
+        return self.execute_command('CLIENT KILL', *filter_options)
+
     def client_list(self, _type=None):
         """
         Returns a list of currently connected clients.
diff --git a/tests/conftest.py b/tests/conftest.py
index 5a43968..7a36954 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -47,6 +47,15 @@ def r(request, **kwargs):
     return _get_client(redis.Redis, request, **kwargs)
 
 
+@pytest.fixture()
+def r3(request, **kwargs):
+    return [
+        _get_client(redis.Redis, request, **kwargs),
+        _get_client(redis.Redis, request, **kwargs),
+        _get_client(redis.Redis, request, **kwargs)
+    ]
+
+
 def _gen_cluster_mock_resp(r, response):
     mock_connection_pool = Mock()
     connection = Mock()
diff --git a/tests/test_commands.py b/tests/test_commands.py
index 733ec4f..b6177f7 100644
--- a/tests/test_commands.py
+++ b/tests/test_commands.py
@@ -95,6 +95,43 @@ class TestRedisCommands(object):
         assert r.client_setname('redis_py_test')
         assert r.client_getname() == 'redis_py_test'
 
+    @skip_if_server_version_lt('2.4.0')
+    def test_client_kill(self, r, r3):
+        clients = r.client_list()
+        r.client_kill(clients[1].get('addr'))
+        clients = r.client_list()
+        assert len(clients) == 3
+
+    @skip_if_server_version_lt('2.8.12')
+    def test_client_kill_filter_invalid_params(self, r):
+        # invalid type
+        with pytest.raises(exceptions.DataError):
+            r.client_kill_filter(123)
+
+        # empty list
+        with pytest.raises(exceptions.DataError):
+            r.client_kill_filter([])
+
+        # empty tuple
+        with pytest.raises(exceptions.DataError):
+            r.client_kill_filter(())
+
+        # missing values
+        with pytest.raises(exceptions.DataError):
+            r.client_kill_filter(*["type", "master", "skipme"])
+
+        # invalid filter
+        with pytest.raises(exceptions.DataError):
+            r.client_kill_filter(*["type", "master", "allow", "yes"])
+
+        # invalid type
+        with pytest.raises(exceptions.DataError):
+            r.client_kill_filter(*["type", "caster", "skipme", "yes"])
+
+        # invalid skipme
+        with pytest.raises(exceptions.DataError):
+            r.client_kill_filter(*["type", "master", "skipme", "yeah"])
+
     @skip_if_server_version_lt('2.6.9')
     def test_client_list_after_client_setname(self, r):
         r.client_setname('redis_py_test')