move storage into module

axelarnetwork · Feb 7, 2025 · 801bdce · 801bdce
1 parent 08cb8fc
commit 801bdce
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 17 deletions.
diff --git 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::{self, 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 = utils::data(&env);
+        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
@@ -38,14 +38,18 @@ impl DummyContract {
     }
 }
 
-#[contractstorage]
-pub enum DataKey {
-    #[instance]
-    #[value(soroban_sdk::String)]
-    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
@@ -39,20 +39,24 @@ impl DummyContract {
 
     pub fn migrate(env: Env, migration_data: soroban_sdk::String) -> Result<(), ContractError> {
         Self::owner(&env).require_auth();
-        set_data(&env, &migration_data);
+        storage::set_data(&env, &migration_data);
 
         Ok(())
     }
 }
 
-#[contractstorage]
-pub enum DataKey {
-    #[instance]
-    #[value(soroban_sdk::String)]
-    Data,
-}
-
 #[contracterror]
 pub enum ContractError {
     SomeFailure = 1,
 }
+
+mod storage {
+    use super::*;
+
+    #[contractstorage]
+    pub enum DataKey {
+        #[instance]
+        #[value(soroban_sdk::String)]
+        Data,
+    }
+}