From d94def6c6df48dd61e285e9890e45260285c7438 Mon Sep 17 00:00:00 2001 From: Noah Saso Date: Mon, 24 Jun 2024 20:14:02 -0400 Subject: [PATCH] support custom stargate query in tests --- .../src/tests/multitest/app.rs | 79 +++++++++++++++++++ .../src/tests/multitest/mod.rs | 2 + .../src/tests/multitest/stargate.rs | 64 +++++++++++++++ .../src/tests/multitest/tests.rs | 12 +-- 4 files changed, 152 insertions(+), 5 deletions(-) create mode 100644 contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/app.rs create mode 100644 contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/stargate.rs diff --git a/contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/app.rs b/contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/app.rs new file mode 100644 index 000000000..94f4e0561 --- /dev/null +++ b/contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/app.rs @@ -0,0 +1,79 @@ +use std::ops::{Deref, DerefMut}; + +use crate::tests::multitest::stargate::StargateKeeper; +use cosmwasm_std::{ + testing::MockApi, Api, Empty, GovMsg, IbcMsg, IbcQuery, MemoryStorage, Storage, +}; +use cw_multi_test::{ + no_init, App, AppBuilder, BankKeeper, DistributionKeeper, FailingModule, GovFailingModule, + IbcFailingModule, Router, StakeKeeper, WasmKeeper, +}; +#[allow(clippy::type_complexity)] +pub struct CustomApp( + App< + BankKeeper, + MockApi, + MemoryStorage, + FailingModule, + WasmKeeper, + StakeKeeper, + DistributionKeeper, + FailingModule, + FailingModule, + StargateKeeper, + >, +); +impl Deref for CustomApp { + type Target = App< + BankKeeper, + MockApi, + MemoryStorage, + FailingModule, + WasmKeeper, + StakeKeeper, + DistributionKeeper, + FailingModule, + FailingModule, + StargateKeeper, + >; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl DerefMut for CustomApp { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.0 + } +} +impl Default for CustomApp { + fn default() -> Self { + Self::new(no_init) + } +} + +impl CustomApp { + pub fn new(init_fn: F) -> Self + where + F: FnOnce( + &mut Router< + BankKeeper, + FailingModule, + WasmKeeper, + StakeKeeper, + DistributionKeeper, + IbcFailingModule, + GovFailingModule, + StargateKeeper, + >, + &dyn Api, + &mut dyn Storage, + ), + { + let app_builder = AppBuilder::default(); + let stargate = StargateKeeper {}; + let app = app_builder.with_stargate(stargate).build(init_fn); + CustomApp(app) + } +} diff --git a/contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/mod.rs b/contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/mod.rs index 14f00389d..3735cbfa0 100644 --- a/contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/mod.rs +++ b/contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/mod.rs @@ -1 +1,3 @@ +mod app; mod tests; +mod stargate; diff --git a/contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/stargate.rs b/contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/stargate.rs new file mode 100644 index 000000000..1428d527a --- /dev/null +++ b/contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/stargate.rs @@ -0,0 +1,64 @@ +use cosmwasm_schema::cw_serde; +use cosmwasm_std::{ + from_json, storage_keys::to_length_prefixed, to_json_binary, Addr, Api, Binary, BlockInfo, + Decimal, Querier, Storage, +}; +use cw_multi_test::{error::AnyResult, AppResponse, CosmosRouter, Stargate}; +use cw_storage_plus::Map; +use osmosis_std::types::cosmos::staking::v1beta1::{Pool, QueryPoolResponse}; + +use super::tests::{DELEGATOR, VALIDATOR}; + +// from StakingKeeper +#[cw_serde] +struct Shares { + stake: Decimal, + rewards: Decimal, +} +const NAMESPACE_STAKING: &[u8] = b"staking"; +const STAKES: Map<(&Addr, &Addr), Shares> = Map::new("stakes"); +pub struct StargateKeeper {} + +impl StargateKeeper {} + +impl Stargate for StargateKeeper { + fn execute( + &self, + _api: &dyn Api, + _storage: &mut dyn Storage, + _router: &dyn CosmosRouter, + _block: &BlockInfo, + _sender: Addr, + _type_url: String, + _value: Binary, + ) -> AnyResult { + Ok(AppResponse::default()) + } + + fn query( + &self, + _api: &dyn Api, + storage: &dyn Storage, + _querier: &dyn Querier, + _block: &BlockInfo, + path: String, + data: Binary, + ) -> AnyResult { + if path == *"/cosmos.staking.v1beta1.Query/Pool" { + // since we can't access the app or staking module, let's just raw + // access the storage key we know is set for staking in the test :D + let mut key = to_length_prefixed(NAMESPACE_STAKING); + let map_key = STAKES.key((&Addr::unchecked(DELEGATOR), &Addr::unchecked(VALIDATOR))); + key.extend_from_slice(&map_key); + let data: Shares = from_json(storage.get(&key).unwrap())?; + + return Ok(to_json_binary(&QueryPoolResponse { + pool: Some(Pool { + not_bonded_tokens: "0".to_string(), + bonded_tokens: data.stake.to_string(), + }), + })?); + } + Ok(data) + } +} diff --git a/contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/tests.rs b/contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/tests.rs index 567a40ad0..b7bb212e8 100644 --- a/contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/tests.rs +++ b/contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/tests.rs @@ -1,13 +1,15 @@ use cosmwasm_std::{ coin, coins, testing::mock_env, Addr, CosmosMsg, Decimal, Empty, StakingMsg, Uint128, Validator, }; -use cw_multi_test::{App, Contract, ContractWrapper, Executor}; +use cw_multi_test::{Contract, ContractWrapper, Executor}; use dao_interface::voting::{ InfoResponse, TotalPowerAtHeightResponse, VotingPowerAtHeightResponse, }; use crate::msg::{InstantiateMsg, QueryMsg}; +use super::app::CustomApp; + fn cosmos_staked_contract() -> Box> { let contract = ContractWrapper::new( crate::contract::execute, @@ -18,13 +20,13 @@ fn cosmos_staked_contract() -> Box> { } const DAO: &str = "dao"; -const DELEGATOR: &str = "delegator"; -const VALIDATOR: &str = "validator"; +pub const DELEGATOR: &str = "delegator"; +pub const VALIDATOR: &str = "validator"; const VALIDATOR2: &str = "validator2"; const DENOM: &str = "TOKEN"; -fn setup_test_env() -> App { - App::new(|router, api, storage| { +fn setup_test_env() -> CustomApp { + CustomApp::new(|router, api, storage| { router .staking .add_validator(