Add DelegationTotalRewards query

CosmWasm · Jul 21, 2023 · 46472c6 · 46472c6
1 parent 981cb46
commit 46472c6
Show file tree

Hide file tree

Showing 5 changed files with 99 additions and 6 deletions.
diff --git a/packages/std/src/lib.rs b/packages/std/src/lib.rs
@@ -74,7 +74,9 @@ pub use crate::query::{
 #[cfg(feature = "stargate")]
 pub use crate::query::{ChannelResponse, IbcQuery, ListChannelsResponse, PortIdResponse};
 #[cfg(feature = "cosmwasm_1_4")]
-pub use crate::query::{DecCoin, DelegationRewardsResponse};
+pub use crate::query::{
+    DecCoin, DelegationRewardsResponse, DelegationTotalRewardsResponse, DelegatorReward,
+};
 #[allow(deprecated)]
 pub use crate::results::SubMsgExecutionResponse;
 #[cfg(all(feature = "stargate", feature = "cosmwasm_1_2"))]

diff --git a/packages/std/src/query/distribution.rs b/packages/std/src/query/distribution.rs
@@ -19,6 +19,11 @@ pub enum DistributionQuery {
         delegator_address: String,
         validator_address: String,
     },
+    // https://github.com/cosmos/cosmos-sdk/blob/e3482f2d4142c55f9dc3f47a321b56610a11492c/x/distribution/types/query.pb.go#L614-L619
+    #[cfg(feature = "cosmwasm_1_4")]
+    DelegationTotalRewards {
+        delegator_address: String,
+    },
 }
 
 // https://github.com/cosmos/cosmos-sdk/blob/4f6f6c00021f4b5ee486bbb71ae2071a8ceb47c9/x/distribution/types/query.pb.go#L832-L835
@@ -55,3 +60,28 @@ pub struct DecCoin {
     pub denom: String,
     pub amount: crate::Decimal,
 }
+
+// https://github.com/cosmos/cosmos-sdk/blob/e3482f2d4142c55f9dc3f47a321b56610a11492c/x/distribution/types/query.pb.go#L654-L661
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
+#[non_exhaustive]
+#[cfg(feature = "cosmwasm_1_4")]
+pub struct DelegationTotalRewardsResponse {
+    pub rewards: Vec<DelegatorReward>,
+    pub total: Vec<DecCoin>,
+}
+
+#[cfg(feature = "cosmwasm_1_4")]
+impl_response_constructor!(
+    DelegationTotalRewardsResponse,
+    rewards: Vec<DelegatorReward>,
+    total: Vec<DecCoin>
+);
+#[cfg(feature = "cosmwasm_1_4")]
+impl QueryResponseType for DelegationTotalRewardsResponse {}
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
+#[cfg(feature = "cosmwasm_1_4")]
+pub struct DelegatorReward {
+    pub validator_address: String,
+    pub reward: Vec<DecCoin>,
+}
diff --git a/packages/std/src/query/mod.rs b/packages/std/src/query/mod.rs
@@ -37,7 +37,9 @@ pub use bank::{AllBalanceResponse, BalanceResponse, BankQuery};
 #[cfg(feature = "cosmwasm_1_3")]
 pub use bank::{AllDenomMetadataResponse, DenomMetadataResponse};
 #[cfg(feature = "cosmwasm_1_4")]
-pub use distribution::{DecCoin, DelegationRewardsResponse};
+pub use distribution::{
+    DecCoin, DelegationRewardsResponse, DelegationTotalRewardsResponse, DelegatorReward,
+};
 #[cfg(feature = "cosmwasm_1_3")]
 pub use distribution::{DelegatorWithdrawAddressResponse, DistributionQuery};
 #[cfg(feature = "stargate")]

diff --git a/packages/std/src/testing/mock.rs b/packages/std/src/testing/mock.rs
@@ -45,7 +45,7 @@ use crate::{Attribute, DenomMetadata};
 #[cfg(feature = "stargate")]
 use crate::{ChannelResponse, IbcQuery, ListChannelsResponse, PortIdResponse};
 #[cfg(feature = "cosmwasm_1_4")]
