refactoring classes and rewriting redis backend

Author: nitxiodev

pybloom/tests/__init__.py

@@ -1,21 +1,193 @@
 import unittest
 
 import redis
-from hamcrest import assert_that, equal_to, raises, is_, instance_of, empty, is_not, greater_than
+from fakeredis import FakeStrictRedis
+from hamcrest import assert_that, equal_to, raises, is_, instance_of, greater_than, empty, is_not
 from mock import mock
-from mockredis import mock_strict_redis_client
 from redis import StrictRedis
+from redis.exceptions import LockError
+from redis.lock import LuaLock
 
 from pybloom.src.backends.bitarraybackend import BitArrayBackend
 from pybloom.src.backends.numpybackend import NumpyBackend
 from pybloom.src.backends.redisbackend import RedisBackend, RedisProxy
 from pybloom.src.bloomfilter import BloomFilter, BloomFilterException, Options, Size, size_to_human_format
 
+LUA_ADD_SCRIPT = """
+
+    -- https://gist.github.com/tylerneylon/59f4bcf316be525b30ab
+    local json = {}
+
+    local function kind_of(obj)
+      if type(obj) ~= 'table' then return type(obj) end
+      local i = 1
+      for _ in pairs(obj) do
+        if obj[i] ~= nil then i = i + 1 else return 'table' end
+      end
+      if i == 1 then return 'table' else return 'array' end
+    end
+
+    local function escape_str(s)
+      local in_char  =  { '\\\\', '\\"', '/', '\b', '\f', '\\\n', '\\\r', '\t' }
+      local out_char = {'\\\\', '\\"', '/',  'b',  'f',  'n',  'r',  't'}
+
+      for i, c in ipairs(in_char) do
+        s = s:gsub(c, '\\\\' .. out_char[i])
+      end
+      return s
+    end
+
+
+    local function skip_delim(str, pos, delim, err_if_missing)
+      pos = pos + #str:match('^%s*', pos)
+      if str:sub(pos, pos) ~= delim then
+        if err_if_missing then
+          error('Expected ' .. delim .. ' near position ' .. pos)
+        end
+        return pos, false
+      end
+      return pos + 1, true
+    end
+
+    local function parse_str_val(str, pos, val)
+      val = val or ''
+      local early_end_error = 'End of input found while parsing string.'
+      if pos > #str then error(early_end_error) end
+      local c = str:sub(pos, pos)
+      if c == '\\"'  then return val, pos + 1 end
+      if c ~= '\\\\' then return parse_str_val(str, pos + 1, val .. c) end
+      -- We must have a \ character.
+      local esc_map = {b = '\b', f = '\f', n = '\\\n', r = '\\\r', t = '\t'}
+      local nextc = str:sub(pos + 1, pos + 1)
+      if not nextc then error(early_end_error) end
+      return parse_str_val(str, pos + 2, val .. (esc_map[nextc] or nextc))
+    end
+
+    local function parse_num_val(str, pos)
+      local num_str = str:match('^-?%d+%.?%d*[eE]?[+-]?%d*', pos)
+      local val = tonumber(num_str)
+      if not val then error('Error parsing number at position ' .. pos .. '.') end
+      return val, pos + #num_str
+    end
+
+    function json.stringify(obj, as_key)
+      local s = {}  -- We'll build the string as an array of strings to be concatenated.
+      local kind = kind_of(obj)  -- This is 'array' if it's an array or type(obj) otherwise.
+      if kind == 'array' then
+        if as_key then error('Can\\\'t encode array as key.') end
+        s[#s + 1] = '['
+        for i, val in ipairs(obj) do
+          if i > 1 then s[#s + 1] = ', ' end
+          s[#s + 1] = json.stringify(val)
+        end
+        s[#s + 1] = ']'
+      elseif kind == 'table' then
+        if as_key then error('Can\\\'t encode table as key.') end
+        s[#s + 1] = '{'
+        for k, v in pairs(obj) do
+          if #s > 1 then s[#s + 1] = ', ' end
+          s[#s + 1] = json.stringify(k, true)
+          s[#s + 1] = ':'
+          s[#s + 1] = json.stringify(v)
+        end
+        s[#s + 1] = '}'
+      elseif kind == 'string' then
+        return '"' .. escape_str(obj) .. '"'
+      elseif kind == 'number' then
+        if as_key then return '"' .. tostring(obj) .. '"' end
+        return tostring(obj)
+      elseif kind == 'boolean' then
+        return tostring(obj)
+      elseif kind == 'nil' then
+        return 'null'
+      else
+        error('Unjsonifiable type: ' .. kind .. '.')
+      end
+      return table.concat(s)
+    end
+
+    json.null = {}
+
+    function json.parse(str, pos, end_delim)
+      pos = pos or 1
+      if pos > #str then error('Reached unexpected end of input.') end
+      local pos = pos + #str:match('^%s*', pos)  -- Skip whitespace.
+      local first = str:sub(pos, pos)
+      if first == '{' then  -- Parse an object.
+        local obj, key, delim_found = {}, true, true
+        pos = pos + 1
+        while true do
+          key, pos = json.parse(str, pos, '}')
+          if key == nil then return obj, pos end
+          if not delim_found then error('Comma missing between object items.') end
+          pos = skip_delim(str, pos, ':', true)  -- true -> error if missing.
+          obj[key], pos = json.parse(str, pos)
+          pos, delim_found = skip_delim(str, pos, ',')
+        end
+      elseif first == '[' then  -- Parse an array.
+        local arr, val, delim_found = {}, true, true
+        pos = pos + 1
+        while true do
+          val, pos = json.parse(str, pos, ']')
+          if val == nil then return arr, pos end
+          if not delim_found then error('Comma missing between array items.') end
+          arr[#arr + 1] = val
+          pos, delim_found = skip_delim(str, pos, ',')
+        end
+      elseif first == '"' then  -- Parse a string.
+        return parse_str_val(str, pos + 1)
+      elseif first == '-' or first:match('%d') then  -- Parse a number.
+        return parse_num_val(str, pos)
+      elseif first == end_delim then  -- End of an object or array.
+        return nil, pos + 1
+      else  -- Parse true, false, or null.
+        local literals = {['true'] = true, ['false'] = false, ['null'] = json.null}
+        for lit_str, lit_val in pairs(literals) do
+          local lit_end = pos + #lit_str - 1
+          if str:sub(pos, lit_end) == lit_str then return lit_val, lit_end + 1 end
+        end
+        local pos_info_str = 'position ' .. pos .. ': ' .. str:sub(pos, pos + 10)
+        error('Invalid json syntax starting at ' .. pos_info_str)
+      end
+    end
+
+    local capacity = tonumber(redis.call('HGET', KEYS[1], 'capacity'))
+    local filter_size = tonumber(redis.call('HGET', KEYS[1], 'filter_size'))
+    local hash_size = tonumber(redis.call('HGET', KEYS[1], 'hash_size'))
+
+    -- This means that filter has been reset
+    if capacity == nil or filter_size == nil then
+        return false
+    end
+
+    if capacity >= filter_size then
+        return false
+    end
+
+    local sum = 0
+    for i=1, #ARGV do
+        local args = json.parse(ARGV[i])
+        if redis.call('GETBIT', args['key'], args['offset']) == 1 then
+            sum = sum + 1
+        end
+
+        redis.call('SETBIT', args['key'], args['offset'], 1)
+    end
+
+    -- Only if the element don't exists, increase the capacity (i.e. add it)
+    if sum ~= hash_size then
+        capacity = capacity + 1
+        redis.call('HSET', KEYS[1], 'capacity', capacity)
+    end
+
+    return capacity
+"""
+
 
 class MockRedisProxy(object):
     def __init__(self, *args, **kwargs):
         self._connection = mock.patch('pybloom.src.backends.redisbackend.redis.StrictRedis',
-                                      new_callable=mock_strict_redis_client).start()
+                                      new_callable=FakeStrictRedis).start()
 
     def as_pipeline(self):
         return self._connection.pipeline()
