diff --git a/Cargo.lock b/Cargo.lock
index 239c255d..ae0dd0de 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2105,6 +2105,7 @@ name = "stellar-upgrader"
 version = "1.0.0"
 dependencies = [
  "cfg-if",
+ "goldie",
  "soroban-sdk",
  "stellar-axelar-std",
 ]
diff --git a/contracts/stellar-upgrader/Cargo.toml b/contracts/stellar-upgrader/Cargo.toml
index 8f58907f..1ca65b43 100644
--- a/contracts/stellar-upgrader/Cargo.toml
+++ b/contracts/stellar-upgrader/Cargo.toml
@@ -15,6 +15,7 @@ soroban-sdk = { workspace = true }
 stellar-axelar-std = { workspace = true }
 
 [dev-dependencies]
+goldie = { workspace = true }
 soroban-sdk = { workspace = true, features = ["testutils"] }
 
 [features]
diff --git a/contracts/stellar-upgrader/src/lib.rs b/contracts/stellar-upgrader/src/lib.rs
index c91b1ca1..a8288af0 100644
--- a/contracts/stellar-upgrader/src/lib.rs
+++ b/contracts/stellar-upgrader/src/lib.rs
@@ -1,4 +1,8 @@
 #![no_std]
+
+#[cfg(any(test, feature = "testutils"))]
+extern crate std;
+
 #[cfg(test)]
 extern crate alloc;
 
diff --git a/contracts/stellar-upgrader/src/tests/atomic_upgrades.rs b/contracts/stellar-upgrader/src/tests/atomic_upgrades.rs
index 7c78bac8..890d4290 100644
--- a/contracts/stellar-upgrader/src/tests/atomic_upgrades.rs
+++ b/contracts/stellar-upgrader/src/tests/atomic_upgrades.rs
@@ -2,8 +2,9 @@ use soroban_sdk::testutils::Address as _;
 use soroban_sdk::{Address, BytesN, Env, String};
 use stellar_axelar_std::{assert_contract_err, mock_auth};
 
-use super::utils::{DataKey, DummyContract, DummyContractClient};
+use super::utils::{DummyContract, DummyContractClient};
 use crate::error::ContractError;
+use crate::tests::utils;
 use crate::{Upgrader, UpgraderClient};
 
 const WASM_AFTER_UPGRADE: &[u8] = include_bytes!("testdata/dummy.wasm");
@@ -42,7 +43,7 @@ fn upgrade_and_migrate_are_atomic() {
 
     // ensure migration was successful
     env.as_contract(&contract_address, || {
-        let data: String = env.storage().instance().get(&DataKey::Data).unwrap();
+        let data = utils::storage::data(&env);
         assert_eq!(data, expected_data);
     });
 }
diff --git a/contracts/stellar-upgrader/src/tests/utils/dummy_contract.rs b/contracts/stellar-upgrader/src/tests/utils/dummy_contract.rs
index f0e2b7b2..7b0606a6 100644
--- a/contracts/stellar-upgrader/src/tests/utils/dummy_contract.rs
+++ b/contracts/stellar-upgrader/src/tests/utils/dummy_contract.rs
@@ -1,8 +1,8 @@
 //! Dummy contract to test the [crate::Upgrader]
 
 use soroban_sdk::{contract, contracterror, contractimpl, contracttype, Address, BytesN, Env};
-use stellar_axelar_std::interfaces;
 use stellar_axelar_std::interfaces::{OwnableInterface, UpgradableInterface};
+use stellar_axelar_std::{contractstorage, interfaces};
 
 #[contract]
 pub struct DummyContract;
@@ -38,12 +38,18 @@ impl DummyContract {
     }
 }
 
-#[contracttype]
-pub enum DataKey {
-    Data,
-}
-
 #[contracterror]
 pub enum ContractError {
     SomeFailure = 1,
 }
+
+pub mod storage {
+    use super::*;
+
+    #[contractstorage]
+    pub enum DataKey {
+        #[instance]
+        #[value(soroban_sdk::String)]
+        Data,
+    }
+}
diff --git a/contracts/stellar-upgrader/src/tests/utils/dummy_contract_after_upgrade.rs b/contracts/stellar-upgrader/src/tests/utils/dummy_contract_after_upgrade.rs
index 2c1f9a4b..f99ad4c9 100644
--- a/contracts/stellar-upgrader/src/tests/utils/dummy_contract_after_upgrade.rs
+++ b/contracts/stellar-upgrader/src/tests/utils/dummy_contract_after_upgrade.rs
@@ -1,8 +1,8 @@
 //! Base for the dummy.wasm file. This is the dummy contract after upgrade.
 
 use soroban_sdk::{contract, contracterror, contractimpl, contracttype, Address, BytesN, Env};
-use stellar_axelar_std::interfaces;
 use stellar_axelar_std::interfaces::{OwnableInterface, UpgradableInterface};
+use stellar_axelar_std::{contractstorage, interfaces};
 
 #[contract]
 pub struct DummyContract;
@@ -39,19 +39,24 @@ impl DummyContract {
 
     pub fn migrate(env: Env, migration_data: soroban_sdk::String) -> Result<(), ContractError> {
         Self::owner(&env).require_auth();
-        env.storage()
-            .instance()
-            .set(&DataKey::Data, &migration_data);
+        storage::set_data(&env, &migration_data);
+
         Ok(())
     }
 }
 
-#[contracttype]
-pub enum DataKey {
-    Data,
-}
-
 #[contracterror]
 pub enum ContractError {
     SomeFailure = 1,
 }
+
+mod storage {
+    use super::*;
+
+    #[contractstorage]
+    pub enum DataKey {
+        #[instance]
+        #[value(soroban_sdk::String)]
+        Data,
+    }
+}
diff --git a/contracts/stellar-upgrader/src/tests/utils/testdata/ensure_data_key_storage_schema_is_unchanged.golden b/contracts/stellar-upgrader/src/tests/utils/testdata/ensure_data_key_storage_schema_is_unchanged.golden
new file mode 100644
index 00000000..75a133d8
--- /dev/null
+++ b/contracts/stellar-upgrader/src/tests/utils/testdata/ensure_data_key_storage_schema_is_unchanged.golden
@@ -0,0 +1,6 @@
+pub enum DataKey {
+
+    #[instance]
+    #[value(soroban_sdk::String)]
+    Data,
+}