From 2eba99ee445a4aee34d0919041388e7e727f7ed8 Mon Sep 17 00:00:00 2001
From: Luke Schoen <ltfschoen@users.noreply.github.com>
Date: Tue, 13 Jun 2023 21:05:50 +1000
Subject: [PATCH] Allow setting block number for off-chain testing (#1806)

* Add set_block_number for offchain test api env to resolve #1805

* Update CHANGELOG with PR#1806
---
 CHANGELOG.md                                |  1 +
 crates/engine/src/exec_context.rs           |  5 +++++
 crates/engine/src/test_api.rs               |  6 ++++++
 crates/engine/src/tests.rs                  | 18 ++++++++++++++++++
 crates/env/src/engine/off_chain/test_api.rs | 10 ++++++++++
 5 files changed, 40 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 92d07cfb7a..5abe96f342 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 - Schema generation - [#1765](https://github.com/paritytech/ink/pull/1765)
+- Add `set_block_number` to off-chain test api `Engine` - [#1806](https://github.com/paritytech/ink/pull/1806)
 
 ### Changed
 - E2E: improve call API, remove `build_message` + callback - [#1782](https://github.com/paritytech/ink/pull/1782)
diff --git a/crates/engine/src/exec_context.rs b/crates/engine/src/exec_context.rs
index c3cd1a3507..27ded1022d 100644
--- a/crates/engine/src/exec_context.rs
+++ b/crates/engine/src/exec_context.rs
@@ -71,6 +71,11 @@ impl ExecContext {
     pub fn set_block_timestamp(&mut self, block_timestamp: BlockTimestamp) {
         self.block_timestamp = block_timestamp
     }
+
+    /// Set the block number for the execution context.
+    pub fn set_block_number(&mut self, block_number: BlockNumber) {
+        self.block_number = block_number
+    }
 }
 
 #[cfg(test)]
diff --git a/crates/engine/src/test_api.rs b/crates/engine/src/test_api.rs
index 04aa9da6a6..2c6d48859e 100644
--- a/crates/engine/src/test_api.rs
+++ b/crates/engine/src/test_api.rs
@@ -17,6 +17,7 @@ use crate::{
     types::{
         AccountId,
         Balance,
+        BlockNumber,
         BlockTimestamp,
     },
     AccountError,
@@ -273,6 +274,11 @@ impl Engine {
     pub fn set_block_timestamp(&mut self, new_block_timestamp: BlockTimestamp) {
         self.exec_context.block_timestamp = new_block_timestamp;
     }
+
+    /// Set the block number for the execution context.
+    pub fn set_block_number(&mut self, new_block_number: BlockNumber) {
+        self.exec_context.block_number = new_block_number;
+    }
 }
 
 #[cfg(test)]
diff --git a/crates/engine/src/tests.rs b/crates/engine/src/tests.rs
index b6bc5f6601..0d0e54359a 100644
--- a/crates/engine/src/tests.rs
+++ b/crates/engine/src/tests.rs
@@ -290,3 +290,21 @@ fn setting_getting_block_timestamp() {
         .expect("decoding value transferred failed");
     assert_eq!(output, new_block_timestamp);
 }
+
+#[test]
+fn setting_getting_block_number() {
+    // given
+    let mut engine = Engine::new();
+    let new_block_number: u32 = 1000;
+    let output = &mut &mut get_buffer()[..];
+
+    // when
+    engine.advance_block();
+    engine.set_block_number(new_block_number);
+    engine.block_number(output);
+
+    // then
+    let output = <u32 as scale::Decode>::decode(&mut &output[..16])
+        .expect("decoding value transferred failed");
+    assert_eq!(output, new_block_number);
+}
diff --git a/crates/env/src/engine/off_chain/test_api.rs b/crates/env/src/engine/off_chain/test_api.rs
index 283d89cdf0..bb4fdce1d0 100644
--- a/crates/env/src/engine/off_chain/test_api.rs
+++ b/crates/env/src/engine/off_chain/test_api.rs
@@ -279,6 +279,16 @@ where
     })
 }
 
+/// Sets the block number for the next [`advance_block`] invocation.
+pub fn set_block_number<T>(value: T::BlockNumber)
+where
+    T: Environment<BlockNumber = u32>,
+{
+    <EnvInstance as OnInstance>::on_instance(|instance| {
+        instance.engine.set_block_number(value);
+    })
+}
+
 /// Runs the given closure test function with the default configuration
 /// for the off-chain environment.
 pub fn run_test<T, F>(f: F) -> Result<()>