[py27] Fixing a bug with bad conditions (#337)

2 files changed, 42 insertions, 2 deletions

diff --git a/pymemcache/client/base.py b/pymemcache/client/base.py
index 773fead..dc5e16c 100644
--- a/pymemcache/client/base.py
+++ b/pymemcache/client/base.py
@@ -113,7 +113,7 @@ def check_key_helper(key, allow_unicode_keys, key_prefix=b''):
     if len(key) > 250:
         raise MemcacheIllegalInputError("Key is too long: %r" % key)
     # second statement catches leading or trailing whitespace
-    elif len(parts) > 1 or parts[0] != key:
+    elif len(parts) > 1 or (parts and parts[0] != key):
         raise MemcacheIllegalInputError("Key contains whitespace: %r" % key)
     elif b'\00' in key:
         raise MemcacheIllegalInputError("Key contains null: %r" % key)
diff --git a/pymemcache/test/test_client.py b/pymemcache/test/test_client.py
index 118a229..8b323de 100644
--- a/pymemcache/test/test_client.py
+++ b/pymemcache/test/test_client.py
@@ -24,6 +24,7 @@ import mock
 import platform
 import re
 import socket
+import sys
 import unittest
 
 import pytest
@@ -32,7 +33,8 @@ from pymemcache.client.base import (
     PooledClient,
     Client,
     normalize_server_spec,
-    KeepaliveOpts
+    KeepaliveOpts,
+    check_key_helper
 )
 from pymemcache.exceptions import (
     MemcacheClientError,
@@ -52,6 +54,44 @@ def is_ipv6(address):
     return re.match(r'^[0-9a-f:]+$', address)
 
 
+@pytest.mark.parametrize(
+    'key,allow_unicode_keys,key_prefix,ex_exception,ex_excinfo,ignore_py27',
+    [
+        (u'b'*251, True, b'',
+         MemcacheIllegalInputError, 'Key is too long', False),
+        (u'foo bar', True, b'',
+         MemcacheIllegalInputError, 'Key contains whitespace', False),
+        (u'\00', True, b'',
+         MemcacheIllegalInputError, 'Key contains null', False),
+        (None, True, b'', TypeError, None, False),
+        # The following test won't fail with a TypeError with python 2.7
+        (b"", False, '', TypeError, None, True),
+    ])
+@pytest.mark.unit()
+def test_check_key_helper_failing_conditions(key, allow_unicode_keys,
+                                             key_prefix, ex_exception,
+                                             ex_excinfo, ignore_py27):
+
+    if ignore_py27 and sys.version_info < (3, 0, 0):
+        pytest.skip("skipping for Python 2.7")
+    with pytest.raises(ex_exception) as excinfo:
+        check_key_helper(key, allow_unicode_keys, key_prefix)
+
+    # Allow to ignore the excinfo value. Different implementation (2.7,
+    # pypy, 3.x) comes with various error messages for the same thing.
+    # We just check the kind of the exception.
+    if ex_excinfo:
+        assert ex_excinfo in str(excinfo.value)
+
+
+@pytest.mark.unit()
+def test_check_key_helper():
+    assert check_key_helper(b"key", True, b'') == b"key"
+    assert check_key_helper("key", True) == b"key"
+    assert isinstance(check_key_helper(b"key", True), bytes)
+    assert check_key_helper("", True) == b""
+
+
 class MockSocket(object):
     def __init__(self, recv_bufs, connect_failure=None, close_failure=None):
         self.recv_bufs = collections.deque(recv_bufs)