added tests

contracts/proxy_to_stt/src/testing/mock_querier.rs

@@ -0,0 +1,175 @@
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+
+use astroport_generator_proxy::stt_staking::StakerInfoResponse;
+use cosmwasm_std::testing::{MockApi, MockQuerier, MockStorage, MOCK_CONTRACT_ADDR};
+use cosmwasm_std::{
+    from_binary, from_slice, to_binary, Coin, ContractResult, Decimal, Empty, OwnedDeps, Querier,
+    QuerierResult, QueryRequest, SystemError, SystemResult, Uint128, WasmQuery,
+};
+use cw20::{BalanceResponse as Cw20BalanceResponse, Cw20QueryMsg};
+
+/// mock_dependencies is a drop-in replacement for cosmwasm_std::testing::mock_dependencies
+/// this uses our CustomQuerier.
+pub fn mock_dependencies(
+    contract_balance: &[Coin],
+) -> OwnedDeps<MockStorage, MockApi, WasmMockQuerier> {
+    let custom_querier: WasmMockQuerier =
+        WasmMockQuerier::new(MockQuerier::new(&[(MOCK_CONTRACT_ADDR, contract_balance)]));
+
+    OwnedDeps {
+        api: MockApi::default(),
+        storage: MockStorage::default(),
+        querier: custom_querier,
+    }
+}
+
+pub struct WasmMockQuerier {
+    base: MockQuerier<Empty>,
+    token_querier: TokenQuerier,
+    reward_querier: RewardQuerier,
+}
+
+#[derive(Clone, Default)]
+pub struct TokenQuerier {
+    // this lets us iterate over all pairs that match the first string
+    balances: HashMap<String, HashMap<String, Uint128>>,
+}
+
+impl TokenQuerier {
+    pub fn new(balances: &[(&String, &[(&String, &Uint128)])]) -> Self {
+        TokenQuerier {
+            balances: balances_to_map(balances),
+        }
+    }
+}
+
+pub(crate) fn balances_to_map(
+    balances: &[(&String, &[(&String, &Uint128)])],
+) -> HashMap<String, HashMap<String, Uint128>> {
+    let mut balances_map: HashMap<String, HashMap<String, Uint128>> = HashMap::new();
+    for (contract_addr, balances) in balances.iter() {
+        let mut contract_balances_map: HashMap<String, Uint128> = HashMap::new();
+        for (addr, balance) in balances.iter() {
+            contract_balances_map.insert(addr.to_string(), **balance);
+        }
+
+        balances_map.insert(contract_addr.to_string(), contract_balances_map);
+    }
+    balances_map
+}
+
+#[derive(Clone, Default)]
+pub struct RewardQuerier {
+    pending_reward: Uint128,
+    deposit_amount: Uint128,
+}
+
+impl RewardQuerier {
+    pub fn new(pending_reward: Uint128, deposit_amount: Uint128) -> Self {
+        RewardQuerier {
+            pending_reward,
+            deposit_amount,
+        }
+    }
+}
+
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, JsonSchema)]
+#[serde(rename_all = "snake_case")]
+pub enum QueryMsg {
+    StakerInfo {
+        staker: String,
+        block_height: Option<u64>,
+    },
+}
+
+impl Querier for WasmMockQuerier {
+    fn raw_query(&self, bin_request: &[u8]) -> QuerierResult {
+        // MockQuerier doesn't support Custom, so we ignore it completely here
+        let request: QueryRequest<Empty> = match from_slice(bin_request) {
+            Ok(v) => v,
+            Err(e) => {
+                return SystemResult::Err(SystemError::InvalidRequest {
+                    error: format!("Parsing query request: {}", e),
+                    request: bin_request.into(),
+                })
+            }
+        };
+        self.handle_query(&request)
+    }
+}
+
+impl WasmMockQuerier {
+    pub fn handle_query(&self, request: &QueryRequest<Empty>) -> QuerierResult {
+        match &request {
+            QueryRequest::Wasm(WasmQuery::Smart { contract_addr, msg }) => match from_binary(msg) {
+                Ok(QueryMsg::StakerInfo {
+                    staker: _,
+                    block_height: _,
+                }) => SystemResult::Ok(ContractResult::from(to_binary(&StakerInfoResponse {
+                    staker: "generator0000".to_string(),
+                    reward_index: Decimal::zero(),
+                    bond_amount: self.reward_querier.deposit_amount,
+                    pending_reward: self.reward_querier.pending_reward,
+                    rewards_per_fee: vec![],
+                    time_to_best_fee: None,
+                    pending_unbond_left: None,
+                }))),
+                _ => match from_binary(msg).unwrap() {
+                    Cw20QueryMsg::Balance { address } => {
+                        let balances: &HashMap<String, Uint128> =
+                            match self.token_querier.balances.get(contract_addr) {
+                                Some(balances) => balances,
+                                None => {
+                                    return SystemResult::Err(SystemError::InvalidRequest {
+                                        error: format!(
+                                            "No balance info exists for the contract {}",
+                                            contract_addr
+                                        ),
+                                        request: msg.as_slice().into(),
+                                    })
+                                }
+                            };
+
+                        let balance = match balances.get(&address) {
+                            Some(v) => *v,
+                            None => {
+                                return SystemResult::Ok(ContractResult::Ok(
+                                    to_binary(&Cw20BalanceResponse {
+                                        balance: Uint128::zero(),
+                                    })
+                                    .unwrap(),
+                                ));
+                            }
+                        };
+
+                        SystemResult::Ok(ContractResult::Ok(
+                            to_binary(&Cw20BalanceResponse { balance }).unwrap(),
+                        ))
+                    }
+                    _ => panic!("Query Not Mocked"),
+                },
+            },
+            _ => self.base.handle_query(request),
+        }
+    }
+}
+
+impl WasmMockQuerier {
+    pub fn new(base: MockQuerier<Empty>) -> Self {
+        WasmMockQuerier {
+            base,
+            token_querier: TokenQuerier::default(),
+            reward_querier: RewardQuerier::default(),
+        }
+    }
+
+    pub fn with_token_balances(&mut self, balances: &[(&String, &[(&String, &Uint128)])]) {
+        self.token_querier = TokenQuerier::new(balances);
+    }
+
+    pub fn with_reward_info(&mut self, pending_reward: Uint128, deposit_amount: Uint128) {
+        self.reward_querier = RewardQuerier::new(pending_reward, deposit_amount);
+    }
+}

contracts/proxy_to_stt/src/testing/mod.rs

@@ -0,0 +1,2 @@
+mod mock_querier;
+mod tests;