From 371f58299614d458f5b7266b88a0b7211dc2cb48 Mon Sep 17 00:00:00 2001
From: Ermal Kaleci <ermalkaleci@gmail.com>
Date: Thu, 26 Jan 2023 19:36:04 +0100
Subject: [PATCH 1/2] fix upcoming block

---
 e2e/block.test.ts        | 37 +++++++++++++++++++++++++++++++++++++
 e2e/helper.ts            |  3 ++-
 src/blockchain/index.ts  |  4 ++--
 src/blockchain/txpool.ts | 15 +++++++--------
 4 files changed, 48 insertions(+), 11 deletions(-)
 create mode 100644 e2e/block.test.ts

diff --git a/e2e/block.test.ts b/e2e/block.test.ts
new file mode 100644
index 00000000..0bd2aaaa
--- /dev/null
+++ b/e2e/block.test.ts
@@ -0,0 +1,37 @@
+import { describe, expect, it } from 'vitest'
+
+import { chain, delay, dev, setupApi } from './helper'
+
+setupApi({ endpoint: 'wss://rpc.polkadot.io' })
+
+describe('block', () => {
+  it('upcoming block works', async () => {
+    const blockNumber = chain.head.number
+
+    setTimeout(() => {
+      dev.newBlock()
+    }, 1000)
+    {
+      const next = await chain.upcomingBlock()
+      expect(next.number).toEqual(blockNumber + 1)
+    }
+
+    setTimeout(() => {
+      dev.newBlock()
+    }, 1000)
+    {
+      const next = await chain.upcomingBlock()
+      expect(next.number).toEqual(blockNumber + 2)
+    }
+
+    setTimeout(() => {
+      dev.newBlock({ count: 3 })
+    }, 1000)
+    {
+      const next = await chain.upcomingBlock(2)
+      expect(next.number).toEqual(blockNumber + 5)
+    }
+
+    await delay(1000)
+  })
+})
diff --git a/e2e/helper.ts b/e2e/helper.ts
index 1bad0b4d..76d84e48 100644
--- a/e2e/helper.ts
+++ b/e2e/helper.ts
@@ -121,7 +121,8 @@ export const setupApi = (option: SetupOption) => {
   beforeAll(async () => {
     const res = await setupAll(option)
     setup = res.setup
-    return () => res.teardownAll()
+
+    return res.teardownAll
   })
 
   beforeEach(async () => {
diff --git a/src/blockchain/index.ts b/src/blockchain/index.ts
index 36d4fe00..4f1096f9 100644
--- a/src/blockchain/index.ts
+++ b/src/blockchain/index.ts
@@ -161,8 +161,8 @@ export class Blockchain {
     return this.#head
   }
 
-  async upcomingBlock(count = 1) {
-    return this.#txpool.upcomingBlock(count)
+  async upcomingBlock(skipCount = 0) {
+    return this.#txpool.upcomingBlock(skipCount)
   }
 
   async dryRunExtrinsic(
diff --git a/src/blockchain/txpool.ts b/src/blockchain/txpool.ts
index 725c2d0a..b475a9c4 100644
--- a/src/blockchain/txpool.ts
+++ b/src/blockchain/txpool.ts
@@ -1,6 +1,6 @@
+import { BehaviorSubject, firstValueFrom } from 'rxjs'
 import { HexString } from '@polkadot/util/types'
-import { ReplaySubject, firstValueFrom, share } from 'rxjs'
-import { first, skip } from 'rxjs/operators'
+import { skip, take } from 'rxjs/operators'
 import _ from 'lodash'
 
 import { Block } from './block'
@@ -37,13 +37,12 @@ export class TxPool {
   readonly #mode: BuildBlockMode
   readonly #inherentProvider: InherentProvider
 
+  #last: BehaviorSubject<Block>
   #lastBuildBlockPromise: Promise<void> = Promise.resolve()
 
-  #last = new ReplaySubject<Block>(1)
-  #upcoming = this.#last.pipe(share())
-
   constructor(chain: Blockchain, inherentProvider: InherentProvider, mode: BuildBlockMode = BuildBlockMode.Batch) {
     this.#chain = chain
+    this.#last = new BehaviorSubject<Block>(chain.head)
     this.#mode = mode
     this.#inherentProvider = inherentProvider
   }
@@ -77,9 +76,9 @@ export class TxPool {
     this.#last.next(this.#chain.head)
   }
 
-  async upcomingBlock(count = 1) {
-    if (count < 1) throw new Error('count needs to be greater than 0')
-    return firstValueFrom(this.#upcoming.pipe(skip(count - 1), first()))
+  async upcomingBlock(skipCount = 0) {
+    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)))
   }
 
   async #buildBlock(wait: Promise<void>, params?: BuildBlockParams) {

From 359ef210fc2cbf1df31690b22a35dd62ba24ddd0 Mon Sep 17 00:00:00 2001
From: Ermal Kaleci <ermalkaleci@gmail.com>
Date: Thu, 26 Jan 2023 19:48:39 +0100
Subject: [PATCH 2/2] bump

---
 package.json | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/package.json b/package.json
index e950657a..e645a7dc 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
 	"name": "@acala-network/chopsticks",
-	"version": "0.3.8",
+	"version": "0.3.9",
 	"main": "./dist/index.js",
 	"types": "./dist/index.d.ts",
 	"author": "Bryan Chen <xlchen1291@gmail.com>",