From 42a5104a1bf8ee4cd9c0f984400189a819f0d873 Mon Sep 17 00:00:00 2001
From: Ermal Kaleci <ermalkaleci@gmail.com>
Date: Tue, 28 Feb 2023 12:35:47 +0100
Subject: [PATCH] upcoming block with timeout (#208)

---
 e2e/src/block.test.ts                        | 15 ++++++++++++++-
 packages/chopsticks/src/blockchain/index.ts  |  6 +++---
 packages/chopsticks/src/blockchain/txpool.ts | 18 +++++++++++++++---
 3 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/e2e/src/block.test.ts b/e2e/src/block.test.ts
index 0bd2aaaa..cede30a3 100644
--- a/e2e/src/block.test.ts
+++ b/e2e/src/block.test.ts
@@ -28,10 +28,23 @@ describe('block', () => {
       dev.newBlock({ count: 3 })
     }, 1000)
     {
-      const next = await chain.upcomingBlock(2)
+      const next = await chain.upcomingBlock({ skipCount: 2 })
       expect(next.number).toEqual(blockNumber + 5)
     }
 
+    setTimeout(() => {
+      dev.newBlock()
+    }, 1000)
+    {
+      const start = Date.now()
+      // no block is build within 1 sec
+      await expect(chain.upcomingBlock({ timeout: 1_000 })).rejects.toThrowError('Timeout has occurred')
+      expect(Date.now() - start).to.be.approximately(1_000, 50)
+
+      const next = await chain.upcomingBlock({ timeout: 10_000 })
+      expect(next.number).toEqual(blockNumber + 6)
+    }
+
     await delay(1000)
   })
 })
diff --git a/packages/chopsticks/src/blockchain/index.ts b/packages/chopsticks/src/blockchain/index.ts
index 061fc006..c7ecc166 100644
--- a/packages/chopsticks/src/blockchain/index.ts
+++ b/packages/chopsticks/src/blockchain/index.ts
@@ -8,7 +8,7 @@ import type { TransactionValidity } from '@polkadot/types/interfaces/txqueue'
 
 import { Api } from '../api'
 import { Block } from './block'
-import { BuildBlockMode, BuildBlockParams, HorizontalMessage, TxPool } from './txpool'
+import { BuildBlockMode, BuildBlockParams, HorizontalMessage, TxPool, UpcomingBlockParams } from './txpool'
 import { HeadState } from './head-state'
 import { InherentProvider } from './inherent'
 import { defaultLogger } from '../logger'
@@ -166,8 +166,8 @@ export class Blockchain {
     return this.#head
   }
 
-  async upcomingBlock(skipCount = 0) {
-    return this.#txpool.upcomingBlock(skipCount)
+  async upcomingBlock(params?: UpcomingBlockParams) {
+    return this.#txpool.upcomingBlock(params)
   }
 
   async dryRunExtrinsic(
diff --git a/packages/chopsticks/src/blockchain/txpool.ts b/packages/chopsticks/src/blockchain/txpool.ts
index 7fc25098..915ec3b9 100644
--- a/packages/chopsticks/src/blockchain/txpool.ts
+++ b/packages/chopsticks/src/blockchain/txpool.ts
@@ -1,7 +1,7 @@
 import { BehaviorSubject, firstValueFrom } from 'rxjs'
 import { EventEmitter } from 'node:stream'
 import { HexString } from '@polkadot/util/types'
-import { skip, take } from 'rxjs/operators'
+import { skip, take, timeout } from 'rxjs/operators'
 import _ from 'lodash'
 
 import { Block } from './block'
@@ -34,6 +34,13 @@ export interface BuildBlockParams {
   }
 }
 
+export interface UpcomingBlockParams {
+  /// upcoming blocks to skip
+  skipCount?: number
+  /// how long to wait in milliseconds
+  timeout?: number
+}
+
 export class TxPool {
   readonly #chain: Blockchain
   readonly #pool: HexString[] = []
@@ -81,9 +88,14 @@ export class TxPool {
     this.#last.next(this.#chain.head)
   }
 
-  async upcomingBlock(skipCount = 0) {
+  async upcomingBlock(params?: UpcomingBlockParams) {
+    const { skipCount, timeout: millisecs } = { skipCount: 0, ...(params || {}) }
     if (skipCount < 0) throw new Error('skipCount needs to be greater or equal to 0')
-    return firstValueFrom(this.#last.pipe(skip(1 + skipCount), take(1)))
+    let stream$ = this.#last.pipe()
+    if (millisecs) {
+      stream$ = stream$.pipe(timeout(millisecs))
+    }
+    return firstValueFrom(stream$.pipe(skip(1 + skipCount), take(1)))
   }
 
   async #buildBlock(wait: Promise<void>, params?: BuildBlockParams) {