From fd05d92a538203f6beecebf3560a25e8e5b45ce4 Mon Sep 17 00:00:00 2001
From: Kayanski <kowalski.kowalskin@gmail.com>
Date: Tue, 22 Oct 2024 08:47:50 +0000
Subject: [PATCH] State from state file load

---
 cw-orch-daemon/src/state.rs                    | 7 +++++++
 cw-orch-daemon/src/sync/core.rs                | 4 ++++
 packages/clone-testing/src/core.rs             | 4 ++++
 packages/cw-orch-core/src/environment/state.rs | 6 ++++++
 4 files changed, 21 insertions(+)

diff --git a/cw-orch-daemon/src/state.rs b/cw-orch-daemon/src/state.rs
index fa1295223..97b9a91f5 100644
--- a/cw-orch-daemon/src/state.rs
+++ b/cw-orch-daemon/src/state.rs
@@ -357,7 +357,9 @@ pub trait DeployedChains<Chain: CwEnv>: cw_orch_core::contract::Deploy<Chain> {
     /// Set the default contract state for a contract, so that users can retrieve it in their application when importing the library
     /// If a state is provided, it is used for all contracts, otherwise, the state is loaded from the crate's state file.
     fn set_contracts_state(&mut self, custom_state: Option<Value>) {
+        let mut is_loading_from_file = false;
         let Some(maybe_old_state) = custom_state.or_else(|| {
+            is_loading_from_file = true;
             Self::deployed_state_file_path()
                 .and_then(|state_file| crate::json_lock::read(&state_file).ok())
         }) else {
@@ -368,6 +370,11 @@ pub trait DeployedChains<Chain: CwEnv>: cw_orch_core::contract::Deploy<Chain> {
         let all_contracts = self.get_contracts_mut();
 
         for contract in all_contracts {
+            // If we're loading the state from file and the environment doesn't allow that, we don't load addresses and code-ids
+            if is_loading_from_file && !contract.environment().can_load_state_from_state_file() {
+                return;
+            }
+
             // We set the code_id and/or address of the contract in question if they are not present already
             let env_info = contract.environment().env_info();
             // We load the file
diff --git a/cw-orch-daemon/src/sync/core.rs b/cw-orch-daemon/src/sync/core.rs
index 7eb8615fa..63323a842 100644
--- a/cw-orch-daemon/src/sync/core.rs
+++ b/cw-orch-daemon/src/sync/core.rs
@@ -143,6 +143,10 @@ impl<Sender> ChainState for DaemonBase<Sender> {
     fn state(&self) -> Self::Out {
         self.daemon.state.clone()
     }
+
+    fn can_load_state_from_state_file(&self) -> bool {
+        true
+    }
 }
 
 // Execute on the real chain, returns tx response
diff --git a/packages/clone-testing/src/core.rs b/packages/clone-testing/src/core.rs
index 58772801f..9ff00b395 100644
--- a/packages/clone-testing/src/core.rs
+++ b/packages/clone-testing/src/core.rs
@@ -271,6 +271,10 @@ impl<S: StateInterface> ChainState for CloneTesting<S> {
     fn state(&self) -> Self::Out {
         self.state.clone()
     }
+
+    fn can_load_state_from_state_file(&self) -> bool {
+        true
+    }
 }
 
 // Execute on the test chain, returns test response type
diff --git a/packages/cw-orch-core/src/environment/state.rs b/packages/cw-orch-core/src/environment/state.rs
index 613712fa7..4e9c3511e 100644
--- a/packages/cw-orch-core/src/environment/state.rs
+++ b/packages/cw-orch-core/src/environment/state.rs
@@ -11,6 +11,12 @@ pub trait ChainState {
     type Out: StateInterface;
     /// Get the underlying state.
     fn state(&self) -> Self::Out;
+    /// Returns wether this environment can load state from the state file.
+    ///
+    /// This can be used within the Deploy trait to load the contracts from file if necessary
+    fn can_load_state_from_state_file(&self) -> bool {
+        false
+    }
 }
 
 /// This Interface allows for managing the local state of a deployment on any CosmWasm-supported environment.