Add DelegationRewards query

CosmWasm · Jul 21, 2023 · 981cb46 · 981cb46
1 parent 5d07829
commit 981cb46
Show file tree

Hide file tree

Showing 5 changed files with 87 additions and 2 deletions.
diff --git a/packages/std/src/lib.rs b/packages/std/src/lib.rs
@@ -73,6 +73,8 @@ 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};
 #[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
@@ -10,7 +10,15 @@ use super::query_response::QueryResponseType;
 #[serde(rename_all = "snake_case")]
 pub enum DistributionQuery {
     // https://github.com/cosmos/cosmos-sdk/blob/4f6f6c00021f4b5ee486bbb71ae2071a8ceb47c9/x/distribution/types/query.pb.go#L792-L795
-    DelegatorWithdrawAddress { delegator_address: String },
+    DelegatorWithdrawAddress {
+        delegator_address: String,
+    },
+    // https://github.com/cosmos/cosmos-sdk/blob/e3482f2d4142c55f9dc3f47a321b56610a11492c/x/distribution/types/query.pb.go#L525-L532
+    #[cfg(feature = "cosmwasm_1_4")]
+    DelegationRewards {
+        delegator_address: String,
+        validator_address: String,
+    },
 }
 
 // https://github.com/cosmos/cosmos-sdk/blob/4f6f6c00021f4b5ee486bbb71ae2071a8ceb47c9/x/distribution/types/query.pb.go#L832-L835
@@ -24,3 +32,26 @@ pub struct DelegatorWithdrawAddressResponse {
 impl_response_constructor!(DelegatorWithdrawAddressResponse, withdraw_address: Addr);
 
 impl QueryResponseType for DelegatorWithdrawAddressResponse {}
+
+// https://github.com/cosmos/cosmos-sdk/blob/e3482f2d4142c55f9dc3f47a321b56610a11492c/x/distribution/types/query.pb.go#L567-L572
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
+#[serde(rename_all = "snake_case")]
+#[cfg(feature = "cosmwasm_1_4")]
+pub struct DelegationRewardsResponse {
+    pub rewards: Vec<DecCoin>,
+}
+
+#[cfg(feature = "cosmwasm_1_4")]
+impl_response_constructor!(DelegationRewardsResponse, rewards: Vec<DecCoin>);
+#[cfg(feature = "cosmwasm_1_4")]
+impl QueryResponseType for DelegationRewardsResponse {}
+
+/// A coin type with decimal amount.
+/// Modeled after the Cosmos SDK's `DecCoin` type
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
+#[serde(rename_all = "snake_case")]
+#[cfg(feature = "cosmwasm_1_4")]
+pub struct DecCoin {
+    pub denom: String,
+    pub amount: crate::Decimal,
+}
diff --git a/packages/std/src/query/mod.rs b/packages/std/src/query/mod.rs
@@ -36,6 +36,8 @@ pub use bank::SupplyResponse;
 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};
 #[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
@@ -44,6 +44,8 @@ use crate::{
 use crate::{Attribute, DenomMetadata};
 #[cfg(feature = "stargate")]
 use crate::{ChannelResponse, IbcQuery, ListChannelsResponse, PortIdResponse};
+#[cfg(feature = "cosmwasm_1_4")]
+use crate::{DecCoin, DelegationRewardsResponse};
 
 use super::riffle_shuffle;
 
@@ -934,12 +936,17 @@ impl StakingQuerier {
 #[derive(Clone, Default)]
 pub struct DistributionQuerier {
     withdraw_addresses: HashMap<String, String>,
+    #[cfg(feature = "cosmwasm_1_4")]
+    rewards: BTreeMap<(String, String), Vec<crate::DecCoin>>,
 }
 
 #[cfg(feature = "cosmwasm_1_3")]
 impl DistributionQuerier {
     pub fn new(withdraw_addresses: HashMap<String, String>) -> Self {
-        DistributionQuerier { withdraw_addresses }
+        DistributionQuerier {
+            withdraw_addresses,
+            ..Default::default()
+        }
     }
 
     pub fn set_withdraw_address(
@@ -968,6 +975,17 @@ impl DistributionQuerier {
         self.withdraw_addresses.clear();
     }
 
+    /// Sets accumulated rewards for a given validator and delegator pair.
+    pub fn set_rewards(
+        &mut self,
+        validator: impl Into<String>,
+        delegator: impl Into<String>,
+        rewards: Vec<DecCoin>,
+    ) {
+        self.rewards
+            .insert((validator.into(), delegator.into()), rewards);
+    }
+
     pub fn query(&self, request: &DistributionQuery) -> QuerierResult {
         let contract_result: ContractResult<Binary> = match request {
             DistributionQuery::DelegatorWithdrawAddress { delegator_address } => {
@@ -980,6 +998,20 @@ impl DistributionQuerier {
                 };
                 to_binary(&res).into()
             }
+            #[cfg(feature = "cosmwasm_1_4")]
+            DistributionQuery::DelegationRewards {
+                delegator_address,
+                validator_address,
+            } => {
+                let res = DelegationRewardsResponse {
+                    rewards: self
+                        .rewards
+                        .get(&(validator_address.clone(), delegator_address.clone()))
+                        .cloned()
+                        .unwrap_or_default(),
+                };
+                to_binary(&res).into()
+            }
         };
         // system result is always ok in the mock implementation
         SystemResult::Ok(contract_result)

diff --git a/packages/std/src/traits.rs b/packages/std/src/traits.rs
@@ -313,6 +313,24 @@ impl<'a, C: CustomQuery> QuerierWrapper<'a, C> {
         self.query(&request)
     }
 
+    #[cfg(feature = "cosmwasm_1_4")]
+    pub fn query_delegation_rewards(
+        &self,
+        delegator: impl Into<String>,
+        validator: impl Into<String>,
+    ) -> StdResult<Vec<crate::DecCoin>> {
+        use crate::DelegationRewardsResponse;
+
+        let request = DistributionQuery::DelegationRewards {
+            delegator_address: delegator.into(),
+            validator_address: validator.into(),
+        }
+        .into();
+        let DelegationRewardsResponse { rewards } = self.query(&request)?;
+
+        Ok(rewards)
+    }
+
     // 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>(