From 1d5281ae7b24a9a4bcc53ceccfa246b85392dbc7 Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Fri, 8 Sep 2023 15:33:50 +0100
Subject: [PATCH 01/14] WIP

---
 .../set-code-hash-migration/Cargo.toml        |  23 +++
 .../set-code-hash-migration/lib.rs            | 133 ++++++++++++++++++
 .../migration/Cargo.toml                      |  19 +++
 .../set-code-hash-migration/migration/lib.rs  |  51 +++++++
 .../updated-incrementer/Cargo.toml            |  19 +++
 .../updated-incrementer/lib.rs                |  75 ++++++++++
 6 files changed, 320 insertions(+)
 create mode 100644 integration-tests/upgradeable-contracts/set-code-hash-migration/Cargo.toml
 create mode 100644 integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
 create mode 100644 integration-tests/upgradeable-contracts/set-code-hash-migration/migration/Cargo.toml
 create mode 100644 integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs
 create mode 100644 integration-tests/upgradeable-contracts/set-code-hash-migration/updated-incrementer/Cargo.toml
 create mode 100644 integration-tests/upgradeable-contracts/set-code-hash-migration/updated-incrementer/lib.rs

diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/Cargo.toml b/integration-tests/upgradeable-contracts/set-code-hash-migration/Cargo.toml
new file mode 100644
index 00000000000..79c4aad15f0
--- /dev/null
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/Cargo.toml
@@ -0,0 +1,23 @@
+[package]
+name = "incrementer"
+version = "4.2.0"
+authors = ["Parity Technologies <admin@parity.io>"]
+edition = "2021"
+publish = false
+
+[dependencies]
+ink = { path = "../../../crates/ink", default-features = false }
+
+[dev-dependencies]
+ink_e2e = { path = "../../../crates/e2e" }
+
+[lib]
+path = "lib.rs"
+
+[features]
+default = ["std"]
+std = [
+    "ink/std",
+]
+ink-as-dependency = []
+e2e-tests = []
diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
new file mode 100644
index 00000000000..86851852d56
--- /dev/null
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
@@ -0,0 +1,133 @@
+#![cfg_attr(not(feature = "std"), no_std, no_main)]
+
+//! Demonstrates how to use [`set_code_hash`](https://docs.rs/ink_env/latest/ink_env/fn.set_code_hash.html)
+//! to swap out the `code_hash` of an on-chain contract.
+//!
+//! We will swap the code of our `Incrementer` contract with that of the an `Incrementer`
+//! found in the `updated_incrementer` folder.
+//!
+//! See the included End-to-End tests an example update workflow.
+
+#[ink::contract]
+pub mod incrementer {
+
+    /// Track a counter in storage.
+    ///
+    /// # Note
+    ///
+    /// Is is important to realize that after the call to `set_code_hash` the contract's
+    /// storage remains the same.
+    ///
+    /// If you change the storage layout in your storage struct you may introduce
+    /// undefined behavior to your contract!
+    #[ink(storage)]
+    #[derive(Default)]
+    pub struct Incrementer {
+        count: u32,
+    }
+
+    impl Incrementer {
+        /// Creates a new counter smart contract initialized with the given base value.
+        #[ink(constructor)]
+        pub fn new() -> Self {
+            Default::default()
+        }
+
+        /// Increments the counter value which is stored in the contract's storage.
+        #[ink(message)]
+        pub fn inc(&mut self) {
+            self.count = self.count.checked_add(1).unwrap();
+            ink::env::debug_println!(
+                "The new count is {}, it was modified using the original contract code.",
+                self.count
+            );
+        }
+
+        /// Returns the counter value which is stored in this contract's storage.
+        #[ink(message)]
+        pub fn get(&self) -> u32 {
+            self.count
+        }
+
+        /// Modifies the code which is used to execute calls to this contract address
+        /// (`AccountId`).
+        ///
+        /// We use this to upgrade the contract logic. We don't do any authorization here,
+        /// any caller can execute this method.
+        ///
+        /// In a production contract you would do some authorization here!
+        #[ink(message)]
+        pub fn set_code(&mut self, code_hash: [u8; 32]) {
+            ink::env::set_code_hash(&code_hash).unwrap_or_else(|err| {
+                panic!("Failed to `set_code_hash` to {code_hash:?} due to {err:?}")
+            });
+            ink::env::debug_println!("Switched code hash to {:?}.", code_hash);
+        }
+    }
+
+    #[cfg(all(test, feature = "e2e-tests"))]
+    mod e2e_tests {
+        use super::*;
+        use ink_e2e::ContractsBackend;
+
+        type E2EResult<T> = std::result::Result<T, Box<dyn std::error::Error>>;
+
+        #[ink_e2e::test(additional_contracts = "./updated-incrementer/Cargo.toml")]
+        async fn set_code_works<Client: E2EBackend>(mut client: Client) -> E2EResult<()> {
+            // Given
+            let constructor = IncrementerRef::new();
+            let contract = client
+                .instantiate("incrementer", &ink_e2e::alice(), constructor, 0, None)
+                .await
+                .expect("instantiate failed");
+            let mut call = contract.call::<Incrementer>();
+
+            let get = call.get();
+            let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
+            assert!(matches!(get_res.return_value(), 0));
+
+            let inc = call.inc();
+            let _inc_result = client
+                .call(&ink_e2e::alice(), &inc, 0, None)
+                .await
+                .expect("`inc` failed");
+
+            let get = call.get();
+            let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
+            assert!(matches!(get_res.return_value(), 1));
+
+            // When
+            let new_code_hash = client
+                .upload("updated_incrementer", &ink_e2e::alice(), None)
+                .await
+                .expect("uploading `updated_incrementer` failed")
+                .code_hash;
+
+            let new_code_hash = new_code_hash.as_ref().try_into().unwrap();
+            let set_code = call.set_code(new_code_hash);
+
+            let _set_code_result = client
+                .call(&ink_e2e::alice(), &set_code, 0, None)
+                .await
+                .expect("`set_code` failed");
+
+            // Then
+            // Note that our contract's `AccountId` (so `contract_acc_id`) has stayed the
+            // same between updates!
+            let inc = call.inc();
+
+            let _inc_result = client
+                .call(&ink_e2e::alice(), &inc, 0, None)
+                .await
+                .expect("`inc` failed");
+
+            let get = call.get();
+            let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
+
+            // Remember, we updated our incrementer contract to increment by `4`.
+            assert!(matches!(get_res.return_value(), 5));
+
+            Ok(())
+        }
+    }
+}
diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/Cargo.toml b/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/Cargo.toml
new file mode 100644
index 00000000000..ecc27a0e158
--- /dev/null
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/Cargo.toml
@@ -0,0 +1,19 @@
+[package]
+name = "updated-incrementer"
+version = "4.2.0"
+authors = ["Parity Technologies <admin@parity.io>"]
+edition = "2021"
+publish = false
+
+[dependencies]
+ink = { path = "../../../../crates/ink", default-features = false }
+
+[lib]
+path = "lib.rs"
+
+[features]
+default = ["std"]
+std = [
+    "ink/std",
+]
+ink-as-dependency = []
diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs
new file mode 100644
index 00000000000..b7ab7b7d98b
--- /dev/null
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs
@@ -0,0 +1,51 @@
+#![cfg_attr(not(feature = "std"), no_std, no_main)]
+#![allow(clippy::new_without_default)]
+
+#[ink::contract]
+pub mod incrementer {
+
+    /// Storage struct matches exactly that of the original `incrementer` contract, from which
+    /// we are migrating.
+    #[ink(storage)]
+    pub struct Incrementer {
+        count: u32,
+    }
+
+    impl Incrementer {
+        /// Creates a new counter smart contract initialized with the given base value.
+        ///
+        /// # Note
+        ///
+        /// When upgrading using the `set_code_hash` workflow we only need to point to a
+        /// contract's uploaded code hash, **not** an instantiated contract's
+        /// `AccountId`.
+        ///
+        /// Because of this we will never actually call the constructor of this contract.
+        #[ink(constructor)]
+        pub fn new() -> Self {
+            unreachable!(
+                "Constructors are not called when upgrading using `set_code_hash`."
+            )
+        }
+
+        /// Run the migration to the data layout for the upgraded contract.
+        #[ink(message)]
+        pub fn migrate(&self) {
+        }
+
+        /// Modifies the code which is used to execute calls to this contract address
+        /// (`AccountId`).
+        ///
+        /// We use this to upgrade the contract logic. We don't do any authorization here,
+        /// any caller can execute this method.
+        ///
+        /// In a production contract you would do some authorization here!
+        #[ink(message)]
+        pub fn set_code(&mut self, code_hash: Hash) {
+            self.env().set_code_hash(&code_hash).unwrap_or_else(|err| {
+                panic!("Failed to `set_code_hash` to {code_hash:?} due to {err:?}")
+            });
+            ink::env::debug_println!("Switched code hash to {:?}.", code_hash);
+        }
+    }
+}
diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/updated-incrementer/Cargo.toml b/integration-tests/upgradeable-contracts/set-code-hash-migration/updated-incrementer/Cargo.toml
new file mode 100644
index 00000000000..ecc27a0e158
--- /dev/null
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/updated-incrementer/Cargo.toml
@@ -0,0 +1,19 @@
+[package]
+name = "updated-incrementer"
+version = "4.2.0"
+authors = ["Parity Technologies <admin@parity.io>"]
+edition = "2021"
+publish = false
+
+[dependencies]
+ink = { path = "../../../../crates/ink", default-features = false }
+
+[lib]
+path = "lib.rs"
+
+[features]
+default = ["std"]
+std = [
+    "ink/std",
+]
+ink-as-dependency = []
diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/updated-incrementer/lib.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/updated-incrementer/lib.rs
new file mode 100644
index 00000000000..584001ffa3a
--- /dev/null
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/updated-incrementer/lib.rs
@@ -0,0 +1,75 @@
+#![cfg_attr(not(feature = "std"), no_std, no_main)]
+#![allow(clippy::new_without_default)]
+
+#[ink::contract]
+pub mod incrementer {
+
+    /// Track a counter in storage.
+    ///
+    /// # Note
+    ///
+    /// We have changed the storage layout:
+    /// - `count` is now a `u64` instead of a `u32`.
+    /// - We have added a new field `inc_by` which controls how many to increment by.
+    #[ink(storage)]
+    pub struct Incrementer {
+        count: u64,
+        inc_by: u8,
+    }
+
+    impl Incrementer {
+        /// Creates a new counter smart contract initialized with the given base value.
+        ///
+        /// # Note
+        ///
+        /// When upgrading using the `set_code_hash` workflow we only need to point to a
+        /// contract's uploaded code hash, **not** an instantiated contract's
+        /// `AccountId`.
+        ///
+        /// Because of this we will never actually call the constructor of this contract.
+        #[ink(constructor)]
+        pub fn new() -> Self {
+            unreachable!(
+                "Constructors are not called when upgrading using `set_code_hash`."
+            )
+        }
+
+        /// Increments the counter value which is stored in the contract's storage.
+        ///
+        /// # Note
+        ///
+        /// In this upgraded contract the value is incremented by the value in the
+        /// `inc_by` field.
+        #[ink(message)]
+        pub fn inc(&mut self) {
+            self.count = self.count.checked_add(self.inc_by).unwrap();
+        }
+
+        /// Set the value by which the counter will be incremented.
+        #[ink(message)]
+        pub fn set_inc_by(&mut self, inc_by: u8) {
+            self.inc_by = inc_by;
+        }
+
+        /// Returns the counter value which is stored in this contract's storage.
+        #[ink(message)]
+        pub fn get(&self) -> u32 {
+            self.count
+        }
+
+        /// Modifies the code which is used to execute calls to this contract address
+        /// (`AccountId`).
+        ///
+        /// We use this to upgrade the contract logic. We don't do any authorization here,
+        /// any caller can execute this method.
+        ///
+        /// In a production contract you would do some authorization here!
+        #[ink(message)]
+        pub fn set_code(&mut self, code_hash: Hash) {
+            self.env().set_code_hash(&code_hash).unwrap_or_else(|err| {
+                panic!("Failed to `set_code_hash` to {code_hash:?} due to {err:?}")
+            });
+            ink::env::debug_println!("Switched code hash to {:?}.", code_hash);
+        }
+    }
+}

