diff --git a/src/blockio/__init__.py b/src/blockio/__init__.py
index 72940344..f88116e2 100644
--- a/src/blockio/__init__.py
+++ b/src/blockio/__init__.py
@@ -2,6 +2,7 @@
Wrap safedb for storing and fetching blocks
"""
+from .store import store_block
"""
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
@@ -15,4 +16,4 @@ GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see .
-"""
\ No newline at end of file
+"""
diff --git a/src/blockio/store/__init__.py b/src/blockio/store/__init__.py
index 251b361e..9e01af1a 100644
--- a/src/blockio/store/__init__.py
+++ b/src/blockio/store/__init__.py
@@ -2,9 +2,7 @@
Store blocks and cache meta info such as block type
"""
-from typing import TYPE_CHECKING, Union, NewType
-
-from safedb import DBProtectionOpeningModeError
+from typing import TYPE_CHECKING
if TYPE_CHECKING:
from kasten import Kasten
@@ -24,10 +22,20 @@ You should have received a copy of the GNU General Public License
along with this program. If not, see .
"""
-RawBlock = NewType('RawBlock', bytes)
+def store_block(block: 'Kasten', safe_db: 'SafeDB'):
+
+ block_type = block.get_data_type()
+ try:
+ block_list_for_type = safe_db.get(f'bl-{block_type}')
+ if block.id in block_list_for_type:
+ raise ValueError("Cannot store duplicate block")
+ except KeyError:
+ block_list_for_type = b''
-def store_block(block: Kasten, safe_db: SafeDB):
-
safe_db.put(block.id, block.get_packed())
+ # Append the block to the list of blocks for this given type
+ block_list_for_type += block.id
+ safe_db.put(f'bl-{block_type}', block_list_for_type)
+
diff --git a/tests/test_blockio.py b/tests/test_blockio.py
new file mode 100644
index 00000000..8a61a0ee
--- /dev/null
+++ b/tests/test_blockio.py
@@ -0,0 +1,52 @@
+#!/usr/bin/env python3
+import sys, os
+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 identifyhome, createdirs, bettersleep
+from onionrsetup import setup_config, setup_default_plugins
+
+createdirs.create_dirs()
+setup_config()
+setup_default_plugins()
+
+import kasten
+from onionrblocks.generators import anonvdf
+from utils import identifyhome
+
+import safedb
+import blockio
+
+def _remove_db(path):
+ try:
+ os.remove(path)
+ except FileNotFoundError:
+ pass
+
+class TestBlockIO(unittest.TestCase):
+ def test_store_block(self):
+ packed = kasten.generator.pack.pack(b"test", "tst")
+ bl: kasten.Kasten = anonvdf.AnonVDFGenerator.generate(packed, rounds=1000)
+ db_file = identifyhome.identify_home() + 'test.db'
+ db = safedb.SafeDB(db_file)
+ blockio.store_block(bl, db)
+ db.close()
+ _remove_db(db_file)
+
+ def test_store_dupe(self):
+ packed = kasten.generator.pack.pack(b"test", "tst")
+ bl: kasten.Kasten = anonvdf.AnonVDFGenerator.generate(packed, rounds=1000)
+ db_file = identifyhome.identify_home() + 'test.db'
+ db = safedb.SafeDB(db_file)
+ blockio.store_block(bl, db)
+ self.assertRaises(ValueError, blockio.store_block, bl, db)
+ db.close()
+ _remove_db(db_file)
+
+
+
+unittest.main()