implement migration and update old migration tests

contracts/transmuter/Cargo.toml

@@ -2,12 +2,12 @@
 authors = ["Supanat Potiwarakorn <supanat.ptk@gmail.com>"]
 edition = "2021"
 name = "transmuter"
-version = "3.2.0"
+version = "4.0.0"
 
 exclude = [
-    # Those files are rust-optimizer artifacts. You might want to commit them for convenience but they should not be part of the source code publication.
-    "contract.wasm",
-    "hash.txt",
+  # Those files are rust-optimizer artifacts. You might want to commit them for convenience but they should not be part of the source code publication.
+  "contract.wasm",
+  "hash.txt",
 ]
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
@@ -42,25 +42,25 @@ optimize = """docker run --rm -v "$(pwd)":/code \
 """
 
 [dependencies]
-cosmwasm-schema = { workspace = true }
-cosmwasm-std = { workspace = true, features = ["cosmwasm_1_1"] }
+cosmwasm-schema = {workspace = true}
+cosmwasm-std = {workspace = true, features = ["cosmwasm_1_1"]}
 cosmwasm-storage = "1.3.1"
 cw-storage-plus = "1.1.0"
 cw2 = "1.1.0"
 osmosis-std = "0.22.0"
 schemars = "0.8.12"
-serde = { version = "1.0.183", default-features = false, features = ["derive"] }
+serde = {version = "1.0.183", default-features = false, features = ["derive"]}
 sylvia = "0.10.1"
-thiserror = { version = "1.0.44" }
-transmuter_math = { version = "1.0.0", path = "../../packages/transmuter_math" }
+thiserror = {version = "1.0.44"}
+transmuter_math = {version = "1.0.0", path = "../../packages/transmuter_math"}
 
 [dev-dependencies]
 itertools = "0.12.0"
 osmosis-test-tube = "22.1.0"
 rstest = "0.18.2"
 
 [lints.rust]
-unexpected_cfgs = { level = "warn", check-cfg = [
-    'cfg(tarpaulin)',
-    'cfg(tarpaulin_include)',
-] }
+unexpected_cfgs = {level = "warn", check-cfg = [
+  'cfg(tarpaulin)',
+  'cfg(tarpaulin_include)',
+]}

contracts/transmuter/src/lib.rs

@@ -102,9 +102,9 @@ mod entry_points {
     pub fn migrate(
         deps: DepsMut,
         _env: Env,
-        _msg: migrations::v3_2_0::MigrateMsg,
+        _msg: migrations::v4_0_0::MigrateMsg,
     ) -> Result<Response, ContractError> {
-        migrations::v3_2_0::execute_migration(deps)
+        migrations::v4_0_0::execute_migration(deps)
     }
 }
 

contracts/transmuter/src/migrations/mod.rs

@@ -1 +1 @@
-pub mod v3_2_0;
+pub mod v4_0_0;

contracts/transmuter/src/migrations/v4_0_0.rs

