From 3733891b4d551c69452e71b917408adbdcd22e2c Mon Sep 17 00:00:00 2001
From: Till Ziegler <tz@schoeneweide.tk>
Date: Mon, 15 Apr 2024 13:09:59 +0200
Subject: [PATCH] test: add migration test for terraport token

---
 contracts/cw20-taxed/src/contract.rs | 128 +++++++++++++++++++++++++++
 contracts/cw20-taxed/src/state.rs    |   2 +-
 2 files changed, 129 insertions(+), 1 deletion(-)

diff --git a/contracts/cw20-taxed/src/contract.rs b/contracts/cw20-taxed/src/contract.rs
index 6002a4d..2e98bcf 100644
--- a/contracts/cw20-taxed/src/contract.rs
+++ b/contracts/cw20-taxed/src/contract.rs
@@ -2098,6 +2098,10 @@ mod tests {
     }
 
     mod migration {
+        use std::{borrow::{Borrow, BorrowMut}, mem};
+
+        use self::state::migrate_v1::tests::mock_dependencies_with_terraport_balances;
+
         use super::*;
 
         use cosmwasm_std::Empty;
@@ -2215,6 +2219,130 @@ mod tests {
                 }]
             );
         }
+
+        #[test]
+        fn test_migrate_from_terraport() {
+            let mut deps = mock_dependencies_with_terraport_balances(vec![
+                // initial balances
+                (Addr::unchecked("addr1"), Uint128::new(1234), 123),
+                (Addr::unchecked("addr2"), Uint128::new(1234), 123),
+                (Addr::unchecked("addr3"), Uint128::new(4455), 123),
+
+                // mock a transfer at later height
+                (Addr::unchecked("addr1"), Uint128::new(1233), 456),
+                (Addr::unchecked("addr2"), Uint128::new(1235), 456),
+            ]);
+            
+            let env = mock_env();
+            crate::contract::migrate(deps.as_mut(), env, MigrateMsg { tax_map: None }).unwrap();
+
+            // balances are ok
+            match query_balance(deps.as_ref(), "addr1".to_string()){
+                Ok(balance) => {
+                    assert_eq!(
+                        balance,
+                        BalanceResponse{
+                            balance: Uint128::new(1233),
+                        }
+                    )
+                },
+                Err(e) => panic!("Error querying balance: {:?}", e),
+            }
+            match query_balance(deps.as_ref(), "addr2".to_string()){
+                Ok(balance) => {
+                    assert_eq!(
+                        balance,
+                        BalanceResponse{
+                            balance: Uint128::new(1235),
+                        }
+                    )
+                },
+                Err(e) => panic!("Error querying balance: {:?}", e),
+            }
+
+            // tax map is set
+            match TAX_INFO.load(deps.as_ref().storage) {
+                Ok(tax_info) => {
+                    assert_eq!(tax_info, TaxMap::default() )
+                },
+                Err(_) => panic!("Expected Tax map to be available!"),
+            }
+            
+        }
+
+        #[test]
+        fn test_migrate_from_terraport_with_explicit_tax_map() {
+            let mut deps = mock_dependencies_with_terraport_balances(vec![
+                // initial balances
+                (Addr::unchecked("addr1"), Uint128::new(1234), 123),
+                (Addr::unchecked("addr2"), Uint128::new(1234), 123),
+                (Addr::unchecked("addr3"), Uint128::new(4455), 123),
+
+                // mock a transfer at later height
+                (Addr::unchecked("addr1"), Uint128::new(1233), 456),
+                (Addr::unchecked("addr2"), Uint128::new(1235), 456),
+            ]);
+            let tax = TaxMap {
+                admin: Addr::unchecked("admin"),
+                on_transfer: TaxInfo{
+                    src_cond: TaxCondition::Never(TaxNeverCondition{}),
+                    dst_cond: TaxCondition::Never(TaxNeverCondition{}),
+                    proceeds: Addr::unchecked("proceeds1"),
+                },
+                on_transfer_from: TaxInfo {
+                    src_cond: TaxCondition::Never(TaxNeverCondition{}),
+                    dst_cond: TaxCondition::Never(TaxNeverCondition{}),
+                    proceeds: Addr::unchecked("proceeds2"),
+                },
+                on_send: TaxInfo {
+                    src_cond: TaxCondition::Never(TaxNeverCondition{}),
+                    dst_cond: TaxCondition::Never(TaxNeverCondition{}),
+                    proceeds: Addr::unchecked("proceeds3"),
+                
+                },
+                on_send_from: TaxInfo{
+                    src_cond: TaxCondition::Always(TaxAlwaysCondition{tax_rate: Decimal::percent(1)}),
+                    dst_cond: TaxCondition::Never(TaxNeverCondition{}),
+                    proceeds: Addr::unchecked("proceeds4"),
+                }
+            };
+            
+            let env = mock_env();
+            crate::contract::migrate(deps.as_mut(), env, MigrateMsg { tax_map: Some(tax.clone()) }).unwrap();
+
+            // balances are ok
+            match query_balance(deps.as_ref(), "addr1".to_string()){
+                Ok(balance) => {
+                    assert_eq!(
+                        balance,
+                        BalanceResponse{
+                            balance: Uint128::new(1233),
+                        }
+                    )
+                },
+                Err(e) => panic!("Error querying balance: {:?}", e),
+            }
+            match query_balance(deps.as_ref(), "addr2".to_string()){
+                Ok(balance) => {
+                    assert_eq!(
+                        balance,
+                        BalanceResponse{
+                            balance: Uint128::new(1235),
+                        }
+                    )
+                },
+                Err(e) => panic!("Error querying balance: {:?}", e),
+            }
+
+            // tax map is set
+            match TAX_INFO.load(deps.as_ref().storage) {
+                Ok(tax_info) => {
+                    assert_eq!(tax_info, tax )
+                },
+                Err(_) => panic!("Expected Tax map to be available!"),
+            }
+            
+        }
     }
 
     mod marketing {
diff --git a/contracts/cw20-taxed/src/state.rs b/contracts/cw20-taxed/src/state.rs
index d8d882a..c3f21e6 100644
--- a/contracts/cw20-taxed/src/state.rs
+++ b/contracts/cw20-taxed/src/state.rs
@@ -94,7 +94,7 @@ pub mod migrate_v1 {
         use cw2::set_contract_version;
 
         // mock a terraport style store
-        fn mock_dependencies_with_terraport_balances(
+        pub fn mock_dependencies_with_terraport_balances(
             balances: Vec<(Addr, Uint128, u64)>
         ) -> OwnedDeps<MockStorage, MockApi, MockQuerier> {
             let mut deps = mock_dependencies();