Tests: Rework blockstore to avoid re-serialization.

2 files changed, 53 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
diff --git a/qa/rpc-tests/test_framework/mininode.py b/qa/rpc-tests/test_framework/mininode.py
index 1617daa200..e85399c960 100755
--- a/qa/rpc-tests/test_framework/mininode.py
+++ b/qa/rpc-tests/test_framework/mininode.py
@@ -556,6 +556,7 @@ class CBlock(CBlockHeader):
             self.nNonce += 1
             self.rehash()
 
+
     def __repr__(self):
         return "CBlock(nVersion=%i hashPrevBlock=%064x hashMerkleRoot=%064x nTime=%s nBits=%08x nNonce=%08x vtx=%s)" \
             % (self.nVersion, self.hashPrevBlock, self.hashMerkleRoot,
@@ -836,6 +837,18 @@ class msg_block(object):
     def __repr__(self):
         return "msg_block(block=%s)" % (repr(self.block))
 
+# for cases where a user needs tighter control over what is sent over the wire
+# note that the user must supply the name of the command, and the data
+class msg_generic(object):
+    def __init__(self, command, data=None):
+        self.command = command
+        self.data = data
+
+    def serialize(self):
+        return self.data
+
+    def __repr__(self):
+        return "msg_generic()"
 
 class msg_getaddr(object):
     command = b"getaddr"