refactor: Validator Staking Adjustments

contracts/finance/andromeda-validator-staking/src/contract.rs

@@ -2,8 +2,9 @@ use std::str::FromStr;
 
 use crate::state::{DEFAULT_VALIDATOR, UNSTAKING_QUEUE};
 use cosmwasm_std::{
-    ensure, entry_point, Addr, BankMsg, Binary, Coin, CosmosMsg, Deps, DepsMut, DistributionMsg,
-    Env, FullDelegation, MessageInfo, Reply, Response, StakingMsg, StdError, SubMsg, Timestamp,
+    coin, ensure, entry_point, Addr, BankMsg, Binary, Coin, CosmosMsg, Deps, DepsMut,
+    DistributionMsg, Env, FullDelegation, MessageInfo, Reply, Response, StakingMsg, StdError,
+    SubMsg, Timestamp, Uint128,
 };
 use cw2::set_contract_version;
 
@@ -78,7 +79,7 @@ pub fn execute(
 pub fn handle_execute(ctx: ExecuteContext, msg: ExecuteMsg) -> Result<Response, ContractError> {
     match msg {
         ExecuteMsg::Stake { validator } => execute_stake(ctx, validator),
-        ExecuteMsg::Unstake { validator } => execute_unstake(ctx, validator),
+        ExecuteMsg::Unstake { validator, amount } => execute_unstake(ctx, validator, amount),
         ExecuteMsg::Claim {
             validator,
             recipient,
@@ -138,6 +139,7 @@ fn execute_stake(ctx: ExecuteContext, validator: Option<Addr>) -> Result<Respons
 fn execute_unstake(
     ctx: ExecuteContext,
     validator: Option<Addr>,
+    amount: Option<Uint128>,
 ) -> Result<Response, ContractError> {
     let ExecuteContext {
         deps, info, env, ..
@@ -166,8 +168,10 @@ fn execute_unstake(
         });
     };
 
+    let unstake_amount = amount.unwrap_or(res.amount.amount);
+
     ensure!(
-        !res.amount.amount.is_zero(),
+        !unstake_amount.is_zero() && unstake_amount <= res.amount.amount,
         ContractError::InvalidValidatorOperation {
             operation: "Unstake".to_string(),
             validator: validator.to_string(),
@@ -176,13 +180,14 @@ fn execute_unstake(
 
     let undelegate_msg = CosmosMsg::Staking(StakingMsg::Undelegate {
         validator: validator.to_string(),
-        amount: res.amount,
+        amount: coin(unstake_amount.u128(), res.amount.denom),
     });
     let undelegate_msg = SubMsg::reply_on_success(undelegate_msg, ReplyId::ValidatorUnstake.repr());
 
     let res = Response::new()
         .add_submessage(undelegate_msg)
         .add_attribute("action", "validator-unstake")
+        .add_attribute("amount", unstake_amount)
         .add_attribute("from", info.sender)
         .add_attribute("to", validator.to_string());
 
@@ -204,15 +209,15 @@ fn execute_claim(
     // Check if the validator is valid before unstaking
     is_validator(&deps, &validator)?;
 
-    let recipient = if let Some(recipient) = recipient {
+    let recipient_address = if let Some(ref recipient) = recipient {
         recipient.get_raw_address(&deps.as_ref())?
     } else {
         info.sender
     };
 
     // Ensure recipient is the contract owner
     ensure!(
-        ADOContract::default().is_contract_owner(deps.storage, recipient.as_str())?,
+        ADOContract::default().is_contract_owner(deps.storage, recipient_address.as_str())?,
         ContractError::Unauthorized {}
     );
 
@@ -235,13 +240,13 @@ fn execute_claim(
 
     let res = Response::new()
         .add_message(DistributionMsg::SetWithdrawAddress {
-            address: recipient.to_string(),
+            address: recipient_address.to_string(),
         })
         .add_message(DistributionMsg::WithdrawDelegatorReward {
             validator: validator.to_string(),
         })
         .add_attribute("action", "validator-claim-reward")
-        .add_attribute("recipient", recipient)
+        .add_attribute("recipient", recipient_address)
         .add_attribute("validator", validator.to_string());
 
     Ok(res)

contracts/finance/andromeda-validator-staking/src/mock.rs

@@ -1,5 +1,5 @@
 use andromeda_finance::validator_staking::{ExecuteMsg, InstantiateMsg, QueryMsg, UnstakingTokens};
-use cosmwasm_std::{Addr, Coin, Delegation, Empty};
+use cosmwasm_std::{Addr, Coin, Delegation, Empty, Uint128};
 
 use crate::contract::{execute, instantiate, query, reply};
 use andromeda_testing::{
@@ -31,8 +31,9 @@ impl MockValidatorStaking {
         app: &mut MockApp,
         sender: Addr,
         validator: Option<Addr>,
+        amount: Option<Uint128>,
     ) -> ExecuteResult {
-        let msg = mock_execute_unstake(validator);
+        let msg = mock_execute_unstake(validator, amount);
         self.execute(app, &msg, sender, &[])
     }
 
@@ -94,8 +95,8 @@ pub fn mock_execute_stake(validator: Option<Addr>) -> ExecuteMsg {
     ExecuteMsg::Stake { validator }
 }
 
-pub fn mock_execute_unstake(validator: Option<Addr>) -> ExecuteMsg {
-    ExecuteMsg::Unstake { validator }
+pub fn mock_execute_unstake(validator: Option<Addr>, amount: Option<Uint128>) -> ExecuteMsg {
+    ExecuteMsg::Unstake { validator, amount }
 }
 
 pub fn mock_execute_claim_reward(

contracts/finance/andromeda-validator-staking/src/testing/tests.rs

@@ -143,6 +143,7 @@ fn test_unauthorized_unstake() {
 
     let msg = ExecuteMsg::Unstake {
         validator: Some(valid_validator),
+        amount: None,
     };
 
     let info = mock_info("other", &[coin(100, "uandr")]);

packages/andromeda-finance/src/validator_staking.rs

@@ -1,6 +1,6 @@
 use andromeda_std::{amp::AndrAddr, andr_exec, andr_instantiate, andr_query, error::ContractError};
 use cosmwasm_schema::{cw_serde, QueryResponses};
-use cosmwasm_std::{Addr, Coin, DepsMut, Timestamp};
+use cosmwasm_std::{Addr, Coin, DepsMut, Timestamp, Uint128};
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 
@@ -18,6 +18,7 @@ pub enum ExecuteMsg {
     },
     Unstake {
         validator: Option<Addr>,
+        amount: Option<Uint128>,
     },
     Claim {
         validator: Option<Addr>,

tests-integration/tests/validator_staking.rs

@@ -14,7 +14,7 @@ use andromeda_validator_staking::mock::{
 use andromeda_std::error::ContractError::Std;
 use andromeda_testing::MockContract;
 use cosmwasm_std::StdError::GenericErr;
-use cosmwasm_std::{coin, to_json_binary, Addr, BlockInfo};
+use cosmwasm_std::{coin, to_json_binary, Addr, BlockInfo, Delegation, Uint128};
 
 #[test]
 fn test_validator_stake() {
@@ -123,6 +123,7 @@ fn test_validator_stake() {
             &mut router,
             owner.clone(),
             Some(Addr::unchecked("fake_validator")),
+            None,
         )
         .unwrap_err();
     // let _err = err.root_cause().downcast_ref::<ContractError>().unwrap();
@@ -136,6 +137,7 @@ fn test_validator_stake() {
             &mut router,
             Addr::unchecked("other"),
             Some(Addr::unchecked("fake_validator")),
+            None,
         )
         .unwrap_err();
     // let _err = err.root_cause().downcast_ref::<ContractError>().unwrap();
@@ -144,7 +146,7 @@ fn test_validator_stake() {
     // assert_eq!(err, &expected_err);
 
     validator_staking
-        .execute_unstake(&mut router, owner.clone(), None)
+        .execute_unstake(&mut router, owner.clone(), None, None)
         .unwrap();
 
     // Test staked token query from undelegated validator
@@ -191,3 +193,179 @@ fn test_validator_stake() {
     let owner_balance = router.wrap().query_balance(owner, "TOKEN").unwrap();
     assert_eq!(owner_balance, coin(1050, "TOKEN"));
 }
+
+#[test]
+fn test_validator_stake_and_unstake_specific_amount() {
+    let mut router = mock_app(Some(vec!["TOKEN"]));
+
+    let andr = MockAndromedaBuilder::new(&mut router, "admin")
+        .with_wallets(vec![("owner", vec![coin(1000, "TOKEN")])])
+        .with_contracts(vec![
+            ("app-contract", mock_andromeda_app()),
+            ("validator-staking", mock_andromeda_validator_staking()),
+        ])
+        .build(&mut router);
+    let owner = andr.get_wallet("owner");
+    let recipient = AndrAddr::from_string(owner.to_string());
+    let validator_1 = router.api().addr_make("validator1");
+
+    let validator_staking_init_msg = mock_validator_staking_instantiate_msg(
+        validator_1.clone(),
+        None,
+        andr.kernel.addr().to_string(),
+    );
+
+    let validator_staking_component = AppComponent::new(
+        "staking".to_string(),
+        "validator-staking".to_string(),
+        to_json_binary(&validator_staking_init_msg).unwrap(),
+    );
+
+    let app_components = vec![validator_staking_component.clone()];
+    let app = MockAppContract::instantiate(
+        andr.get_code_id(&mut router, "app-contract"),
+        owner,
+        &mut router,
+        "Validator Staking App",
+        app_components,
+        andr.kernel.addr(),
+        Some(owner.to_string()),
+    );
+
+    let validator_staking: MockValidatorStaking =
+        app.query_ado_by_component_name(&router, validator_staking_component.name);
+
+    let funds = vec![coin(1000, "TOKEN")];
+
+    validator_staking
+        .execute_stake(&mut router, owner.clone(), None, funds)
+        .unwrap();
+
+    let stake_info = validator_staking
+        .query_staked_tokens(&router, None)
+        .unwrap();
+    assert_eq!(stake_info.validator, validator_1.to_string());
+
+    // Testing when there is no reward to claim
+    let _err = validator_staking
+        .execute_claim_reward(
+            &mut router,
+            owner.clone(),
+            Some(validator_1.clone()),
+            Some(recipient.clone()),
+        )
+        .unwrap_err();
+    // assert_eq!(may_err.unwrap(), &expected_err);
+
+    // wait 1/2 year
+    router.set_block(BlockInfo {
+        height: router.block_info().height,
+        time: router
+            .block_info()
+            .time
+            .plus_seconds(60 * 60 * 24 * 365 / 2),
+        chain_id: router.block_info().chain_id,
+    });
+
+    // only owner can become a recipient
+    let _err = validator_staking
+        .execute_claim_reward(
+            &mut router,
+            owner.clone(),
+            Some(validator_1.clone()),
+            Some(AndrAddr::from_string("some_address")),
+        )
+        .unwrap_err();
+
+    validator_staking
+        .execute_claim_reward(
+            &mut router,
+            owner.clone(),
+            Some(validator_1),
+            Some(recipient),
+        )
+        .unwrap();
+
+    // Default APR 10% by cw-multi-test -> StakingInfo
+    // should now have 1000 * 10% / 2 - 0% commission = 50 tokens reward
+    let owner_balance = router.wrap().query_balance(owner.clone(), "TOKEN").unwrap();
+    assert_eq!(owner_balance, coin(50, "TOKEN"));
+
+    // Test unstake with invalid validator
+    let _err = validator_staking
+        .execute_unstake(
+            &mut router,
+            owner.clone(),
+            Some(Addr::unchecked("fake_validator")),
+            None,
+        )
+        .unwrap_err();
+
+    // Test unstake from invalid owner
+    let _err = validator_staking
+        .execute_unstake(
+            &mut router,
+            Addr::unchecked("other"),
+            Some(Addr::unchecked("fake_validator")),
+            None,
+        )
+        .unwrap_err();
+
+    // Try unstaking an amount larger than staked
+    validator_staking
+        .execute_unstake(&mut router, owner.clone(), None, Some(Uint128::new(2_000)))
+        .unwrap_err();
+
+    // Try unstaking a zero amount
+    validator_staking
+        .execute_unstake(&mut router, owner.clone(), None, Some(Uint128::zero()))
+        .unwrap_err();
+
+    validator_staking
+        .execute_unstake(&mut router, owner.clone(), None, Some(Uint128::new(200)))
+        .unwrap();
+
+    // Test staked token query from undelegated validator
+    let delegation = validator_staking
+        .query_staked_tokens(&router, None)
+        .unwrap();
+    assert_eq!(
+        delegation,
+        Delegation {
+            delegator: Addr::unchecked(
+                "andr12sjjp2n9epg6x8g4adaumlfp4ulaj5jvcu5rj29uk02zcdvrfkuqdfct9e"
+            ),
+            validator: "andr1qcxce9c4thzxnfmpr2dqnnlqea9ey35y7tnke37fymfcgzte0zwshp76a9"
+                .to_string(),
+            amount: coin(800_u128, "TOKEN")
+        }
+    );
+
+    // Test withdraw before payout period
+    let _err = validator_staking
+        .execute_withdraw_fund(&mut router, owner.clone())
+        .unwrap_err();
+
+    let unstaked_tokens = validator_staking.query_unstaked_tokens(&router).unwrap();
+    let unbonding_period =
+        unstaked_tokens[0].payout_at.seconds() - router.block_info().time.seconds();
+    // Update block to payout period
+    router.set_block(BlockInfo {
+        height: router.block_info().height,
+        time: router.block_info().time.plus_seconds(unbonding_period),
+        chain_id: router.block_info().chain_id,
+    });
+
+    router.set_block(BlockInfo {
+        height: router.block_info().height,
+        time: router.block_info().time.plus_seconds(1),
+        chain_id: router.block_info().chain_id,
+    });
+
+    validator_staking
+        .execute_withdraw_fund(&mut router, owner.clone())
+        .unwrap();
+
+    let owner_balance = router.wrap().query_balance(owner, "TOKEN").unwrap();
+    assert_eq!(owner_balance, coin(250, "TOKEN"));
+}