From 801bdce31e1275a2eef4938cc0e57a3dc59ed074 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Noah=20Bayindirli=20=F0=9F=A5=82?= <noah@primeprotocol.xyz>
Date: Thu, 6 Feb 2025 19:35:17 -0500
Subject: [PATCH] move storage into module

---
 .../src/tests/atomic_upgrades.rs              |  5 +++--
 .../src/tests/utils/dummy_contract.rs         | 18 ++++++++++-------
 .../utils/dummy_contract_after_upgrade.rs     | 20 +++++++++++--------
 3 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/contracts/stellar-upgrader/src/tests/atomic_upgrades.rs b/contracts/stellar-upgrader/src/tests/atomic_upgrades.rs
index 710f8dc8..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::{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
index d680f897..7b0606a6 100644
--- 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
index 1e71985f..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
@@ -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,
+    }
+}