diff --git a/src/blockdb/__init__.py b/src/blockdb/__init__.py index 6b98c5ed..900684f9 100644 --- a/src/blockdb/__init__.py +++ b/src/blockdb/__init__.py @@ -11,7 +11,7 @@ def store_vdf_block(block: Block): db.set(block_db_path, block.id, block.raw) -def get_blocks_by_type(block_type): +def get_blocks_by_type(block_type: str): block_db = db.get_db_obj(block_db_path, 'u') for block_hash in db.list_keys(block_db_path): block = Block(block_hash, block_db[block_hash], auto_verify=False) @@ -19,8 +19,9 @@ def get_blocks_by_type(block_type): yield block -def get_blocks_after_timestamp(timestamp: int, block_type: bytes = ''): +def get_blocks_after_timestamp(timestamp: int, block_type: str = ''): block_db = db.get_db_obj(block_db_path, 'u') + for block_hash in db.list_keys(block_db_path): block = Block(block_hash, block_db[block_hash], auto_verify=False) if block.timestamp > timestamp: diff --git a/tests/test_blockdb.py b/tests/test_blockdb.py new file mode 100644 index 00000000..0b565b2a --- /dev/null +++ b/tests/test_blockdb.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 +import sys, os +import dbm +sys.path.append(".") +sys.path.append("src/") +import uuid +TEST_DIR = 'testdata/%s-%s' % (uuid.uuid4(), os.path.basename(__file__)) + '/' +print("Test directory:", TEST_DIR) +os.environ["ONIONR_HOME"] = TEST_DIR +import unittest + +from utils import createdirs +createdirs.create_dirs() + + +import onionrblocks + +import blockdb + + +class TestBlockDB(unittest.TestCase): + def test_store_vdf_block(self): + bl: Block = onionrblocks.create_anonvdf_block(os.urandom(10), b'bin', 2500) + blockdb.store_vdf_block(bl) + + with dbm.open(blockdb.block_db_path, 'r') as b_db: + b_db[bl.id] + + def test_get_blocks_by_type(self): + with dbm.open(blockdb.block_db_path, 'c') as b_db: + bl: Block = onionrblocks.create_anonvdf_block('test', b'txt', 2500) + b_db[bl.id] = bl.raw + looped = False + for block in blockdb.get_blocks_by_type('txt'): + looped = True + self.assertEqual(bl.id, block.id) + self.assertTrue(looped) + + + +unittest.main()