1 files changed, 40 insertions, 19 deletions
diff --git a/qa/rpc-tests/test_framework/blockstore.py b/qa/rpc-tests/test_framework/blockstore.py
index 4bc279032b..6120dd574b 100644
--- a/qa/rpc-tests/test_framework/blockstore.py
+++ b/qa/rpc-tests/test_framework/blockstore.py
@@ -13,20 +13,31 @@ class BlockStore(object):
         self.blockDB = dbm.ndbm.open(datadir + "/blocks", 'c')
         self.currentBlock = 0
         self.headers_map = dict()
-    
+
     def close(self):
         self.blockDB.close()
 
+    def erase(self, blockhash):
+        del self.blockDB[repr(blockhash)]
+
+    # lookup an entry and return the item as raw bytes
     def get(self, blockhash):
-        serialized_block = None
+        value = None
         try:
-            serialized_block = self.blockDB[repr(blockhash)]
+            value = self.blockDB[repr(blockhash)]
         except KeyError:
             return None
-        f = BytesIO(serialized_block)
-        ret = CBlock()
-        ret.deserialize(f)
-        ret.calc_sha256()
+        return value
+
+    # lookup an entry and return it as a CBlock
+    def get_block(self, blockhash):
+        ret = None
+        serialized_block = self.get(blockhash)
+        if serialized_block is not None:
+            f = BytesIO(serialized_block)
+            ret = CBlock()
+            ret.deserialize(f)
+            ret.calc_sha256()
         return ret
 
     def get_header(self, blockhash):
@@ -75,13 +86,16 @@ class BlockStore(object):
     def add_header(self, header):
         self.headers_map[header.sha256] = header
 
+    # lookup the hashes in "inv", and return p2p messages for delivering
+    # blocks found.
     def get_blocks(self, inv):
         responses = []
         for i in inv:
             if (i.type == 2): # MSG_BLOCK
-                block = self.get(i.hash)
-                if block is not None:
-                    responses.append(msg_block(block))
+                data = self.get(i.hash)
+                if data is not None:
+                    # Use msg_generic to avoid re-serialization
+                    responses.append(msg_generic(b"block", data))
         return responses
 
     def get_locator(self, current_tip=None):
@@ -90,11 +104,11 @@ class BlockStore(object):
         r = []
         counter = 0
         step = 1
-        lastBlock = self.get(current_tip)
+        lastBlock = self.get_block(current_tip)
         while lastBlock is not None:
             r.append(lastBlock.hashPrevBlock)
             for i in range(step):
-                lastBlock = self.get(lastBlock.hashPrevBlock)
+                lastBlock = self.get_block(lastBlock.hashPrevBlock)
                 if lastBlock is None:
                     break
             counter += 1
@@ -111,16 +125,23 @@ class TxStore(object):
     def close(self):
         self.txDB.close()
 
+    # lookup an entry and return the item as raw bytes
     def get(self, txhash):
-        serialized_tx = None
+        value = None
         try:
-            serialized_tx = self.txDB[repr(txhash)]
+            value = self.txDB[repr(txhash)]
         except KeyError:
             return None
-        f = BytesIO(serialized_tx)
-        ret = CTransaction()
-        ret.deserialize(f)
-        ret.calc_sha256()
+        return value
+
+    def get_transaction(self, txhash):
+        ret = None
+        serialized_tx = self.get(txhash)
+        if serialized_tx is not None:
+            f = BytesIO(serialized_tx)
+            ret = CTransaction()
+            ret.deserialize(f)
+            ret.calc_sha256()
         return ret
 
     def add_transaction(self, tx):
@@ -136,5 +157,5 @@ class TxStore(object):
             if (i.type == 1): # MSG_TX
                 tx = self.get(i.hash)
                 if tx is not None:
-                    responses.append(msg_tx(tx))
+                    responses.append(msg_generic(b"tx", tx))
         return responses