improve public interface (#336)

* improve public interface

* fmt

* fix

packages/chopsticks/src/cli.ts

@@ -9,7 +9,7 @@ import yaml from 'js-yaml'
 import yargs from 'yargs'
 
 import { Blockchain, BuildBlockMode, connectParachains, connectVertical, setup, setupWithServer } from '.'
-import { configSchema } from './schema'
+import { Config, configSchema } from './schema'
 import { decodeKey } from './utils/decoder'
 import { dryRun } from './dry-run'
 import { dryRunPreimage } from './dry-run-preimage'
@@ -22,7 +22,7 @@ dotenvConfig()
 
 const CONFIGS_BASE_URL = 'https://raw.githubusercontent.com/AcalaNetwork/chopsticks/master/configs/'
 
-const processConfig = async (path: string) => {
+export const fetchConfig = async (path: string): Promise<Config> => {
   let file: string
   if (isUrl(path)) {
     file = await axios.get(path).then((x) => x.data)
@@ -48,7 +48,7 @@ const processConfig = async (path: string) => {
 
 const processArgv = async (argv: any) => {
   if (argv.config) {
-    argv = { ...(await processConfig(argv.config)), ...argv }
+    argv = { ...(await fetchConfig(argv.config)), ...argv }
   }
   argv.port = argv.port ?? (process.env.PORT ? Number(process.env.PORT) : 8000)
   return argv
@@ -264,7 +264,7 @@ yargs(hideBin(process.argv))
     async (argv) => {
       const parachains: Blockchain[] = []
       for (const config of argv.parachain) {
-        const { chain } = await setupWithServer(await processConfig(config))
+        const { chain } = await setupWithServer(await fetchConfig(config))
         parachains.push(chain)
       }
 
@@ -273,7 +273,7 @@ yargs(hideBin(process.argv))
       }
 
       if (argv.relaychain) {
-        const { chain: relaychain } = await setupWithServer(await processConfig(argv.relaychain))
+        const { chain: relaychain } = await setupWithServer(await fetchConfig(argv.relaychain))
         for (const parachain of parachains) {
           await connectVertical(relaychain, parachain)
         }

packages/chopsticks/src/index.ts

@@ -4,4 +4,5 @@ export { BuildBlockMode } from './blockchain/txpool'
 export { connectParachains, connectVertical } from './xcm'
 export { setup } from './setup'
 export { setupWithServer } from './setup-with-server'
+export { fetchConfig } from './cli'
 export * from './blockchain/inherent'

packages/testing/src/index.ts

@@ -1,6 +1,13 @@
 import { ApiPromise, WsProvider } from '@polkadot/api'
-import { BuildBlockMode, setupWithServer } from '@acala-network/chopsticks'
+import {
+  BuildBlockMode,
+  connectParachains,
+  connectVertical,
+  fetchConfig,
+  setupWithServer,
+} from '@acala-network/chopsticks'
 import { Codec } from '@polkadot/types/types'
+import { Config } from '@acala-network/chopsticks/schema'
 import { HexString } from '@polkadot/util/types'
 import { Keyring, createTestKeyring } from '@polkadot/keyring'
 import { StorageValues } from '@acala-network/chopsticks/utils/set-storage'
@@ -16,11 +23,24 @@ export type SetupOption = {
   wasmOverride?: string
   db?: string
   timeout?: number
+  port?: number
 }
 
-export const setupContext = async ({ endpoint, blockNumber, blockHash, wasmOverride, db, timeout }: SetupOption) => {
-  // random port
-  const port = Math.floor(Math.random() * 10000) + 10000
+export type SetupConfig = Config & {
+  timeout?: number
+}
+
+export const createConfig = ({
+  endpoint,
+  blockNumber,
+  blockHash,
+  wasmOverride,
+  db,
+  timeout,
+  port,
+}: SetupOption): SetupConfig => {
+  // random port if not specified
+  port = port ?? Math.floor(Math.random() * 10000) + 10000
   const config = {
     endpoint,
     port,
@@ -29,7 +49,16 @@ export const setupContext = async ({ endpoint, blockNumber, blockHash, wasmOverr
     'build-block-mode': BuildBlockMode.Manual,
     db,
     'wasm-override': wasmOverride,
+    timeout,
   }
+  return config
+}
+
+export const setupContext = async (option: SetupOption) => {
+  return setupContextWithConfig(createConfig(option))
+}
+
+export const setupContextWithConfig = async ({ timeout, ...config }: SetupConfig) => {
   const { chain, listenPort, close } = await setupWithServer(config)
 
   const url = `ws://localhost:${listenPort}`
@@ -80,6 +109,47 @@ export const setupContext = async ({ endpoint, blockNumber, blockHash, wasmOverr
   }
 }
 
+export type NetworkContext = Awaited<ReturnType<typeof setupContext>>
+
+export const setupNetworks = async (networkOptions: Partial<Record<string, Config | string | undefined>>) => {
+  const ret = {} as Record<string, NetworkContext>
+
+  let wasmOverriden = false
+
+  for (const [name, options] of Object.entries(networkOptions) as [string, Config | string | undefined][]) {
+    const config = typeof options === 'string' ? await fetchConfig(options) : options ?? (await fetchConfig(name))
+    ret[name] = await setupContextWithConfig(config)
+    wasmOverriden ||= config['wasm-override'] != null
+  }
+
+  const relaychainName = Object.keys(ret).filter((x) => ['polkadot', 'kusama'].includes(x.toLocaleLowerCase()))[0]
+  const { [relaychainName]: relaychain, ...parachains } = ret
+
+  if (relaychain) {
+    for (const parachain of Object.values(parachains)) {
+      await connectVertical(relaychain.chain, parachain.chain)
+    }
+  }
+
+  const parachainList = Object.values(parachains).map((i) => i.chain)
+  if (parachainList.length > 0) {
+    await connectParachains(parachainList)
+  }
+
+  if (wasmOverriden) {
+    // trigger runtime upgrade if needed (due to wasm override)
+    for (const chain of Object.values(ret)) {
+      await chain.dev.newBlock()
+    }
+    // handle xcm version message if needed (due to wasm override triggered xcm version upgrade)
+    for (const chain of Object.values(ret)) {
+      await chain.dev.newBlock()
+    }
+  }
+
+  return ret
+}
+
 type CodecOrArray = Codec | Codec[]
 
 const processCodecOrArray = (codec: CodecOrArray, fn: (c: Codec) => any) =>