@@ -35,8 +207,8 @@ class testRedisProxy(unittest.TestCase):
     def setUp(self):
         self._proxy = RedisProxy('')
         self._proxy._connection = mock.patch('pybloom.src.backends.redisbackend.redis.StrictRedis',
-                                             new_callable=mock_strict_redis_client).start()
-        self._proxy._connection.reset = self._proxy._connection.pipeline()._reset  # bad signature in mockredis
+                                             new_callable=FakeStrictRedis).start()
+        self._proxy._connection.reset = self._proxy._connection.pipeline().reset  # bad signature in mockredis
 
     @mock.patch('pybloom.src.backends.redisbackend.redis.StrictRedis', spec=StrictRedis)
     def testRetryConnectionError(self, rediss):
@@ -60,7 +232,7 @@ def testRetryTimeoutError(self, rediss):
 
     def testNoRetry(self):
         self._proxy.set('key', 'hello')
-        assert_that(self._proxy.ping(), is_(b'PONG'))
+        assert_that(self._proxy.ping(), is_(True))
         assert_that(self._proxy.get('key'), equal_to(b'hello'))
 
     def testPipeline(self):
@@ -74,8 +246,9 @@ def testPipeline(self):
 class testRedisBackend(unittest.TestCase):
     def setUp(self):
         with mock.patch('pybloom.src.backends.redisbackend.RedisProxy', new=MockRedisProxy):
