update run-block command

packages/chopsticks/src/cli.ts

@@ -64,14 +64,14 @@ yargs(hideBin(process.argv))
     (yargs) =>
       yargs.options({
         ...defaultOptions,
-        port: {
-          desc: 'Port to listen on',
-          number: true,
-        },
         'output-path': {
           desc: 'File path to print output',
           string: true,
         },
+        'import-storage': {
+          desc: 'Pre-defined JSON/YAML storage file path',
+          string: true,
+        },
         html: {
           desc: 'Generate html with storage diff',
         },

packages/chopsticks/src/run-block.ts

@@ -8,7 +8,7 @@ import { runTask, taskHandler } from './executor'
 import { setup } from './setup'
 
 export const runBlock = async (argv: Config) => {
-  const context = await setup(argv)
+  const context = await setup(argv, true)
 
   const header = await context.chain.head.header
   const wasm = await context.chain.head.wasm

packages/chopsticks/src/setup.ts

@@ -21,7 +21,7 @@ import { importStorage, overrideWasm } from './utils/import-storage'
 import { openDb } from './db'
 import { timeTravel } from './utils/time-travel'
 
-export const setup = async (argv: Config) => {
+export const setup = async (argv: Config, runBlock = false) => {
   let provider: ProviderInterface
   if (argv.genesis) {
     if (typeof argv.genesis === 'string') {
@@ -78,10 +78,18 @@ export const setup = async (argv: Config) => {
 
   if (argv.timestamp) await timeTravel(chain, argv.timestamp)
 
+  let at: HexString | undefined
+  if (runBlock) {
+    // in case of run block we need to apply wasm-override and import-storage to parent block
+    const block = await chain.head.parentBlock
+    if (!block) throw new Error('Cannot find parent block')
+    at = block.hash
+  }
+
   // override wasm before importing storage, in case new pallets have been
   // added that have storage imports
-  await overrideWasm(chain, argv['wasm-override'])
-  await importStorage(chain, argv['import-storage'])
+  await overrideWasm(chain, argv['wasm-override'], at)
+  await importStorage(chain, argv['import-storage'], at)
 
   return { chain, api, ws: provider }
 }

packages/chopsticks/src/utils/import-storage.ts

@@ -6,7 +6,7 @@ import { Blockchain } from '../blockchain'
 import { StorageValues, setStorage } from './set-storage'
 import { defaultLogger } from '../logger'
 
-export const importStorage = async (chain: Blockchain, storage?: string | StorageValues) => {
+export const importStorage = async (chain: Blockchain, storage?: string | StorageValues, at?: HexString) => {
   if (storage == null) {
     return
   }
@@ -17,11 +17,11 @@ export const importStorage = async (chain: Blockchain, storage?: string | Storag
   } else {
     storageValue = storage
   }
-  const blockHash = await setStorage(chain, storageValue)
+  const blockHash = await setStorage(chain, storageValue, at)
   defaultLogger.trace({ blockHash, storage }, 'ImportStorage')
 }
 
-export const overrideWasm = async (chain: Blockchain, wasmPath?: string) => {
+export const overrideWasm = async (chain: Blockchain, wasmPath?: string, at?: HexString) => {
   if (wasmPath == null) {
     return
   }
@@ -33,5 +33,11 @@ export const overrideWasm = async (chain: Blockchain, wasmPath?: string) => {
   } else {
     wasmHex = '0x' + wasm.toString('hex')
   }
-  chain.head.setWasm(wasmHex as HexString)
+  if (at) {
+    const block = await chain.getBlock(at)
+    if (!block) throw new Error(`Cannot find block ${at}`)
+    block.setWasm(wasmHex as HexString)
+  } else {
+    chain.head.setWasm(wasmHex as HexString)
+  }
 }