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;
 
@@ -14,7 +15,7 @@ use andromeda_finance::validator_staking::{
 use andromeda_std::{
     ado_base::InstantiateMsg as BaseInstantiateMsg,
     ado_contract::ADOContract,
-    amp::AndrAddr,
+    amp::Recipient,
     common::{context::ExecuteContext, encode_binary},
     error::ContractError,
 };
@@ -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,12 @@ fn execute_unstake(
         });
     };
 
+    let unstake_amount = amount.unwrap_or(res.amount.amount);
+
     ensure!(
-        !res.amount.amount.is_zero(),
+        !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 +182,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());
 
@@ -192,7 +199,7 @@ fn execute_unstake(
 fn execute_claim(
     ctx: ExecuteContext,
     validator: Option<Addr>,
-    recipient: Option<AndrAddr>,
+    recipient: Option<Recipient>,
 ) -> Result<Response, ContractError> {
     let ExecuteContext {
         deps, info, env, ..
@@ -204,15 +211,15 @@ fn execute_claim(
     // Check if the validator is valid before unstaking
     is_validator(&deps, &validator)?;
 
-    let recipient = if let Some(recipient) = recipient {
-        recipient.get_raw_address(&deps.as_ref())?
+    let recipient_address = if let Some(ref recipient) = recipient {
+        recipient.address.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,15 +242,22 @@ 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());
 
+    if let Some(recipient) = recipient {
+        if recipient.msg.is_some() {
+            res.clone()
+                .add_submessage(recipient.generate_direct_msg(&deps.as_ref(), info.funds)?);
+        }
+    };
+
     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::{
@@ -9,7 +9,7 @@ use andromeda_testing::{
 };
 use cw_multi_test::{Contract, ContractWrapper};
 
-use andromeda_std::{amp::AndrAddr, error::ContractError};
+use andromeda_std::{amp::Recipient, error::ContractError};
 
 pub struct MockValidatorStaking(Addr);
 mock_ado!(MockValidatorStaking, ExecuteMsg, QueryMsg);
@@ -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, &[])
     }
 
@@ -41,7 +42,7 @@ impl MockValidatorStaking {
         app: &mut MockApp,
         sender: Addr,
         validator: Option<Addr>,
-        recipient: Option<AndrAddr>,
+        recipient: Option<Recipient>,
     ) -> ExecuteResult {
         let msg = mock_execute_claim_reward(validator, recipient);
         self.execute(app, &msg, sender, &[])
@@ -94,13 +95,13 @@ 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(
     validator: Option<Addr>,
-    recipient: Option<AndrAddr>,
+    recipient: Option<Recipient>,
 ) -> ExecuteMsg {
     ExecuteMsg::Claim {
         validator,

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

@@ -4,7 +4,9 @@ use crate::{
 };
 
 use andromeda_std::{
-    amp::AndrAddr, error::ContractError, testing::mock_querier::MOCK_KERNEL_CONTRACT,
+    amp::{AndrAddr, Recipient},
+    error::ContractError,
+    testing::mock_querier::MOCK_KERNEL_CONTRACT,
 };
 use cosmwasm_std::{
     coin,
@@ -143,6 +145,7 @@ fn test_unauthorized_unstake() {
 
     let msg = ExecuteMsg::Unstake {
         validator: Some(valid_validator),
+        amount: None,
     };
 
     let info = mock_info("other", &[coin(100, "uandr")]);
@@ -167,7 +170,7 @@ fn test_unauthorized_claim() {
 
     let msg = ExecuteMsg::Claim {
         validator: Some(valid_validator.clone()),
-        recipient: Some(AndrAddr::from_string("other")),
+        recipient: Some(Recipient::new(AndrAddr::from_string("other"), None)),
     };
 
     let info = mock_info(OWNER, &[coin(100, "uandr")]);

packages/andromeda-finance/src/validator_staking.rs

@@ -1,6 +1,8 @@
-use andromeda_std::{amp::AndrAddr, andr_exec, andr_instantiate, andr_query, error::ContractError};
+use andromeda_std::{
+    amp::Recipient, 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,10 +20,11 @@ pub enum ExecuteMsg {
     },
     Unstake {
         validator: Option<Addr>,
+        amount: Option<Uint128>,
     },
     Claim {
         validator: Option<Addr>,
-        recipient: Option<AndrAddr>,
+        recipient: Option<Recipient>,
     },
     WithdrawFunds {},
 }