From 947fa42bdbed36124fe182e9992f851d6b429ae0 Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Tue, 12 Sep 2023 13:03:00 +0100
Subject: [PATCH 02/14] Update versions

---
 .../upgradeable-contracts/set-code-hash-migration/Cargo.toml    | 2 +-
 .../set-code-hash-migration/migration/Cargo.toml                | 2 +-
 .../set-code-hash-migration/updated-incrementer/Cargo.toml      | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/Cargo.toml b/integration-tests/upgradeable-contracts/set-code-hash-migration/Cargo.toml
index 79c4aad15f0..6a9d79de8f2 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/Cargo.toml
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "incrementer"
-version = "4.2.0"
+version = "5.0.0-alpha"
 authors = ["Parity Technologies <admin@parity.io>"]
 edition = "2021"
 publish = false
diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/Cargo.toml b/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/Cargo.toml
index ecc27a0e158..49b24990c9d 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/Cargo.toml
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "updated-incrementer"
-version = "4.2.0"
+version = "5.0.0-alpha"
 authors = ["Parity Technologies <admin@parity.io>"]
 edition = "2021"
 publish = false
diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/updated-incrementer/Cargo.toml b/integration-tests/upgradeable-contracts/set-code-hash-migration/updated-incrementer/Cargo.toml
index ecc27a0e158..49b24990c9d 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/updated-incrementer/Cargo.toml
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/updated-incrementer/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "updated-incrementer"
-version = "4.2.0"
+version = "5.0.0-alpha"
 authors = ["Parity Technologies <admin@parity.io>"]
 edition = "2021"
 publish = false