-            self._backend = RedisBackend(array_size=10, hash_size=3, redis_connection='')
-        #         # self._backend._redis = MockRedisProxy()
+            self._backend = RedisBackend(array_size=10, hash_size=3, redis_connection='', filter_size=5)
+            self._backend._lua_add = self._backend._redis.register_script(LUA_ADD_SCRIPT)
+            # self._backend._lua_add = self._backend._redis.register_script(LUA_ADD_KEY)
 
     def testRightOffset(self):
         # First, a simple offset (first 2^32)
@@ -88,16 +261,41 @@ def testRightOffset(self):
         assert_that(k, equal_to(2))  # 2 offsets in total (1: [0, 2^32-1], 2: [2^32, 2^33 - 1])
         assert_that(offset, equal_to(2 ** 33 - 1))
 
-    def testResetWithNoData(self):
-        self._backend.reset()
-        assert_that(list(self._backend._redis.scan_iter('bloomfilter:*')), is_(empty()))
+    @mock.patch('pybloom.src.backends.redisbackend.lock', spec=LuaLock)
+    def testMetadataError(self, mock_lock):
+        mock_lock.side_effect = LockError
+        with mock.patch('pybloom.src.backends.redisbackend.RedisProxy', new=MockRedisProxy):
+            with self.assertRaises(BloomFilterException) as cm:
+                self._backend = RedisBackend(array_size=10, hash_size=3, redis_connection='', filter_size=5)
 
-    def testResetWithData(self):
+            assert_that(str(cm.exception), equal_to("Cannot retrieve metadata from redis. Seems another process has "
+                                                    "acquired the lock and did not released. Check if "
+                                                    "'bloom_filter_lock' key is in your redis server."))
+
+    # @mock.patch('pybloom.src.backends.redisbackend.lock', spec=LuaLock)
+    def testMetadataOk(self):
+        # Check we dont have any metadata yet
+        response = {key.decode(): val.decode() for key, val in
+                    self._backend._redis.hgetall(self._backend._metadata_key).items()}
+        assert_that(response, equal_to(dict(array_size='10', hash_size='3', filter_size='5', capacity='0')))
+
+        # Add data
         self._backend.add(4)
-        assert_that(list(self._backend._redis.scan_iter('bloomfilter:*')), is_not(empty()))
+
+        # Check metadata again
+        response = {key.decode(): val.decode() for key, val in
+                    self._backend._redis.hgetall(self._backend._metadata_key).items()}
+        assert_that(response, equal_to(dict(array_size='10', hash_size='3', filter_size='5', capacity='1')))
+
+    def testResetWithData(self):
+        self._backend.add(45)
+        assert_that(list(self._backend._redis.scan_iter('bloom_filter:*')), is_not(empty()))
 
         self._backend.reset()
-        assert_that(list(self._backend._redis.scan_iter('bloomfilter:*')), is_(empty()))
+        assert_that(list(self._backend._redis.scan_iter('bloom_filter:*')), is_(empty()))
+        assert_that(len(self._backend), is_(0))
+        response = self._backend._redis.hgetall(self._backend._metadata_key).items()
+        assert_that(response, is_(empty()))
 
     def testAddandCheck(self):
         self._backend.add('house')
@@ -106,9 +304,10 @@ def testAddandCheck(self):
         self._backend += 'horse'
         assert_that('horse' in self._backend, is_(True))
 
+
 class testNumpyBackend(unittest.TestCase):
     def setUp(self):
-        self._backend = NumpyBackend(array_size=10, hash_size=3)
+        self._backend = NumpyBackend(array_size=10, hash_size=3, filter_size=5)
 
     def testResetWithNoData(self):
         self._backend.reset()
@@ -128,12 +327,12 @@ def testAddandCheck(self):
         self._backend += 'horse'
         assert_that('horse' in self._backend, is_(True))
 
-        assert_that(self._backend._capacity, equal_to(2))
+        assert_that(len(self._backend), equal_to(2))
 
 
 class testBitArrayBackend(unittest.TestCase):
     def setUp(self):
-        self._backend = BitArrayBackend(array_size=10, hash_size=3)
+        self._backend = BitArrayBackend(array_size=10, hash_size=3, filter_size=5)
 
     def testResetWithNoData(self):
         self._backend.reset()
@@ -153,7 +352,7 @@ def testAddandCheck(self):
         self._backend += 'horse'
         assert_that('horse' in self._backend, is_(True))
 
-        assert_that(self._backend._capacity, equal_to(2))
+        assert_that(len(self._backend), equal_to(2))
 
 
 class testBloomFilter(unittest.TestCase):