From b7b5f5e7ff25c231b562f957edbf295f7513b11a Mon Sep 17 00:00:00 2001
From: Scotty <66335769+ScottyPoi@users.noreply.github.com>
Date: Fri, 13 May 2022 05:29:29 -0600
Subject: [PATCH] blockchain: fix iterator return type, reinstate tests (#1877)

* blockchain: remove 'void' from iterator() return

* vm: remove void from vm.runBlockchain() return types

* client: adapt VMExecution to stricter blockchain.iterator() call

* vm: reintroduce deleted tests for blockchain.getIteratorHead()

* vm: correct method name in test
---
 packages/blockchain/src/index.ts             | 2 +-
 packages/client/lib/execution/vmexecution.ts | 7 +++----
 packages/vm/src/runBlockchain.ts             | 2 +-
 packages/vm/tests/api/runBlockchain.spec.ts  | 9 ++++++++-
 4 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/packages/blockchain/src/index.ts b/packages/blockchain/src/index.ts
index 19f100778c6..54d5327fbc6 100644
--- a/packages/blockchain/src/index.ts
+++ b/packages/blockchain/src/index.ts
@@ -41,7 +41,7 @@ export interface BlockchainInterface {
    * @param onBlock - Function called on each block with params (block: Block,
    * reorg: boolean)
    */
-  iterator(name: string, onBlock: OnBlock): Promise<void | number>
+  iterator(name: string, onBlock: OnBlock): Promise<number>
 }
 
 /**
diff --git a/packages/client/lib/execution/vmexecution.ts b/packages/client/lib/execution/vmexecution.ts
index 525f202321e..b22fa4431a8 100644
--- a/packages/client/lib/execution/vmexecution.ts
+++ b/packages/client/lib/execution/vmexecution.ts
@@ -24,7 +24,7 @@ export class VMExecution extends Execution {
 
   public receiptsManager?: ReceiptsManager
   private pendingReceipts?: Map<string, TxReceipt[]>
-  private vmPromise?: Promise<number | undefined>
+  private vmPromise?: Promise<number>
 
   /** Number of maximum blocks to run per iteration of {@link VMExecution.run} */
   private NUM_BLOCKS_PER_ITERATION = 50
@@ -157,7 +157,6 @@ export class VMExecution extends Execution {
       headBlock = undefined
       parentState = undefined
       errorBlock = undefined
-
       this.vmPromise = blockchain.iterator(
         'vm',
         async (block: Block, reorg: boolean) => {
@@ -247,7 +246,7 @@ export class VMExecution extends Execution {
         },
         this.NUM_BLOCKS_PER_ITERATION
       )
-      numExecuted = (await this.vmPromise) as number
+      numExecuted = await this.vmPromise
 
       if (errorBlock) {
         await this.chain.blockchain.setIteratorHead('vm', (errorBlock as Block).header.parentHash)
@@ -255,7 +254,7 @@ export class VMExecution extends Execution {
       }
 
       const endHeadBlock = await this.vm.blockchain.getIteratorHead('vm')
-      if (numExecuted > 0) {
+      if (numExecuted && numExecuted > 0) {
         const firstNumber = startHeadBlock.header.number
         const firstHash = short(startHeadBlock.hash())
         const lastNumber = endHeadBlock.header.number
diff --git a/packages/vm/src/runBlockchain.ts b/packages/vm/src/runBlockchain.ts
index 4a61a9d2670..900c0f25dc1 100644
--- a/packages/vm/src/runBlockchain.ts
+++ b/packages/vm/src/runBlockchain.ts
@@ -9,7 +9,7 @@ export default async function runBlockchain(
   this: VM,
   blockchain?: Blockchain,
   maxBlocks?: number
-): Promise<void | number> {
+): Promise<number> {
   let headBlock: Block
   let parentState: Buffer
 
diff --git a/packages/vm/tests/api/runBlockchain.spec.ts b/packages/vm/tests/api/runBlockchain.spec.ts
index 1b8563b091a..c3931ec4e01 100644
--- a/packages/vm/tests/api/runBlockchain.spec.ts
+++ b/packages/vm/tests/api/runBlockchain.spec.ts
@@ -101,6 +101,9 @@ tape('runBlockchain', (t) => {
 
     await vm.runBlockchain()
 
+    const head = await vm.blockchain.getIteratorHead()
+    st.equal(head.hash().toString('hex'), testData.blocks[0].blockHeader.hash.slice(2))
+
     st.end()
   })
 
@@ -138,12 +141,16 @@ tape('runBlockchain', (t) => {
     await blockchain.putBlock(b2)
     await blockchain.putBlock(b3)
 
+    let head = await blockchain.getIteratorHead()
+    st.deepEqual(head.hash(), genesisBlock.hash(), 'Iterator head should still be at genesis')
+
     try {
       await vm.runBlockchain()
       st.fail('should have returned error')
     } catch (e: any) {
       st.equal(e.message, 'test')
-
+      head = await blockchain.getIteratorHead()
+      st.deepEqual(head.hash(), b2.hash(), 'should have removed invalid block from head')
       st.end()
     }
   })