From f44ca60c6ef6140873aa5668cd7a1c9489bd1d54 Mon Sep 17 00:00:00 2001
From: Kevin Froman <beardog@mailbox.org>
Date: Sun, 17 Jan 2021 19:44:57 +0000
Subject: [PATCH] implemented fetch block in blockio

---
 src/blockio/__init__.py       |  3 +-
 src/blockio/store/__init__.py | 20 ++++++++++----
 tests/test_blockio.py         | 52 +++++++++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 7 deletions(-)
 create mode 100644 tests/test_blockio.py

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