From 4b0444f6fa8bc2af5c31f570f9424d1c62dfa7c2 Mon Sep 17 00:00:00 2001
From: Kevin Yang <5478483+k-yang@users.noreply.github.com>
Date: Thu, 14 Mar 2024 00:47:28 -0500
Subject: [PATCH] refactor(core-token-vesting-v2): deregister returns all
 unclaimed funds (#141)

---
 .../core-token-vesting-v2/src/contract.rs     | 14 +++--
 .../core-token-vesting-v2/src/testing.rs      | 54 +++++++++----------
 2 files changed, 30 insertions(+), 38 deletions(-)

diff --git a/contracts/core-token-vesting-v2/src/contract.rs b/contracts/core-token-vesting-v2/src/contract.rs
index 32ba435..acbfc13 100644
--- a/contracts/core-token-vesting-v2/src/contract.rs
+++ b/contracts/core-token-vesting-v2/src/contract.rs
@@ -308,18 +308,14 @@ fn deregister_vesting_account(
     VESTING_ACCOUNTS.remove(storage, address);
 
     let vested_amount = account.vested_amount(timestamp)?;
-    let claimed_amount = account.claimed_amount;
+    let left_vesting_amount = account.vesting_amount.checked_sub(vested_amount)?;
 
-    // transfer already vested amount to the user
-    let claimable_amount = vested_amount.checked_sub(claimed_amount)?;
-    send_if_amount_is_not_zero(messages, claimable_amount, &denom, address)?;
+    let recoverable_amount = account.vesting_amount - account.claimed_amount;
+    // transfer all that's unclaimed to the admin
 
-    // transfer left vesting amount to the admin
-    let left_vesting_amount =
-        account.vesting_amount.checked_sub(vested_amount)?;
     send_if_amount_is_not_zero(
         messages,
-        left_vesting_amount,
+        recoverable_amount,
         &denom,
         admin_address,
     )?;
@@ -330,6 +326,8 @@ fn deregister_vesting_account(
         ("vesting_amount", &account.vesting_amount.to_string()),
         ("vested_amount", &vested_amount.to_string()),
         ("left_vesting_amount", &left_vesting_amount.to_string()),
+        ("claimed_amount", &account.claimed_amount.to_string()),
+        ("recoverable_amount", &recoverable_amount.to_string()),
     ]))
 }
 
diff --git a/contracts/core-token-vesting-v2/src/testing.rs b/contracts/core-token-vesting-v2/src/testing.rs
index e73585d..01de747 100644
--- a/contracts/core-token-vesting-v2/src/testing.rs
+++ b/contracts/core-token-vesting-v2/src/testing.rs
@@ -942,36 +942,40 @@ fn deregister_successful() -> TestResult {
     // Set up the environment with a block time before the vesting start time
     let (mut deps, env) = setup_with_block_time(105)?;
 
-    let register_msg = ExecuteMsg::RewardUsers {
-        rewards: vec![RewardUserRequest {
-            user_address: "addr0001".to_string(),
-            vesting_amount: Uint128::new(5000u128),
-            cliff_amount: Uint128::new(1250u128),
-        }],
-        vesting_schedule: VestingSchedule::LinearVestingWithCliff {
-            start_time: Uint64::new(100),
-            cliff_time: Uint64::new(105),
-            end_time: Uint64::new(110),
-        },
-    };
-
     execute(
         deps.as_mut(),
         env.clone(), // Use the custom environment with the adjusted block time
         testing::mock_info("admin-sender", &[]),
-        register_msg,
+        ExecuteMsg::RewardUsers {
+            rewards: vec![RewardUserRequest {
+                user_address: "addr0001".to_string(),
+                vesting_amount: Uint128::new(5000u128),
+                cliff_amount: Uint128::new(1250u128),
+            }],
+            vesting_schedule: VestingSchedule::LinearVestingWithCliff {
+                start_time: Uint64::new(100),
+                cliff_time: Uint64::new(105),
+                end_time: Uint64::new(110),
+            },
+        },
     )?;
 
-    // Deregister with the manager address
-    let msg = ExecuteMsg::DeregisterVestingAccounts {
-        addresses: vec!["addr0001".to_string()],
-    };
+    // claim some of it
+    execute(
+        deps.as_mut(),
+        env.clone(),
+        testing::mock_info("addr0001", &[]),
+        ExecuteMsg::Claim {},
+    )?;
 
+    // Deregister with the manager address
     let res = execute(
         deps.as_mut(),
         env, // Use the custom environment with the adjusted block time
         testing::mock_info("manager-sender", &[]),
-        msg,
+        ExecuteMsg::DeregisterVestingAccounts {
+            addresses: vec!["addr0001".to_string()],
+        },
     )?;
     let data =
         from_json::<Vec<DeregisterUserResponse>>(res.data.unwrap()).unwrap();
@@ -984,19 +988,9 @@ fn deregister_successful() -> TestResult {
             error_msg: "".to_string(),
         }
     );
-    assert_eq!(res.messages.len(), 2);
+    assert_eq!(res.messages.len(), 1);
     assert_eq!(
         res.messages[0],
-        SubMsg::new(BankMsg::Send {
-            to_address: "addr0001".to_string(),
-            amount: vec![Coin {
-                denom: "token".to_string(),
-                amount: Uint128::new(1250u128),
-            }],
-        })
-    );
-    assert_eq!(
-        res.messages[1],
         SubMsg::new(BankMsg::Send {
             to_address: "admin-sender".to_string(),
             amount: vec![Coin {