@@ -0,0 +1,161 @@
+use std::collections::BTreeMap;
+
+use cosmwasm_schema::cw_serde;
+
+use cosmwasm_std::{ensure_eq, DepsMut, Response, Storage};
+use cw2::{ContractVersion, VersionError, CONTRACT};
+use cw_storage_plus::Item;
+
+use crate::{
+    asset::Asset,
+    contract::{key, CONTRACT_NAME, CONTRACT_VERSION},
+    transmuter_pool::TransmuterPool,
+    ContractError,
+};
+
+const FROM_VERSION: &str = "3.2.0";
+const TO_VERSION: &str = "4.0.0";
+
+#[cw_serde]
+pub struct MigrateMsg {}
+
+#[cw_serde]
+pub struct TransmuterPoolV3 {
+    pub pool_assets: Vec<Asset>,
+    // [to-be-added] pub asset_groups: BTreeMap<String, AssetGroup>,
+}
+
+pub fn execute_migration(deps: DepsMut) -> Result<Response, ContractError> {
+    // Assert that the stored contract version matches the expected version before migration
+    assert_contract_versions(deps.storage, CONTRACT_NAME, FROM_VERSION)?;
+
+    // Ensure that the current contract version matches the target version to prevent migration to an incorrect version
+    ensure_eq!(
+        CONTRACT_VERSION,
+        TO_VERSION,
+        cw2::VersionError::WrongVersion {
+            expected: TO_VERSION.to_string(),
+            found: CONTRACT_VERSION.to_string()
+        }
+    );
+
+    // add asset groups to the pool
+    let pool_v3: TransmuterPoolV3 =
+        Item::<'_, TransmuterPoolV3>::new(key::POOL).load(deps.storage)?;
+
+    let pool_v4 = TransmuterPool {
+        pool_assets: pool_v3.pool_assets,
+        asset_groups: BTreeMap::new(),
+    };
+
+    Item::<'_, TransmuterPool>::new(key::POOL).save(deps.storage, &pool_v4)?;
+
+    // Set the contract version to the target version after successful migration
+    cw2::set_contract_version(deps.storage, CONTRACT_NAME, TO_VERSION)?;
+
+    // Return a response with an attribute indicating the method that was executed
+    Ok(Response::new().add_attribute("method", "v4_0_0/execute_migraiton"))
+}
+
+/// Assert that the stored contract version info matches the given value.
+/// This is useful during migrations, for making sure that the correct contract
+/// is being migrated, and it's being migrated from the correct version.
+fn assert_contract_versions(
+    storage: &dyn Storage,
+    expected_contract: &str,
+    expected_version: &str,
+) -> Result<(), VersionError> {
+    let ContractVersion { contract, version } = match CONTRACT.may_load(storage)? {
+        Some(contract) => contract,
+        None => return Err(VersionError::NotFound),
+    };
+
+    if contract != expected_contract {
+        return Err(VersionError::WrongContract {
+            expected: expected_contract.into(),
+            found: contract,
+        });
+    }
+
+    if version.as_str() != expected_version {
+        return Err(VersionError::WrongVersion {
+            expected: expected_version.to_string(),
+            found: version,
+        });
+    }
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use cosmwasm_std::{testing::mock_dependencies, Uint128};
+
+    use super::*;
+
+    #[test]
+    fn test_successful_migration() {
+        let mut deps = mock_dependencies();
+
+        cw2::set_contract_version(&mut deps.storage, CONTRACT_NAME, FROM_VERSION).unwrap();
+
+        let pool_assets = vec![
+            Asset::new(Uint128::from(100u128), "uusdt", Uint128::from(1u128)).unwrap(),
+            Asset::new(Uint128::from(200u128), "uusdc", Uint128::from(1u128)).unwrap(),
+        ];
+        let pool_v3 = TransmuterPoolV3 {
+            pool_assets: pool_assets.clone(),
+        };
+
+        Item::new(key::POOL)
+            .save(&mut deps.storage, &pool_v3)
+            .unwrap();
+
+        let res = execute_migration(deps.as_mut()).unwrap();
+
+        let pool = Item::<'_, TransmuterPool>::new(key::POOL)
+            .load(&deps.storage)
+            .unwrap();
+
+        assert_eq!(
+            pool,
+            TransmuterPool {
+                pool_assets,
+                asset_groups: BTreeMap::new() // migrgate with empty asset groups
+            }
+        );
+
+        assert_eq!(
+            res,
+            Response::new().add_attribute("method", "v4_0_0/execute_migraiton")
+        );
+    }
+
+    #[test]
+    fn test_invalid_version() {
+        let mut deps = mock_dependencies();
+
+        let pool_assets = vec![
+            Asset::new(Uint128::from(100u128), "uusdt", Uint128::from(1u128)).unwrap(),
+            Asset::new(Uint128::from(200u128), "uusdc", Uint128::from(1u128)).unwrap(),
+        ];
+        let pool_v3 = TransmuterPoolV3 {
+            pool_assets: pool_assets.clone(),
+        };
+
+        Item::new(key::POOL)
+            .save(&mut deps.storage, &pool_v3)
+            .unwrap();
+
+        cw2::set_contract_version(&mut deps.storage, CONTRACT_NAME, "3.0.0").unwrap();
+
+        let err = execute_migration(deps.as_mut()).unwrap_err();
+        assert_eq!(
+            err,
+            ContractError::VersionError(cw2::VersionError::WrongVersion {
+                expected: FROM_VERSION.to_string(),
+                found: "3.0.0".to_string()
+            })
+        );
+    }
+}

contracts/transmuter/src/swap.rs

@@ -679,7 +679,6 @@ fn construct_scope_value_pairs(
     let mut asset_group_weight_pairs: HashMap<Scope, (Decimal, Decimal)> = HashMap::new();
 
     // Reverse index the asset groups
-    // TODO: handle cases where asset group contains denom that does not exist
     let mut asset_groups_of_denom = HashMap::new();
     for (group, asset_group) in asset_group {
         for denom in asset_group.into_denoms() {