def test_redis_communication_reader_set_batch( mock_read_communication_from_buffer, right_to_left, inferred): redis = Mock( srandmember=Mock(side_effect=[ [sentinel.buf0, sentinel.buf1], ]), type=Mock(return_value=b'set'), ) mock_read_communication_from_buffer.side_effect = [ sentinel.comm0, sentinel.comm1 ] reader = RedisCommunicationReader(redis, sentinel.key, key_type=None if inferred else 'set', add_references=sentinel.add_references, right_to_left=right_to_left) assert reader.batch( sentinel.batch_size) == [sentinel.comm0, sentinel.comm1] if inferred: redis.type.assert_called_once_with(sentinel.key) else: assert not redis.type.called redis.srandmember.assert_called_once_with(sentinel.key, sentinel.batch_size) mock_read_communication_from_buffer.assert_has_calls([ call(sentinel.buf0, add_references=sentinel.add_references), call(sentinel.buf1, add_references=sentinel.add_references), ])
def test_set(self): key = 'dataset' with RedisServer(loglevel='warning') as server: redis_db = Redis(port=server.port) redis_db.sadd(key, write_communication_to_buffer(self.comm1)) redis_db.sadd(key, write_communication_to_buffer(self.comm2)) redis_db.sadd(key, write_communication_to_buffer(self.comm3)) reader = RedisCommunicationReader(redis_db, key, key_type='set') comms = [c for c in reader] ids = [c.id for c in comms] # assert no duplicates self.assertEquals(3, len(ids)) self.assertEquals(set(['comm-1', 'comm-2', 'comm-3']), set(ids)) self.assertEquals(3, len(reader)) self.assertTrue(hasattr(comms[0], 'sentenceForUUID')) self.assertTrue(hasattr(comms[1], 'sentenceForUUID')) self.assertTrue(hasattr(comms[2], 'sentenceForUUID')) batch_ids = [c.id for c in reader.batch(2)] # do this weird thing because set(['foo']) != set([u'foo']) self.assertTrue( ('comm-1' in batch_ids and 'comm-2' in batch_ids) or ('comm-1' in batch_ids and 'comm-3' in batch_ids) or ('comm-2' in batch_ids and 'comm-3' in batch_ids) ) # assert data still there ids = [c.id for c in reader] self.assertEquals(set(['comm-1', 'comm-2', 'comm-3']), set(ids)) self.assertEquals(3, redis_db.scard(key))
def test_redis_communication_reader_set_batch( mock_read_communication_from_buffer, right_to_left, inferred): redis = Mock( srandmember=Mock(side_effect=[ [sentinel.buf0, sentinel.buf1], ]), type=Mock(return_value=b'set'), ) mock_read_communication_from_buffer.side_effect = [ sentinel.comm0, sentinel.comm1 ] reader = RedisCommunicationReader( redis, sentinel.key, key_type=None if inferred else 'set', add_references=sentinel.add_references, right_to_left=right_to_left) assert reader.batch(sentinel.batch_size) == [sentinel.comm0, sentinel.comm1] if inferred: redis.type.assert_called_once_with(sentinel.key) else: assert not redis.type.called redis.srandmember.assert_called_once_with(sentinel.key, sentinel.batch_size) mock_read_communication_from_buffer.assert_has_calls([ call(sentinel.buf0, add_references=sentinel.add_references), call(sentinel.buf1, add_references=sentinel.add_references), ])
def test_redis_communication_reader_failed_infer(): redis = Mock(type=Mock(return_value=b'none'), ) with raises(Exception): RedisCommunicationReader(redis, sentinel.key, add_references=sentinel.add_references) redis.type.assert_called_with(sentinel.key)
def test_hash(self): key = 'dataset' with RedisServer(loglevel='warning') as server: redis_db = Redis(port=server.port) redis_db.hset(key, self.comm1.uuid.uuidString, write_communication_to_buffer(self.comm1)) redis_db.hset(key, self.comm2.uuid.uuidString, write_communication_to_buffer(self.comm2)) redis_db.hset(key, self.comm3.uuid.uuidString, write_communication_to_buffer(self.comm3)) reader = RedisCommunicationReader(redis_db, key, key_type='hash') comms = [c for c in reader] ids = [c.id for c in comms] # assert no duplicates self.assertEquals(3, len(ids)) self.assertEquals(set(['comm-1', 'comm-2', 'comm-3']), set(ids)) self.assertEquals(3, len(reader)) self.assertEquals('comm-2', reader[self.comm2.uuid.uuidString].id) self.assertTrue(hasattr(comms[0], 'sentenceForUUID')) self.assertTrue(hasattr(comms[1], 'sentenceForUUID')) self.assertTrue(hasattr(comms[2], 'sentenceForUUID')) # assert data still there ids = [c.id for c in reader] self.assertEquals(set(['comm-1', 'comm-2', 'comm-3']), set(ids)) self.assertEquals(3, redis_db.hlen(key))
def test_redis_communication_reader_failed_batch(right_to_left, inferred, key_type, pop, block): redis = Mock( type=Mock(return_value=key_type.encode('utf-8')), ) reader = RedisCommunicationReader( redis, sentinel.key, key_type=None if inferred else key_type, add_references=sentinel.add_references, pop=pop, block=block) if inferred: redis.type.assert_called_once_with(sentinel.key) else: assert not redis.type.called with raises(Exception): reader.batch(sentinel.batch_size)
def test_redis_communication_reader_failed_batch(right_to_left, inferred, key_type, pop, block): redis = Mock(type=Mock(return_value=key_type.encode('utf-8')), ) reader = RedisCommunicationReader(redis, sentinel.key, key_type=None if inferred else key_type, add_references=sentinel.add_references, pop=pop, block=block) if inferred: redis.type.assert_called_once_with(sentinel.key) else: assert not redis.type.called with raises(Exception): reader.batch(sentinel.batch_size)
def test_redis_communication_reader_list_empty( mock_read_communication_from_buffer, right_to_left, pop, block): redis = Mock( llen=Mock(return_value=0), lindex=Mock(), lpop=Mock(side_effect=[None]), rpop=Mock(side_effect=[None]), blpop=Mock(side_effect=[None]), brpop=Mock(side_effect=[None]), ) reader = RedisCommunicationReader(redis, sentinel.key, key_type='list', add_references=sentinel.add_references, right_to_left=right_to_left, pop=pop, block=block, block_timeout=sentinel.block_timeout) it = iter(reader) with raises(StopIteration): next(it) if pop: if block: assert not redis.llen.called assert not redis.lindex.called assert not redis.lpop.called assert not redis.rpop.called getattr( redis, 'brpop' if right_to_left else 'blpop').assert_called_once_with( sentinel.key, timeout=sentinel.block_timeout) assert not getattr(redis, 'blpop' if right_to_left else 'brpop').called else: assert not redis.llen.called assert not redis.lindex.called getattr( redis, 'rpop' if right_to_left else 'lpop').assert_called_once_with( sentinel.key) assert not getattr(redis, 'lpop' if right_to_left else 'rpop').called assert not redis.blpop.called assert not redis.brpop.called else: redis.llen.assert_called_with(sentinel.key) assert not redis.lindex.called assert not redis.lpop.called assert not redis.rpop.called assert not redis.blpop.called assert not redis.brpop.called assert not mock_read_communication_from_buffer.called
def test_list_implicit(self): key = 'dataset' with RedisServer(loglevel='warning') as server: redis_db = Redis(port=server.port) redis_db.lpush(key, write_communication_to_buffer(self.comm1)) redis_db.lpush(key, write_communication_to_buffer(self.comm2)) redis_db.lpush(key, write_communication_to_buffer(self.comm3)) reader = RedisCommunicationReader(redis_db, key) ids = [c.id for c in reader] self.assertEquals(['comm-1', 'comm-2', 'comm-3'], ids)
def test_redis_communication_reader_set(mock_read_communication_from_buffer, right_to_left, inferred): def _sadd(_, buf): return { sentinel.buf0: 1, sentinel.buf1: 0, sentinel.buf2: 0, sentinel.buf3: 1, sentinel.buf4: 0, }[buf] redis = Mock( sscan=Mock(side_effect=[ (7, [sentinel.buf0, sentinel.buf1]), (3, [sentinel.buf2]), (0, [sentinel.buf3, sentinel.buf4]), ]), sadd=Mock(side_effect=_sadd), scard=Mock(return_value=2), expire=Mock(), exists=Mock(side_effect=[True, False]), type=Mock(return_value=b'set'), ) mock_read_communication_from_buffer.side_effect = [ sentinel.comm0, sentinel.comm3 ] reader = RedisCommunicationReader(redis, 'my-key', key_type=None if inferred else 'set', add_references=sentinel.add_references, right_to_left=right_to_left) it = iter(reader) assert next(it) == sentinel.comm0 assert next(it) == sentinel.comm3 with raises(StopIteration): next(it) if inferred: redis.type.assert_called_once_with('my-key') else: assert not redis.type.called redis.sscan.assert_has_calls([ call('my-key', 0), call('my-key', 7), call('my-key', 3), ]) mock_read_communication_from_buffer.assert_has_calls([ call(sentinel.buf0, add_references=sentinel.add_references), call(sentinel.buf3, add_references=sentinel.add_references), ])
def test_set_implicit(self): key = 'dataset' with RedisServer(loglevel='warning') as server: redis_db = Redis(port=server.port) redis_db.sadd(key, write_communication_to_buffer(self.comm1)) redis_db.sadd(key, write_communication_to_buffer(self.comm2)) redis_db.sadd(key, write_communication_to_buffer(self.comm3)) reader = RedisCommunicationReader(redis_db, key) ids = [c.id for c in reader] # assert no duplicates self.assertEquals(3, len(ids)) self.assertEquals(set(['comm-1', 'comm-2', 'comm-3']), set(ids))
def test_redis_communication_reader_failed_pop(right_to_left, inferred, block): redis = Mock(type=Mock(return_value=b'hash'), ) with raises(ValueError): RedisCommunicationReader(redis, sentinel.key, key_type=None if inferred else 'hash', add_references=sentinel.add_references, block=block, pop=True) if not inferred: assert not redis.type.called
def test_redis_communication_reader_failed_block(right_to_left, inferred, key_type, pop): redis = Mock(type=Mock(return_value=key_type.encode('utf-8')), ) with raises(ValueError): RedisCommunicationReader(redis, sentinel.key, key_type=None if inferred else key_type, add_references=sentinel.add_references, pop=pop, block=True) if not inferred: assert not redis.type.called
def test_list_no_add_references(self): key = 'dataset' with RedisServer(loglevel='warning') as server: redis_db = Redis(port=server.port) redis_db.lpush(key, write_communication_to_buffer(self.comm1)) redis_db.lpush(key, write_communication_to_buffer(self.comm2)) redis_db.lpush(key, write_communication_to_buffer(self.comm3)) reader = RedisCommunicationReader(redis_db, key, key_type='list', add_references=False) comms = [c for c in reader] ids = [c.id for c in comms] self.assertEquals(['comm-1', 'comm-2', 'comm-3'], ids) self.assertFalse(hasattr(comms[0], 'sentenceForUUID')) self.assertFalse(hasattr(comms[1], 'sentenceForUUID')) self.assertFalse(hasattr(comms[2], 'sentenceForUUID'))
def test_hash_empty(self): key = 'dataset' with RedisServer(loglevel='warning') as server: redis_db = Redis(port=server.port) reader = RedisCommunicationReader(redis_db, key, key_type='hash') self.assertEquals(0, len(reader)) redis_db.hset(key, self.comm1.uuid.uuidString, write_communication_to_buffer(self.comm1)) redis_db.hset(key, self.comm2.uuid.uuidString, write_communication_to_buffer(self.comm2)) redis_db.hset(key, self.comm3.uuid.uuidString, write_communication_to_buffer(self.comm3)) ids = [c.id for c in reader] # assert no duplicates self.assertEquals(3, len(ids)) self.assertEquals(set(['comm-1', 'comm-2', 'comm-3']), set(ids)) self.assertEquals(3, len(reader))
def test_list(self): key = 'dataset' with RedisServer(loglevel='warning') as server: redis_db = Redis(port=server.port) redis_db.lpush(key, write_communication_to_buffer(self.comm1)) redis_db.lpush(key, write_communication_to_buffer(self.comm2)) redis_db.lpush(key, write_communication_to_buffer(self.comm3)) reader = RedisCommunicationReader(redis_db, key, key_type='list') comms = [c for c in reader] ids = [c.id for c in comms] self.assertEquals(['comm-1', 'comm-2', 'comm-3'], ids) self.assertEquals(3, len(reader)) self.assertEquals('comm-2', reader[1].id) self.assertTrue(hasattr(comms[0], 'sentenceForUUID')) self.assertTrue(hasattr(comms[1], 'sentenceForUUID')) self.assertTrue(hasattr(comms[2], 'sentenceForUUID')) # assert data still there ids = [c.id for c in reader] self.assertEquals(['comm-1', 'comm-2', 'comm-3'], ids) self.assertEquals(3, redis_db.llen(key))
def test_list_block_pop_timeout(self): key = 'dataset' with RedisServer(loglevel='warning') as server: redis_db = Redis(port=server.port) redis_db.lpush(key, write_communication_to_buffer(self.comm1)) redis_db.lpush(key, write_communication_to_buffer(self.comm2)) redis_db.lpush(key, write_communication_to_buffer(self.comm3)) reader = RedisCommunicationReader(redis_db, key, key_type='list', pop=True, block=True, block_timeout=1) it = iter(reader) ids = [] ids.append(it.next().id) ids.append(it.next().id) self.assertEquals(1, redis_db.llen(key)) ids.append(it.next().id) self.assertEquals(['comm-1', 'comm-2', 'comm-3'], ids) with self.assertRaises(StopIteration): print 'Waiting for timeout (1 sec)...' it.next()
def test_list_pop_left_to_right(self): key = 'dataset' with RedisServer(loglevel='warning') as server: redis_db = Redis(port=server.port) redis_db.lpush(key, write_communication_to_buffer(self.comm1)) redis_db.lpush(key, write_communication_to_buffer(self.comm2)) redis_db.lpush(key, write_communication_to_buffer(self.comm3)) reader = RedisCommunicationReader(redis_db, key, key_type='list', pop=True, right_to_left=False) it = iter(reader) ids = [] ids.append(it.next().id) ids.append(it.next().id) self.assertEquals(1, redis_db.llen(key)) ids.append(it.next().id) self.assertEquals(['comm-3', 'comm-2', 'comm-1'], ids) # assert data is gone self.assertEquals([], [c.id for c in reader]) self.assertFalse(redis_db.exists(key)) with self.assertRaises(StopIteration): it.next()
def test_list_block_pop(self): key = 'dataset' with RedisServer(loglevel='warning') as server: redis_db = Redis(port=server.port) redis_db.lpush(key, write_communication_to_buffer(self.comm1)) redis_db.lpush(key, write_communication_to_buffer(self.comm2)) redis_db.lpush(key, write_communication_to_buffer(self.comm3)) reader = RedisCommunicationReader(redis_db, key, key_type='list', pop=True, block=True) it = iter(reader) ids = [] ids.append(it.next().id) ids.append(it.next().id) self.assertEquals(1, redis_db.llen(key)) ids.append(it.next().id) self.assertEquals(['comm-1', 'comm-2', 'comm-3'], ids) proc = Process(target=_add_comm_to_list, args=(3, server.port, 'comm-4', key)) proc.start() print 'Waiting for new comm to be added (3 sec)...' self.assertEquals('comm-4', iter(reader).next().id) proc.join()
def test_redis_communication_reader_set_pop( mock_read_communication_from_buffer, right_to_left, inferred): redis = Mock( spop=Mock(side_effect=[sentinel.buf0, sentinel.buf1, None]), type=Mock(return_value=b'set'), ) mock_read_communication_from_buffer.side_effect = [ sentinel.comm0, sentinel.comm1 ] reader = RedisCommunicationReader(redis, sentinel.key, key_type=None if inferred else 'set', add_references=sentinel.add_references, right_to_left=right_to_left, pop=True) it = iter(reader) assert next(it) == sentinel.comm0 assert next(it) == sentinel.comm1 with raises(StopIteration): next(it) if inferred: redis.type.assert_called_once_with(sentinel.key) else: assert not redis.type.called redis.spop.assert_has_calls([ call(sentinel.key), call(sentinel.key), call(sentinel.key), ]) mock_read_communication_from_buffer.assert_has_calls([ call(sentinel.buf0, add_references=sentinel.add_references), call(sentinel.buf1, add_references=sentinel.add_references), ])
def test_redis_communication_reader_set_pop_empty( mock_read_communication_from_buffer, right_to_left): redis = Mock( spop=Mock(side_effect=[None]), type=Mock(return_value=b'set'), ) reader = RedisCommunicationReader(redis, sentinel.key, key_type='set', add_references=sentinel.add_references, right_to_left=right_to_left, pop=True) it = iter(reader) with raises(StopIteration): next(it) assert not redis.type.called redis.spop.assert_called_once_with(sentinel.key) assert not mock_read_communication_from_buffer.called
def test_set_pop(self): key = 'dataset' with RedisServer(loglevel='warning') as server: redis_db = Redis(port=server.port) redis_db.sadd(key, write_communication_to_buffer(self.comm1)) redis_db.sadd(key, write_communication_to_buffer(self.comm2)) redis_db.sadd(key, write_communication_to_buffer(self.comm3)) reader = RedisCommunicationReader(redis_db, key, key_type='set', pop=True) it = iter(reader) ids = [] ids.append(it.next().id) ids.append(it.next().id) self.assertEquals(1, redis_db.scard(key)) ids.append(it.next().id) # assert no duplicates self.assertEquals(3, len(ids)) self.assertEquals(set(['comm-1', 'comm-2', 'comm-3']), set(ids)) # assert data is gone self.assertEquals([], [c.id for c in reader]) self.assertFalse(redis_db.exists(key)) with self.assertRaises(StopIteration): it.next()
def test_redis_communication_reader_len(right_to_left, inferred, key_type, pop, block): redis = Mock( llen=Mock(side_effect=[3]), hlen=Mock(side_effect=[3]), scard=Mock(side_effect=[3]), type=Mock(return_value=key_type.encode('utf-8')), ) reader = RedisCommunicationReader(redis, sentinel.key, key_type=None if inferred else key_type, right_to_left=right_to_left, pop=pop, block=block) if inferred: redis.type.assert_called_once_with(sentinel.key) else: assert not redis.type.called len(reader) if key_type == 'set': assert redis.scard.called_once_with(sentinel.key) assert not redis.llen.called assert not redis.hlen.called elif key_type == 'list': assert not redis.scard.called assert redis.llen.called_once_with(sentinel.key) assert not redis.hlen.called elif key_type == 'hash': assert not redis.scard.called assert not redis.llen.called assert redis.hlen.called_once_with(sentinel.key) else: raise ValueError('unexpected key type {}'.format(key_type))
def get_redis_comm(redis_host, redis_port, redis_comm, comm_lookup_by, redis_comm_map, redis_comm_index, redis_direction, communication_loc): def comm_lookup(comm): if comm_lookup_by == 'id': return comm.id elif comm_lookup_by == 'uuid': return comm.uuid.uuidString else: error('unsupported comm_lookup_by %s' % comm_lookup_by) input_db = None if redis_comm_index is not None: if redis_direction == RIGHT_TO_LEFT: redis_comm_index = -(redis_comm_index + 1) else: redis_comm_index = redis_comm_index else: redis_comm_index = None if redis_comm is not None and redis_comm_index is not None: error("Cannot include both --redis-comm and --redis-comm-index") input_db = Redis(redis_host, redis_port) reader = RedisCommunicationReader( input_db, communication_loc, add_references=True, right_to_left=(redis_direction == RIGHT_TO_LEFT)) if redis_comm: # look up comm in collection by comm field (comm_lookup_by) key_type = input_db.type(communication_loc) if ((key_type == 'list' and redis_comm_map is None) or (key_type == 'set')): # do linear scan for co in reader: if comm_lookup(co) == redis_comm: comm = co break elif key_type == 'list' and redis_comm_map is not None: # look up list index using field->index map def no_comm(comm_idx): error(('Unable to find a communication with identifier "%s"' ' with value "%s" using pivoting map "%s", which' ' returned (list) index %s, under the %s %s') % (comm_lookup_by, redis_comm, redis_comm_map, str(comm_idx), key_type, communication_loc)) comm_idx = input_db.hget(redis_comm_map, redis_comm) if comm_idx is None: no_comm(comm_idx) comm_idx = int(comm_idx) if redis_direction == RIGHT_TO_LEFT: comm_idx = -(comm_idx + 1) comm_buf = input_db.lindex(communication_loc, comm_idx) comm = read_communication_from_buffer(comm_buf) if comm_lookup(comm) != redis_comm: error(('Cannot find the appropriate document with %s' ' indexing') % redis_direction) elif key_type == 'hash': # do O(1) hash lookup comm_buf = input_db.hget(communication_loc, redis_comm) comm = read_communication_from_buffer(comm_buf) else: error('Unknown key type %s' % (key_type)) if comm is None: error(('Unable to find communication with id %s at %s:%s under' ' key %s') % (redis_comm, redis_host, redis_port, communication_loc)) elif redis_comm_index is not None: # lookup comm by index in list comm_buf = input_db.lindex(communication_loc, redis_comm_index) if comm_buf is None: error(('Unable to find communication with id %s at %s:%s under' ' key %s') % (redis_comm, redis_host, redis_port, communication_loc)) else: comm = read_communication_from_buffer(comm_buf) logging.info('%dth Communication has id %s' % (redis_comm_index + 1, comm.id)) else: # take first comm in collection # (or return single comm stored in simple key) for co in reader: comm = co break if comm is None: error('Unable to find any communications at %s:%s under key %s' % (redis_host, redis_port, communication_loc)) return comm
def test_redis_communication_reader_list(mock_read_communication_from_buffer, right_to_left, pop, inferred, block): redis = Mock( llen=Mock(return_value=2), lindex=Mock(side_effect=[sentinel.buf0, sentinel.buf1]), lpop=Mock(side_effect=[sentinel.buf0, sentinel.buf1, None]), rpop=Mock(side_effect=[sentinel.buf0, sentinel.buf1, None]), blpop=Mock(side_effect=[(sentinel.key, sentinel.buf0), (sentinel.key, sentinel.buf1), None]), brpop=Mock(side_effect=[(sentinel.key, sentinel.buf0), (sentinel.key, sentinel.buf1), None]), type=Mock(return_value=b'list'), ) mock_read_communication_from_buffer.side_effect = [ sentinel.comm0, sentinel.comm1 ] reader = RedisCommunicationReader(redis, sentinel.key, key_type=None if inferred else 'list', add_references=sentinel.add_references, right_to_left=right_to_left, pop=pop, block=block, block_timeout=sentinel.block_timeout) it = iter(reader) assert next(it) == sentinel.comm0 assert next(it) == sentinel.comm1 with raises(StopIteration): next(it) if inferred: redis.type.assert_called_once_with(sentinel.key) else: assert not redis.type.called if pop: assert not redis.llen.called assert not redis.lindex.called if block: assert not redis.lpop.called assert not redis.rpop.called getattr(redis, 'brpop' if right_to_left else 'blpop').assert_has_calls([ call(sentinel.key, timeout=sentinel.block_timeout), call(sentinel.key, timeout=sentinel.block_timeout), call(sentinel.key, timeout=sentinel.block_timeout), ]) assert not getattr(redis, 'blpop' if right_to_left else 'brpop').called else: getattr(redis, 'rpop' if right_to_left else 'lpop').assert_has_calls([ call(sentinel.key), call(sentinel.key), call(sentinel.key), ]) assert not getattr(redis, 'lpop' if right_to_left else 'rpop').called assert not redis.blpop.called assert not redis.brpop.called else: redis.llen.assert_called_with(sentinel.key) redis.lindex.assert_has_calls([ call(sentinel.key, -1 if right_to_left else 0), call(sentinel.key, -2 if right_to_left else 1), ]) assert not redis.lpop.called assert not redis.rpop.called assert not redis.blpop.called assert not redis.brpop.called mock_read_communication_from_buffer.assert_has_calls([ call(sentinel.buf0, add_references=sentinel.add_references), call(sentinel.buf1, add_references=sentinel.add_references), ])
def test_redis_communication_reader_failed_key_type(key_type): with raises(ValueError): RedisCommunicationReader(sentinel.redis, sentinel.key, key_type=key_type)
def test_implicit_empty(self): key = 'dataset' with RedisServer(loglevel='warning') as server: redis_db = Redis(port=server.port) with self.assertRaises(Exception): RedisCommunicationReader(redis_db, key)