def test_put(self):
"""Test put method of hook"""
hook = LevelDBHook(leveldb_conn_id='leveldb_default')
db = hook.get_conn(name='/tmp/testdb/', create_if_missing=True)
hook.put(b'test_key2', b'test_value2')
assert db.get(b'test_key2') == b'test_value2'
hook.close_conn()
def test_delete(self):
"""Test delete method of hook"""
hook = LevelDBHook(leveldb_conn_id='leveldb_default')
db = hook.get_conn(name='/tmp/testdb/', create_if_missing=True)
db.put(b'test_key3', b'test_value3')
hook.delete(b'test_key3')
assert db.get(b'test_key3') is None
hook.close_conn()
def test_write_batch(self):
"""Test write batch method of hook"""
hook = LevelDBHook(leveldb_conn_id='leveldb_default')
db = hook.get_conn(name='/tmp/testdb/', create_if_missing=True)
keys = [b'key', b'another-key']
values = [b'value', b'another-value']
hook.write_batch(keys, values)
assert db.get(b'key') == b'value'
assert db.get(b'another-key') == b'another-value'
hook.close_conn()
def test_comparator(self):
"""Test comparator"""
def comparator(a, b):
a = a.lower()
b = b.lower()
if a < b:
# a sorts before b
return -1
if a > b:
# a sorts after b
return 1
# a and b are equal
return 0
hook = LevelDBHook(leveldb_conn_id='leveldb_default')
hook.get_conn(
name='/tmp/testdb2/',
create_if_missing=True,
comparator=comparator,
comparator_name=b'CaseInsensitiveComparator',
)
assert hook.db is not None, "Check existence of DB object(with comparator) in connection creation"
hook.close_conn()
def execute(self, context) -> Optional[str]:
"""
Execute command in LevelDB
:returns: value from get(str, not bytes, to prevent error in json.dumps in serialize_value in xcom.py)
or None(Optional[str])
:rtype: Optional[str]
"""
leveldb_hook = LevelDBHook(leveldb_conn_id=self.leveldb_conn_id)
leveldb_hook.get_conn(name=self.name,
create_if_missing=self.create_if_missing,
**self.create_db_extra_options)
value = leveldb_hook.run(
command=self.command,
key=self.key,
value=self.value,
keys=self.keys,
values=self.values,
)
self.log.info("Done. Returned value was: %s", str(value))
leveldb_hook.close_conn()
value = value if value is None else value.decode()
return value
def test_exception(self):
"""Test raising exception of hook in run method if we have unknown command in input"""
hook = LevelDBHook(leveldb_conn_id='leveldb_default')
hook.get_conn(name='/tmp/testdb/', create_if_missing=True)
with pytest.raises(LevelDBHookException):
hook.run(command='other_command', key=b'key', value=b'value')
def test_run(self):
"""Test run method of hook"""
hook = LevelDBHook(leveldb_conn_id='leveldb_default')
hook.get_conn(name='/tmp/testdb/', create_if_missing=True)
assert hook.run(
'get',
b'test_key0') is None, "Initially, this key in LevelDB is empty"
hook.run('put', b'test_key0', b'test_value0')
assert (hook.run('get', b'test_key0') == b'test_value0'
), 'Connection to LevelDB with PUT and GET works.'
hook.run('delete', b'test_key0')
assert hook.run(
'get',
b'test_key0') is None, 'Connection to LevelDB with DELETE works.'
hook.close_conn()
def test_get_conn_db_is_not_none(self):
"""Test get_conn method of hook"""
hook = LevelDBHook(leveldb_conn_id='leveldb_default')
hook.get_conn(name='/tmp/testdb/', create_if_missing=True)
assert hook.db is not None, "Check existence of DB object in connection creation"
hook.close_conn()