-use crate::{DecCoin, DelegationRewardsResponse};
+use crate::{DecCoin, Decimal, DelegationRewardsResponse};
 
 use super::riffle_shuffle;
 
@@ -937,7 +937,8 @@ impl StakingQuerier {
 pub struct DistributionQuerier {
     withdraw_addresses: HashMap<String, String>,
     #[cfg(feature = "cosmwasm_1_4")]
-    rewards: BTreeMap<(String, String), Vec<crate::DecCoin>>,
+    /// Mock of accumulated rewards, indexed first by delegator and then validator address.
+    rewards: BTreeMap<String, BTreeMap<String, Vec<crate::DecCoin>>>,
 }
 
 #[cfg(feature = "cosmwasm_1_3")]
@@ -983,7 +984,9 @@ impl DistributionQuerier {
         rewards: Vec<DecCoin>,
     ) {
         self.rewards
-            .insert((validator.into(), delegator.into()), rewards);
+            .entry(delegator.into())
+            .or_default()
+            .insert(validator.into(), rewards);
     }
 
     pub fn query(&self, request: &DistributionQuery) -> QuerierResult {
@@ -1006,16 +1009,59 @@ impl DistributionQuerier {
                 let res = DelegationRewardsResponse {
                     rewards: self
                         .rewards
-                        .get(&(validator_address.clone(), delegator_address.clone()))
+                        .get(delegator_address)
+                        .and_then(|v| v.get(validator_address))
                         .cloned()
                         .unwrap_or_default(),
                 };
                 to_binary(&res).into()
             }
+            #[cfg(feature = "cosmwasm_1_4")]
+            DistributionQuery::DelegationTotalRewards { delegator_address } => {
+                let validator_rewards = self
+                    .validator_rewards(delegator_address)
+                    .unwrap_or_default();
+                let res = crate::DelegationTotalRewardsResponse {
+                    total: validator_rewards
+                        .iter()
+                        .fold(BTreeMap::<&str, DecCoin>::new(), |mut acc, rewards| {
+                            for coin in &rewards.reward {
+                                acc.entry(&coin.denom)
+                                    .or_insert_with(|| DecCoin {
+                                        denom: coin.denom.clone(),
+                                        amount: Decimal::zero(),
+                                    })
+                                    .amount += coin.amount;
+                            }
+
+                            acc
+                        })
+                        .into_values()
+                        .collect(),
+                    rewards: validator_rewards,
+                };
+                to_binary(&res).into()
+            }
         };
         // system result is always ok in the mock implementation
         SystemResult::Ok(contract_result)
     }
+
+    /// Helper method to get all rewards for a given delegator.
+    #[cfg(feature = "cosmwasm_1_4")]
+    fn validator_rewards(&self, delegator_address: &str) -> Option<Vec<crate::DelegatorReward>> {
+        let validator_rewards = self.rewards.get(delegator_address)?;
+
+        Some(
+            validator_rewards
+                .iter()
+                .map(|(validator, rewards)| crate::DelegatorReward {
+                    validator_address: validator.clone(),
+                    reward: rewards.clone(),
+                })
+                .collect(),
+        )
+    }
 }
 
 pub fn digit_sum(input: &[u8]) -> usize {

diff --git a/packages/std/src/traits.rs b/packages/std/src/traits.rs
@@ -331,6 +331,19 @@ impl<'a, C: CustomQuery> QuerierWrapper<'a, C> {
         Ok(rewards)
     }
 
+    #[cfg(feature = "cosmwasm_1_4")]
+    pub fn query_delegation_total_rewards(
+        &self,
+        delegator: impl Into<String>,
+        validator: impl Into<String>,
+    ) -> StdResult<crate::DelegationTotalRewardsResponse> {
+        let request = DistributionQuery::DelegationTotalRewards {
+            delegator_address: delegator.into(),
+        }
+        .into();
+        self.query(&request)
+    }
+
     // this queries another wasm contract. You should know a priori the proper types for T and U
     // (response and request) based on the contract API
     pub fn query_wasm_smart<T: DeserializeOwned>(