support custom stargate query in tests

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<Empty, Empty, Empty>,
+        WasmKeeper<Empty, Empty>,
+        StakeKeeper,
+        DistributionKeeper,
+        FailingModule<IbcMsg, IbcQuery, Empty>,
+        FailingModule<GovMsg, Empty, Empty>,
+        StargateKeeper,
+    >,
+);
+impl Deref for CustomApp {
+    type Target = App<
+        BankKeeper,
+        MockApi,
+        MemoryStorage,
+        FailingModule<Empty, Empty, Empty>,
+        WasmKeeper<Empty, Empty>,
+        StakeKeeper,
+        DistributionKeeper,
+        FailingModule<IbcMsg, IbcQuery, Empty>,
+        FailingModule<GovMsg, Empty, Empty>,
+        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<F>(init_fn: F) -> Self
+    where
+        F: FnOnce(
+            &mut Router<
+                BankKeeper,
+                FailingModule<Empty, Empty, Empty>,
+                WasmKeeper<Empty, Empty>,
+                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)
+    }
+}

contracts/voting/dao-voting-cosmos-staked/src/tests/multitest/mod.rs

@@ -1 +1,3 @@
+mod app;
 mod tests;
+mod stargate;

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<ExecC, QueryC>(
+        &self,
+        _api: &dyn Api,
+        _storage: &mut dyn Storage,
+        _router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
+        _block: &BlockInfo,
+        _sender: Addr,
+        _type_url: String,
+        _value: Binary,
+    ) -> AnyResult<AppResponse> {
+        Ok(AppResponse::default())
+    }
+
+    fn query(
+        &self,
+        _api: &dyn Api,
+        storage: &dyn Storage,
+        _querier: &dyn Querier,
+        _block: &BlockInfo,
+        path: String,
+        data: Binary,
+    ) -> AnyResult<Binary> {
+        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)
+    }
+}

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<dyn Contract<Empty>> {
     let contract = ContractWrapper::new(
         crate::contract::execute,
@@ -18,13 +20,13 @@ fn cosmos_staked_contract() -> Box<dyn Contract<Empty>> {
 }
 
 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(