From 46cdb794315cf1fc4cc270399d5758d8d49623c1 Mon Sep 17 00:00:00 2001
From: karim-en <karim.alabtakh@nearone.org>
Date: Tue, 4 Mar 2025 19:06:52 +0000
Subject: [PATCH] chore: verify the state after wrong migration

---
 near-plugins-derive/src/upgradable.rs         | 12 +++--
 .../upgradable_state_migration/src/lib.rs     |  6 +++
 near-plugins-derive/tests/upgradable.rs       | 47 +++++++++++++++++++
 3 files changed, 60 insertions(+), 5 deletions(-)

diff --git a/near-plugins-derive/src/upgradable.rs b/near-plugins-derive/src/upgradable.rs
index 86cc0a1..e53015b 100644
--- a/near-plugins-derive/src/upgradable.rs
+++ b/near-plugins-derive/src/upgradable.rs
@@ -218,17 +218,19 @@ pub fn derive_upgradable(input: TokenStream) -> TokenStream {
                     )
                 }
 
-                let promise = ::near_sdk::Promise::new(::near_sdk::env::current_account_id())
+                let deploy_promise = ::near_sdk::Promise::new(::near_sdk::env::current_account_id())
                     .deploy_contract(code);
-                match function_call_args {
-                    None => promise.function_call("up_verify_state".to_owned(), vec![], near_sdk::NearToken::from_yoctonear(0), near_sdk::Gas::from_tgas(2)),
+                let promise = match function_call_args {
+                    None => deploy_promise,
                     Some(args) => {
                         // Execute the `DeployContract` and `FunctionCall` actions in a batch
                         // transaction to make a failure of the function call roll back the code
                         // deployment.
-                        promise.function_call(args.function_name, args.arguments, args.amount, args.gas)
+                        deploy_promise.function_call(args.function_name, args.arguments, args.amount, args.gas)
                     },
-                }
+                };
+
+                promise.function_call("up_verify_state".to_owned(), vec![], near_sdk::NearToken::from_yoctonear(0), near_sdk::Gas::from_tgas(2))
             }
 
             fn up_verify_state(&self) {}
diff --git a/near-plugins-derive/tests/contracts/upgradable_state_migration/src/lib.rs b/near-plugins-derive/tests/contracts/upgradable_state_migration/src/lib.rs
index 784a499..daf6f8b 100644
--- a/near-plugins-derive/tests/contracts/upgradable_state_migration/src/lib.rs
+++ b/near-plugins-derive/tests/contracts/upgradable_state_migration/src/lib.rs
@@ -65,6 +65,12 @@ impl Contract {
     pub fn is_migrated(&self) -> bool {
         self.is_migrated
     }
+
+    pub fn migrate_empty() -> bool {
+        // This method is used to test the rollback mechanism of `Upgradable::up_deploy_code` when
+        // the migration method is empty.
+        true
+    }
 }
 
 /// Corresponds to the state defined in the initial `../upgradable` contract.
diff --git a/near-plugins-derive/tests/upgradable.rs b/near-plugins-derive/tests/upgradable.rs
index a8b9936..fbe8d69 100644
--- a/near-plugins-derive/tests/upgradable.rs
+++ b/near-plugins-derive/tests/upgradable.rs
@@ -681,6 +681,53 @@ async fn test_deploy_code_with_missed_migration() -> anyhow::Result<()> {
     Ok(())
 }
 
+/// Deploys a new version of the contract with wrong migration
+/// Verifies the failure rolls back the deployment, i.e. the initial
+/// code remains active.
+#[tokio::test]
+async fn test_deploy_code_with_wrong_migration() -> anyhow::Result<()> {
+    let worker = near_workspaces::sandbox().await?;
+    let dao = worker.dev_create_account().await?;
+    let setup = Setup::new(worker.clone(), Some(dao.id().clone()), None).await?;
+
+    // Compile the other version of the contract and stage its code.
+    let code = common::repo::compile_project(
+        Path::new(PROJECT_PATH_STATE_MIGRATION),
+        "upgradable_state_migration",
+    )
+    .await?;
+    let res = setup
+        .upgradable_contract
+        .up_stage_code(&dao, code.clone())
+        .await?;
+    assert_success_with_unit_return(res);
+    setup.assert_staged_code(Some(&code)).await;
+
+    // Deploy staged code with wrong migration function
+    let function_call_args: FunctionCallArgs = FunctionCallArgs {
+        function_name: "migrate_empty".to_string(),
+        arguments: Vec::new(),
+        amount: NearToken::from_yoctonear(0),
+        gas: Gas::from_tgas(3),
+    };
+
+    let res = setup
+        .upgradable_contract
+        .up_deploy_code(
+            &dao,
+            convert_code_to_deploy_hash(&code),
+            Some(function_call_args),
+        )
+        .await?;
+    assert_failure_with(res, "Cannot deserialize the contract state");
+
+    // Verify `code` wasn't deployed by calling a function that is defined only in the initial
+    // contract but not in the contract corresponding to the `code`.
+    setup.assert_is_set_up(&setup.unauth_account).await;
+
+    Ok(())
+}
+
 /// Deploys staged code in a batch transaction with two function call actions:
 ///
 /// 1. `up_deploy_code` with a function call to a migration method that fails