From ef830fa8efbc1673e8580d34b522528137e24e9b Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Tue, 12 Sep 2023 14:05:25 +0100
Subject: [PATCH 03/14] WIP

---
 .../set-code-hash-migration/migration/lib.rs  | 35 +++++++++++--------
 1 file changed, 21 insertions(+), 14 deletions(-)

diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs
index b7ab7b7d98b..c5d086ec277 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs
@@ -11,6 +11,12 @@ pub mod incrementer {
         count: u32,
     }
 
+    #[ink::storage_item]
+    pub struct IncrementerNew {
+        count: u64,
+        inc_by: u8,
+    }
+
     impl Incrementer {
         /// Creates a new counter smart contract initialized with the given base value.
         ///
@@ -29,23 +35,24 @@ pub mod incrementer {
         }
 
         /// Run the migration to the data layout for the upgraded contract.
-        #[ink(message)]
-        pub fn migrate(&self) {
-        }
-
-        /// Modifies the code which is used to execute calls to this contract address
-        /// (`AccountId`).
-        ///
-        /// We use this to upgrade the contract logic. We don't do any authorization here,
-        /// any caller can execute this method.
+        /// Once the storage migration has successfully completed, the contract will be upgraded
+        /// to the supplied code hash.
         ///
         /// In a production contract you would do some authorization here!
+        ///
+        /// # Note
+        ///
+        /// This function necessarily accepts a `&self` instead of a `&mut self` because we are
+        /// modifying storage directly for the migration. Using `&mut self` would overwrite our
+        /// migration changes with the contents of the original `Incrementer`.
         #[ink(message)]
-        pub fn set_code(&mut self, code_hash: Hash) {
-            self.env().set_code_hash(&code_hash).unwrap_or_else(|err| {
-                panic!("Failed to `set_code_hash` to {code_hash:?} due to {err:?}")
-            });
-            ink::env::debug_println!("Switched code hash to {:?}.", code_hash);
+        pub fn migrate(&self, inc_by: u8, code_hash: Hash) {
+            let incrementer_new = IncrementerNew {
+                count: self.count as u64,
+                inc_by,
+            };
+            // ink::env::set_contract_storage::<IncrementerNew>(ink::key_hash::new(), &incrementer_new);
+            ink::env::set_code_hash2(&code_hash);
         }
     }
 }

From cedf1205a57cd7dd377717a82d2d876ba142bc51 Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Tue, 12 Sep 2023 14:50:57 +0100
Subject: [PATCH 04/14] WIP migration

---
 .../set-code-hash-migration/Cargo.toml        |  5 +++
 .../set-code-hash-migration/lib.rs            | 31 ++++++++++++++++---
 .../set-code-hash-migration/migration/lib.rs  |  4 +--
 3 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/Cargo.toml b/integration-tests/upgradeable-contracts/set-code-hash-migration/Cargo.toml
index 6a9d79de8f2..31a6b2cb2b6 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/Cargo.toml
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/Cargo.toml
@@ -8,6 +8,9 @@ publish = false
 [dependencies]
 ink = { path = "../../../crates/ink", default-features = false }
 
+migration = { path = "./migration", default-features = false, features = ["ink-as-dependency"] }
+updated-incrementer = { path = "./updated-incrementer", default-features = false, features = ["ink-as-dependency"] }
+
 [dev-dependencies]
 ink_e2e = { path = "../../../crates/e2e" }
 
@@ -18,6 +21,8 @@ path = "lib.rs"
 default = ["std"]
 std = [
     "ink/std",
+    "migration/std",
+    "updated-incrementer/std",
 ]
 ink-as-dependency = []
 e2e-tests = []
diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
index 86851852d56..ce94f4c5209 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
@@ -72,8 +72,8 @@ pub mod incrementer {
 
         type E2EResult<T> = std::result::Result<T, Box<dyn std::error::Error>>;
 
-        #[ink_e2e::test(additional_contracts = "./updated-incrementer/Cargo.toml")]
-        async fn set_code_works<Client: E2EBackend>(mut client: Client) -> E2EResult<()> {
+        #[ink_e2e::test]
+        async fn migration_works<Client: E2EBackend>(mut client: Client) -> E2EResult<()> {
             // Given
             let constructor = IncrementerRef::new();
             let contract = client
@@ -96,14 +96,35 @@ pub mod incrementer {
             let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
             assert!(matches!(get_res.return_value(), 1));
 
-            // When
+            // Upload the code for the contract to be updated to after the migration.
             let new_code_hash = client
-                .upload("updated_incrementer", &ink_e2e::alice(), None)
+                .upload("updated-incrementer", &ink_e2e::alice(), None)
                 .await
-                .expect("uploading `updated_incrementer` failed")
+                .expect("uploading `updated-incrementer` failed")
                 .code_hash;
 
             let new_code_hash = new_code_hash.as_ref().try_into().unwrap();
+
+            // When
+            let migration_contract = client
+                .upload("migration", &ink_e2e::alice(), None)
+                .await
+                .expect("uploading `migration` failed");
+            let migration_code_hash = migration_contract.code_hash.as_ref().try_into().unwrap();
+            let migrate = migration_contract
+                .call::<migration::incrementer::Incrementer>()
+                .migrate();
+
+            let _migration_result = client
+                .call(
+                    &ink_e2e::alice(),
+                    &call.set_code(migration_code_hash),
+                    0,
+                    None,
+                )
+                .await
+                .expect("`set_code` failed");
+
             let set_code = call.set_code(new_code_hash);
 
             let _set_code_result = client
diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs
index c5d086ec277..a707a10dba3 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs
@@ -51,8 +51,8 @@ pub mod incrementer {
                 count: self.count as u64,
                 inc_by,
             };
-            // ink::env::set_contract_storage::<IncrementerNew>(ink::key_hash::new(), &incrementer_new);
-            ink::env::set_code_hash2(&code_hash);
+            ink::env::set_contract_storage(0x00000000, &incrementer_new);
+            ink::env::set_code_hash2::<<Self as ink::ContractEnv>::Env>(&code_hash);
         }
     }
 }

From b8e621a38d5a4e0195f4f8c5f414b01dec74f288 Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Wed, 13 Sep 2023 11:42:45 +0100
Subject: [PATCH 05/14] WIP

---
 .../set-code-hash-migration/lib.rs            | 38 +++++++++----------
 1 file changed, 19 insertions(+), 19 deletions(-)

diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
index ce94f4c5209..7f675b8dbb6 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
@@ -94,7 +94,8 @@ pub mod incrementer {
 
             let get = call.get();
             let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
-            assert!(matches!(get_res.return_value(), 1));
+            let pre_migration_value = get_res.return_value();
+            assert!(matches!(pre_migration_value, 1));
 
             // Upload the code for the contract to be updated to after the migration.
             let new_code_hash = client
@@ -102,40 +103,39 @@ pub mod incrementer {
                 .await
                 .expect("uploading `updated-incrementer` failed")
                 .code_hash;
-
             let new_code_hash = new_code_hash.as_ref().try_into().unwrap();
 
-            // When
+            // Upload the code for the migration contract.
             let migration_contract = client
                 .upload("migration", &ink_e2e::alice(), None)
                 .await
                 .expect("uploading `migration` failed");
             let migration_code_hash = migration_contract.code_hash.as_ref().try_into().unwrap();
-            let migrate = migration_contract
-                .call::<migration::incrementer::Incrementer>()
-                .migrate();
-
-            let _migration_result = client
-                .call(
-                    &ink_e2e::alice(),
-                    &call.set_code(migration_code_hash),
-                    0,
-                    None,
-                )
-                .await
-                .expect("`set_code` failed");
 
-            let set_code = call.set_code(new_code_hash);
+            // When
 
+            // Set the code hash to the migration contract
+            let set_code = call.set_code(migration_code_hash);
             let _set_code_result = client
                 .call(&ink_e2e::alice(), &set_code, 0, None)
                 .await
                 .expect("`set_code` failed");
 
+            // Call the migration contract with a new value for `inc_by` and the code hash of the
+            // updated contract.
+            const NEW_INC_BY: u8 = 4;
+            let migrate = contract
+                .call::<migration::incrementer::Incrementer>()
+                .migrate(NEW_INC_BY, new_code_hash);
+
+            let _migration_result = client.call(&ink_e2e::alice(), &migrate, 0, None)
+                .await
+                .expect("`migrate` failed");
+
             // Then
             // Note that our contract's `AccountId` (so `contract_acc_id`) has stayed the
             // same between updates!
-            let inc = call.inc();
+            let inc = contract.call::<updated_incrementer::incrementer::Incrementer>().inc();
 
             let _inc_result = client
                 .call(&ink_e2e::alice(), &inc, 0, None)
@@ -146,7 +146,7 @@ pub mod incrementer {
             let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
 
             // Remember, we updated our incrementer contract to increment by `4`.
-            assert!(matches!(get_res.return_value(), 5));
+            assert!(matches!(get_res.return_value(), pre_migration_value + NEW_INC_BY.into()));
 
             Ok(())
         }

From d187a4511fcb455bc36c9c3f334eb206def20239 Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Wed, 13 Sep 2023 12:03:13 +0100
Subject: [PATCH 06/14] Make test pass

---
 .../set-code-hash-migration/lib.rs            | 31 ++++++++++++-------
 .../migration/Cargo.toml                      |  2 +-
 .../set-code-hash-migration/migration/lib.rs  | 24 ++++++++------
 .../updated-incrementer/lib.rs                |  4 +--
 4 files changed, 38 insertions(+), 23 deletions(-)

diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
index 7f675b8dbb6..491160811ae 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
@@ -57,8 +57,8 @@ pub mod incrementer {
         ///
         /// In a production contract you would do some authorization here!
         #[ink(message)]
-        pub fn set_code(&mut self, code_hash: [u8; 32]) {
-            ink::env::set_code_hash(&code_hash).unwrap_or_else(|err| {
+        pub fn set_code(&mut self, code_hash: Hash) {
+            self.env().set_code_hash(&code_hash).unwrap_or_else(|err| {
                 panic!("Failed to `set_code_hash` to {code_hash:?} due to {err:?}")
             });
             ink::env::debug_println!("Switched code hash to {:?}.", code_hash);
@@ -73,7 +73,9 @@ pub mod incrementer {
         type E2EResult<T> = std::result::Result<T, Box<dyn std::error::Error>>;
 
         #[ink_e2e::test]
-        async fn migration_works<Client: E2EBackend>(mut client: Client) -> E2EResult<()> {
+        async fn migration_works<Client: E2EBackend>(
+            mut client: Client,
+        ) -> E2EResult<()> {
             // Given
             let constructor = IncrementerRef::new();
             let contract = client
@@ -84,7 +86,7 @@ pub mod incrementer {
 
             let get = call.get();
             let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
-            assert!(matches!(get_res.return_value(), 0));
+            assert_eq!(get_res.return_value(), 0);
 
             let inc = call.inc();
             let _inc_result = client
@@ -95,7 +97,7 @@ pub mod incrementer {
             let get = call.get();
             let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
             let pre_migration_value = get_res.return_value();
-            assert!(matches!(pre_migration_value, 1));
+            assert_eq!(pre_migration_value, 1);
 
             // Upload the code for the contract to be updated to after the migration.
             let new_code_hash = client
@@ -110,7 +112,8 @@ pub mod incrementer {
                 .upload("migration", &ink_e2e::alice(), None)
                 .await
                 .expect("uploading `migration` failed");
-            let migration_code_hash = migration_contract.code_hash.as_ref().try_into().unwrap();
+            let migration_code_hash =
+                migration_contract.code_hash.as_ref().try_into().unwrap();
 
             // When
 
@@ -121,21 +124,24 @@ pub mod incrementer {
                 .await
                 .expect("`set_code` failed");
 
-            // Call the migration contract with a new value for `inc_by` and the code hash of the
-            // updated contract.
+            // Call the migration contract with a new value for `inc_by` and the code hash
+            // of the updated contract.
             const NEW_INC_BY: u8 = 4;
             let migrate = contract
                 .call::<migration::incrementer::Incrementer>()
                 .migrate(NEW_INC_BY, new_code_hash);
 
-            let _migration_result = client.call(&ink_e2e::alice(), &migrate, 0, None)
+            let _migration_result = client
+                .call(&ink_e2e::alice(), &migrate, 0, None)
                 .await
                 .expect("`migrate` failed");
 
             // Then
             // Note that our contract's `AccountId` (so `contract_acc_id`) has stayed the
             // same between updates!
-            let inc = contract.call::<updated_incrementer::incrementer::Incrementer>().inc();
+            let inc = contract
+                .call::<updated_incrementer::incrementer::Incrementer>()
+                .inc();
 
             let _inc_result = client
                 .call(&ink_e2e::alice(), &inc, 0, None)
@@ -146,7 +152,10 @@ pub mod incrementer {
             let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
 
             // Remember, we updated our incrementer contract to increment by `4`.
-            assert!(matches!(get_res.return_value(), pre_migration_value + NEW_INC_BY.into()));
+            assert_eq!(
+                get_res.return_value(),
+                pre_migration_value + NEW_INC_BY as u32
+            );
 
             Ok(())
         }
diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/Cargo.toml b/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/Cargo.toml
index 49b24990c9d..af5cae8ab2a 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/Cargo.toml
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/Cargo.toml
@@ -1,5 +1,5 @@
 [package]
-name = "updated-incrementer"
+name = "migration"
 version = "5.0.0-alpha"
 authors = ["Parity Technologies <admin@parity.io>"]
 edition = "2021"
diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs
index a707a10dba3..5a56202d401 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs
@@ -4,8 +4,8 @@
 #[ink::contract]
 pub mod incrementer {
 
-    /// Storage struct matches exactly that of the original `incrementer` contract, from which
-    /// we are migrating.
+    /// Storage struct matches exactly that of the original `incrementer` contract, from
+    /// which we are migrating.
     #[ink(storage)]
     pub struct Incrementer {
         count: u32,
@@ -35,24 +35,30 @@ pub mod incrementer {
         }
 
         /// Run the migration to the data layout for the upgraded contract.
-        /// Once the storage migration has successfully completed, the contract will be upgraded
-        /// to the supplied code hash.
+        /// Once the storage migration has successfully completed, the contract will be
+        /// upgraded to the supplied code hash.
         ///
         /// In a production contract you would do some authorization here!
         ///
         /// # Note
         ///
-        /// This function necessarily accepts a `&self` instead of a `&mut self` because we are
-        /// modifying storage directly for the migration. Using `&mut self` would overwrite our
-        /// migration changes with the contents of the original `Incrementer`.
+        /// This function necessarily accepts a `&self` instead of a `&mut self` because
+        /// we are modifying storage directly for the migration. Using `&mut self`
+        /// would overwrite our migration changes with the contents of the
+        /// original `Incrementer`.
         #[ink(message)]
         pub fn migrate(&self, inc_by: u8, code_hash: Hash) {
             let incrementer_new = IncrementerNew {
                 count: self.count as u64,
                 inc_by,
             };
-            ink::env::set_contract_storage(0x00000000, &incrementer_new);
-            ink::env::set_code_hash2::<<Self as ink::ContractEnv>::Env>(&code_hash);
+            const STORAGE_KEY: u32 = 0x00000000;
+            ink::env::set_contract_storage(&STORAGE_KEY, &incrementer_new);
+
+            ink::env::set_code_hash::<<Self as ink::env::ContractEnv>::Env>(&code_hash)
+                .unwrap_or_else(|err| {
+                    panic!("Failed to `set_code_hash` to {code_hash:?} due to {err:?}")
+                })
         }
     }
 }
diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/updated-incrementer/lib.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/updated-incrementer/lib.rs
index 584001ffa3a..7ac05c80e3d 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/updated-incrementer/lib.rs
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/updated-incrementer/lib.rs
@@ -42,7 +42,7 @@ pub mod incrementer {
         /// `inc_by` field.
         #[ink(message)]
         pub fn inc(&mut self) {
-            self.count = self.count.checked_add(self.inc_by).unwrap();
+            self.count = self.count.checked_add(self.inc_by.into()).unwrap();
         }
 
         /// Set the value by which the counter will be incremented.
@@ -53,7 +53,7 @@ pub mod incrementer {
 
         /// Returns the counter value which is stored in this contract's storage.
         #[ink(message)]
-        pub fn get(&self) -> u32 {
+        pub fn get(&self) -> u64 {
             self.count
         }
 

From 35644b6a67a8146c01d2bc489bc978af744c82e0 Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Wed, 13 Sep 2023 12:17:31 +0100
Subject: [PATCH 07/14] Move e2e tests mod to own file

---
 .../set-code-hash-migration/e2e_tests.rs      |  89 ++++++++++++++++
 .../set-code-hash-migration/lib.rs            | 100 +-----------------
 2 files changed, 92 insertions(+), 97 deletions(-)
 create mode 100644 integration-tests/upgradeable-contracts/set-code-hash-migration/e2e_tests.rs

diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/e2e_tests.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/e2e_tests.rs
new file mode 100644
index 00000000000..ad8b048fd65
--- /dev/null
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/e2e_tests.rs
@@ -0,0 +1,89 @@
+use super::incrementer::*;
+use ink_e2e::ContractsBackend;
+
+type E2EResult<T> = std::result::Result<T, Box<dyn std::error::Error>>;
+
+#[ink_e2e::test]
+async fn migration_works<Client: E2EBackend>(mut client: Client) -> E2EResult<()> {
+    // Given
+    let constructor = IncrementerRef::new();
+    let contract = client
+        .instantiate("incrementer", &ink_e2e::alice(), constructor, 0, None)
+        .await
+        .expect("instantiate failed");
+    let mut call = contract.call::<Incrementer>();
+
+    let get = call.get();
+    let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
+    assert_eq!(get_res.return_value(), 0);
+
+    let inc = call.inc();
+    let _inc_result = client
+        .call(&ink_e2e::alice(), &inc, 0, None)
+        .await
+        .expect("`inc` failed");
+
+    let get = call.get();
+    let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
+    let pre_migration_value = get_res.return_value();
+    assert_eq!(pre_migration_value, 1);
+
+    // Upload the code for the contract to be updated to after the migration.
+    let new_code_hash = client
+        .upload("updated-incrementer", &ink_e2e::alice(), None)
+        .await
+        .expect("uploading `updated-incrementer` failed")
+        .code_hash;
+    let new_code_hash = new_code_hash.as_ref().try_into().unwrap();
+
+    // Upload the code for the migration contract.
+    let migration_contract = client
+        .upload("migration", &ink_e2e::alice(), None)
+        .await
+        .expect("uploading `migration` failed");
+    let migration_code_hash = migration_contract.code_hash.as_ref().try_into().unwrap();
+
+    // When
+
+    // Set the code hash to the migration contract
+    let set_code = call.set_code(migration_code_hash);
+    let _set_code_result = client
+        .call(&ink_e2e::alice(), &set_code, 0, None)
+        .await
+        .expect("`set_code` failed");
+
+    // Call the migration contract with a new value for `inc_by` and the code hash
+    // of the updated contract.
+    const NEW_INC_BY: u8 = 4;
+    let migrate = contract
+        .call::<migration::incrementer::Incrementer>()
+        .migrate(NEW_INC_BY, new_code_hash);
+
+    let _migration_result = client
+        .call(&ink_e2e::alice(), &migrate, 0, None)
+        .await
+        .expect("`migrate` failed");
+
+    // Then
+    // Note that our contract's `AccountId` (so `contract_acc_id`) has stayed the
+    // same between updates!
+    let inc = contract
+        .call::<updated_incrementer::incrementer::Incrementer>()
+        .inc();
+
+    let _inc_result = client
+        .call(&ink_e2e::alice(), &inc, 0, None)
+        .await
+        .expect("`inc` failed");
+
+    let get = call.get();
+    let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
+
+    // Remember, we updated our incrementer contract to increment by `4`.
+    assert_eq!(
+        get_res.return_value(),
+        pre_migration_value + NEW_INC_BY as u32
+    );
+
+    Ok(())
+}
diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
index 491160811ae..4425424bacb 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
@@ -10,7 +10,6 @@
 
 #[ink::contract]
 pub mod incrementer {
-
     /// Track a counter in storage.
     ///
     /// # Note
@@ -64,100 +63,7 @@ pub mod incrementer {
             ink::env::debug_println!("Switched code hash to {:?}.", code_hash);
         }
     }
-
-    #[cfg(all(test, feature = "e2e-tests"))]
-    mod e2e_tests {
-        use super::*;
-        use ink_e2e::ContractsBackend;
-
-        type E2EResult<T> = std::result::Result<T, Box<dyn std::error::Error>>;
-
-        #[ink_e2e::test]
-        async fn migration_works<Client: E2EBackend>(
-            mut client: Client,
-        ) -> E2EResult<()> {
-            // Given
-            let constructor = IncrementerRef::new();
-            let contract = client
-                .instantiate("incrementer", &ink_e2e::alice(), constructor, 0, None)
-                .await
-                .expect("instantiate failed");
-            let mut call = contract.call::<Incrementer>();
-
-            let get = call.get();
-            let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
-            assert_eq!(get_res.return_value(), 0);
-
-            let inc = call.inc();
-            let _inc_result = client
-                .call(&ink_e2e::alice(), &inc, 0, None)
-                .await
-                .expect("`inc` failed");
-
-            let get = call.get();
-            let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
-            let pre_migration_value = get_res.return_value();
-            assert_eq!(pre_migration_value, 1);
-
-            // Upload the code for the contract to be updated to after the migration.
-            let new_code_hash = client
-                .upload("updated-incrementer", &ink_e2e::alice(), None)
-                .await
-                .expect("uploading `updated-incrementer` failed")
-                .code_hash;
-            let new_code_hash = new_code_hash.as_ref().try_into().unwrap();
-
-            // Upload the code for the migration contract.
-            let migration_contract = client
-                .upload("migration", &ink_e2e::alice(), None)
-                .await
-                .expect("uploading `migration` failed");
-            let migration_code_hash =
-                migration_contract.code_hash.as_ref().try_into().unwrap();
-
-            // When
-
-            // Set the code hash to the migration contract
-            let set_code = call.set_code(migration_code_hash);
-            let _set_code_result = client
-                .call(&ink_e2e::alice(), &set_code, 0, None)
-                .await
-                .expect("`set_code` failed");
-
-            // Call the migration contract with a new value for `inc_by` and the code hash
-            // of the updated contract.
-            const NEW_INC_BY: u8 = 4;
-            let migrate = contract
-                .call::<migration::incrementer::Incrementer>()
-                .migrate(NEW_INC_BY, new_code_hash);
-
-            let _migration_result = client
-                .call(&ink_e2e::alice(), &migrate, 0, None)
-                .await
-                .expect("`migrate` failed");
-
-            // Then
-            // Note that our contract's `AccountId` (so `contract_acc_id`) has stayed the
-            // same between updates!
-            let inc = contract
-                .call::<updated_incrementer::incrementer::Incrementer>()
-                .inc();
-
-            let _inc_result = client
-                .call(&ink_e2e::alice(), &inc, 0, None)
-                .await
-                .expect("`inc` failed");
-
-            let get = call.get();
-            let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
-
-            // Remember, we updated our incrementer contract to increment by `4`.
-            assert_eq!(
-                get_res.return_value(),
-                pre_migration_value + NEW_INC_BY as u32
-            );
-
-            Ok(())
-        }
-    }
 }
+
+#[cfg(all(test, feature = "e2e-tests"))]
+mod e2e_tests;

From 51fe530fb13bc910e89495503e4eba26189f6c45 Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Wed, 13 Sep 2023 12:18:59 +0100
Subject: [PATCH 08/14] Update comment

---
 .../upgradeable-contracts/set-code-hash-migration/e2e_tests.rs  | 2 --
 1 file changed, 2 deletions(-)

diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/e2e_tests.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/e2e_tests.rs
index ad8b048fd65..445201f688b 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/e2e_tests.rs
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/e2e_tests.rs
@@ -65,8 +65,6 @@ async fn migration_works<Client: E2EBackend>(mut client: Client) -> E2EResult<()
         .expect("`migrate` failed");
 
     // Then
-    // Note that our contract's `AccountId` (so `contract_acc_id`) has stayed the
-    // same between updates!
     let inc = contract
         .call::<updated_incrementer::incrementer::Incrementer>()
         .inc();

From c5e38b6683b772d9ba796fb323f7540b97256b71 Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Thu, 26 Oct 2023 09:50:38 +0100
Subject: [PATCH 09/14] Update example for new e2e API

---
 .../set-code-hash-migration/e2e_tests.rs      | 29 ++++++++++++-------
 1 file changed, 18 insertions(+), 11 deletions(-)

diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/e2e_tests.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/e2e_tests.rs
index 445201f688b..e80f7329c1d 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/e2e_tests.rs
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/e2e_tests.rs
@@ -6,31 +6,34 @@ type E2EResult<T> = std::result::Result<T, Box<dyn std::error::Error>>;
 #[ink_e2e::test]
 async fn migration_works<Client: E2EBackend>(mut client: Client) -> E2EResult<()> {
     // Given
-    let constructor = IncrementerRef::new();
+    let mut constructor = IncrementerRef::new();
     let contract = client
-        .instantiate("incrementer", &ink_e2e::alice(), constructor, 0, None)
+        .instantiate("incrementer", &ink_e2e::alice(), &mut constructor)
+        .submit()
         .await
         .expect("instantiate failed");
     let mut call = contract.call::<Incrementer>();
 
     let get = call.get();
-    let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
+    let get_res = client.call(&ink_e2e::alice(), &get).dry_run().await;
     assert_eq!(get_res.return_value(), 0);
 
     let inc = call.inc();
     let _inc_result = client
-        .call(&ink_e2e::alice(), &inc, 0, None)
+        .call(&ink_e2e::alice(), &inc)
+        .submit()
         .await
         .expect("`inc` failed");
 
     let get = call.get();
-    let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
+    let get_res = client.call(&ink_e2e::alice(), &get).dry_run().await;
     let pre_migration_value = get_res.return_value();
     assert_eq!(pre_migration_value, 1);
 
     // Upload the code for the contract to be updated to after the migration.
     let new_code_hash = client
-        .upload("updated-incrementer", &ink_e2e::alice(), None)
+        .upload("updated-incrementer", &ink_e2e::alice())
+        .submit()
         .await
         .expect("uploading `updated-incrementer` failed")
         .code_hash;
@@ -38,7 +41,8 @@ async fn migration_works<Client: E2EBackend>(mut client: Client) -> E2EResult<()
 
     // Upload the code for the migration contract.
     let migration_contract = client
-        .upload("migration", &ink_e2e::alice(), None)
+        .upload("migration", &ink_e2e::alice())
+        .submit()
         .await
         .expect("uploading `migration` failed");
     let migration_code_hash = migration_contract.code_hash.as_ref().try_into().unwrap();
@@ -48,7 +52,8 @@ async fn migration_works<Client: E2EBackend>(mut client: Client) -> E2EResult<()
     // Set the code hash to the migration contract
     let set_code = call.set_code(migration_code_hash);
     let _set_code_result = client
-        .call(&ink_e2e::alice(), &set_code, 0, None)
+        .call(&ink_e2e::alice(), &set_code)
+        .submit()
         .await
         .expect("`set_code` failed");
 
@@ -60,7 +65,8 @@ async fn migration_works<Client: E2EBackend>(mut client: Client) -> E2EResult<()
         .migrate(NEW_INC_BY, new_code_hash);
 
     let _migration_result = client
-        .call(&ink_e2e::alice(), &migrate, 0, None)
+        .call(&ink_e2e::alice(), &migrate)
+        .submit()
         .await
         .expect("`migrate` failed");
 
@@ -70,12 +76,13 @@ async fn migration_works<Client: E2EBackend>(mut client: Client) -> E2EResult<()
         .inc();
 
     let _inc_result = client
-        .call(&ink_e2e::alice(), &inc, 0, None)
+        .call(&ink_e2e::alice(), &inc)
+        .submit()
         .await
         .expect("`inc` failed");
 
     let get = call.get();
-    let get_res = client.call_dry_run(&ink_e2e::alice(), &get, 0, None).await;
+    let get_res = client.call(&ink_e2e::alice(), &get).dry_run().await;
 
     // Remember, we updated our incrementer contract to increment by `4`.
     assert_eq!(

From 2e8b86b23e830e874c24c60cacd0ac58d51d103d Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Tue, 30 Jan 2024 14:16:54 +0000
Subject: [PATCH 10/14] Update
 integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Michael Müller <michi@parity.io>
---
 .../upgradeable-contracts/set-code-hash-migration/lib.rs        | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
index 4425424bacb..ea67cf199aa 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/lib.rs
@@ -3,7 +3,7 @@
 //! Demonstrates how to use [`set_code_hash`](https://docs.rs/ink_env/latest/ink_env/fn.set_code_hash.html)
 //! to swap out the `code_hash` of an on-chain contract.
 //!
-//! We will swap the code of our `Incrementer` contract with that of the an `Incrementer`
+//! We will swap the code of our `Incrementer` contract with that of the `Incrementer`
 //! found in the `updated_incrementer` folder.
 //!
 //! See the included End-to-End tests an example update workflow.

From 231c5d5b7216b36aae8cec470b8e2cc0ce178de2 Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Wed, 31 Jan 2024 10:41:28 +0000
Subject: [PATCH 11/14] Top level gitignore

---
 integration-tests/upgradeable-contracts/.gitignore | 2 ++
 1 file changed, 2 insertions(+)
 create mode 100644 integration-tests/upgradeable-contracts/.gitignore

diff --git a/integration-tests/upgradeable-contracts/.gitignore b/integration-tests/upgradeable-contracts/.gitignore
new file mode 100644
index 00000000000..ff75a1d9fa2
--- /dev/null
+++ b/integration-tests/upgradeable-contracts/.gitignore
@@ -0,0 +1,2 @@
+**/target/
+Cargo.lock
\ No newline at end of file

From 76c4fc14a5981b294630bde88da9d5612490e1bb Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Wed, 31 Jan 2024 10:59:57 +0000
Subject: [PATCH 12/14] Fix tests update comments

---
 .../set-code-hash-migration/e2e_tests.rs      | 22 +++++++++----------
 .../set-code-hash-migration/migration/lib.rs  | 14 ++++++++----
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/e2e_tests.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/e2e_tests.rs
index e80f7329c1d..dbbd029a8b0 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/e2e_tests.rs
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/e2e_tests.rs
@@ -12,21 +12,21 @@ async fn migration_works<Client: E2EBackend>(mut client: Client) -> E2EResult<()
         .submit()
         .await
         .expect("instantiate failed");
-    let mut call = contract.call::<Incrementer>();
+    let mut call_builder = contract.call_builder::<Incrementer>();
 
-    let get = call.get();
-    let get_res = client.call(&ink_e2e::alice(), &get).dry_run().await;
+    let get = call_builder.get();
+    let get_res = client.call(&ink_e2e::alice(), &get).dry_run().await?;
     assert_eq!(get_res.return_value(), 0);
 
-    let inc = call.inc();
+    let inc = call_builder.inc();
     let _inc_result = client
         .call(&ink_e2e::alice(), &inc)
         .submit()
         .await
         .expect("`inc` failed");
 
-    let get = call.get();
-    let get_res = client.call(&ink_e2e::alice(), &get).dry_run().await;
+    let get = call_builder.get();
+    let get_res = client.call(&ink_e2e::alice(), &get).dry_run().await?;
     let pre_migration_value = get_res.return_value();
     assert_eq!(pre_migration_value, 1);
 
@@ -50,7 +50,7 @@ async fn migration_works<Client: E2EBackend>(mut client: Client) -> E2EResult<()
     // When
 
     // Set the code hash to the migration contract
-    let set_code = call.set_code(migration_code_hash);
+    let set_code = call_builder.set_code(migration_code_hash);
     let _set_code_result = client
         .call(&ink_e2e::alice(), &set_code)
         .submit()
@@ -61,7 +61,7 @@ async fn migration_works<Client: E2EBackend>(mut client: Client) -> E2EResult<()
     // of the updated contract.
     const NEW_INC_BY: u8 = 4;
     let migrate = contract
-        .call::<migration::incrementer::Incrementer>()
+        .call_builder::<migration::incrementer::Incrementer>()
         .migrate(NEW_INC_BY, new_code_hash);
 
     let _migration_result = client
@@ -72,7 +72,7 @@ async fn migration_works<Client: E2EBackend>(mut client: Client) -> E2EResult<()
 
     // Then
     let inc = contract
-        .call::<updated_incrementer::incrementer::Incrementer>()
+        .call_builder::<updated_incrementer::incrementer::Incrementer>()
         .inc();
 
     let _inc_result = client
@@ -81,8 +81,8 @@ async fn migration_works<Client: E2EBackend>(mut client: Client) -> E2EResult<()
         .await
         .expect("`inc` failed");
 
-    let get = call.get();
-    let get_res = client.call(&ink_e2e::alice(), &get).dry_run().await;
+    let get = call_builder.get();
+    let get_res = client.call(&ink_e2e::alice(), &get).dry_run().await?;
 
     // Remember, we updated our incrementer contract to increment by `4`.
     assert_eq!(
diff --git a/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs b/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs
index 5a56202d401..93d13cc6f19 100644
--- a/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs
+++ b/integration-tests/upgradeable-contracts/set-code-hash-migration/migration/lib.rs
@@ -43,16 +43,22 @@ pub mod incrementer {
         /// # Note
         ///
         /// This function necessarily accepts a `&self` instead of a `&mut self` because
-        /// we are modifying storage directly for the migration. Using `&mut self`
-        /// would overwrite our migration changes with the contents of the
-        /// original `Incrementer`.
+        /// we are modifying storage directly for the migration.
+        ///
+        /// The `self` in `&mut self` is the original `Incrementer` storage struct, and
+        /// would be implicitly written to storage following the function execution,
+        /// overwriting the migrated storage.
         #[ink(message)]
         pub fn migrate(&self, inc_by: u8, code_hash: Hash) {
             let incrementer_new = IncrementerNew {
                 count: self.count as u64,
                 inc_by,
             };
-            const STORAGE_KEY: u32 = 0x00000000;
+
+            // overwrite the original storage struct with the migrated storage struct,
+            // which has a layout compatible with the new contract code.
+            const STORAGE_KEY: u32 =
+                <Incrementer as ink::storage::traits::StorageKey>::KEY;
             ink::env::set_contract_storage(&STORAGE_KEY, &incrementer_new);
 
             ink::env::set_code_hash::<<Self as ink::env::ContractEnv>::Env>(&code_hash)

From 6171c97f4f139e3f69b6d0b81224f402d9dd0a66 Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Wed, 31 Jan 2024 11:19:15 +0000
Subject: [PATCH 13/14] Update upgradeable contracts README.md

---
 integration-tests/upgradeable-contracts/README.md | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/integration-tests/upgradeable-contracts/README.md b/integration-tests/upgradeable-contracts/README.md
index 85bd7b5c81f..76e5496a0cb 100644
--- a/integration-tests/upgradeable-contracts/README.md
+++ b/integration-tests/upgradeable-contracts/README.md
@@ -12,6 +12,21 @@ This is exactly what `set_code_hash()` function does.
 However, developers needs to be mindful of storage compatibility.
 You can read more about storage compatibility on [use.ink](https://use.ink/basics/upgradeable-contracts#replacing-contract-code-with-set_code_hash)
 
+## [`set-code-hash`](set-code-hash-migration/)
+
+When upgrading a contract, the new code may have a different storage layout. This example illustrates a method to 
+migrate the storage from the old layout to the new layout. It does so by using an intermediate `migration` contract
+which performs the storage upgrade. Thw workflow is as follows:
+
+
+1. Upload a `migration` contract with a message `migrate` which performs the storage migration.
+2. Set code hash to the `migration` contract.
+3. Upload the upgraded version of the original contract.
+4. Call `migrate` on the `migration` contract, passing the code hash of the new updated incrementer contract from `3.` 
+This must happen as a single message, because following the storage migration, the contract will not be able to be 
+called again, since it will fail to load the migrated storage.
+
+
 ## [Delegator](delegator/)
 
 Delegator patter is based around a low level cross contract call function `delegate_call`.

From f7d0903fc99c3391723f0a645ec63271be69971f Mon Sep 17 00:00:00 2001
From: Andrew Jones <ascjones@gmail.com>
Date: Wed, 31 Jan 2024 11:25:35 +0000
Subject: [PATCH 14/14] spelling

---
 integration-tests/upgradeable-contracts/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/integration-tests/upgradeable-contracts/README.md b/integration-tests/upgradeable-contracts/README.md
index 76e5496a0cb..74a25d764ce 100644
--- a/integration-tests/upgradeable-contracts/README.md
+++ b/integration-tests/upgradeable-contracts/README.md
@@ -16,7 +16,7 @@ You can read more about storage compatibility on [use.ink](https://use.ink/basic
 
 When upgrading a contract, the new code may have a different storage layout. This example illustrates a method to 
 migrate the storage from the old layout to the new layout. It does so by using an intermediate `migration` contract
-which performs the storage upgrade. Thw workflow is as follows:
+which performs the storage upgrade. The workflow is as follows:
 
 
 1. Upload a `migration` contract with a message `migrate` which performs the storage migration.