implemented fetch block in blockio

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 <https://www.gnu.org/licenses/>.
-"""
+"""

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 <https://www.gnu.org/licenses/>.
 """
 
-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)
+
 

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()