use std::collections::HashMap;
use crate::wasm_emulation::channel::RemoteChannel;
use crate::wasm_emulation::query::bank::BankQuerier;
use crate::wasm_emulation::query::staking::StakingQuerier;
use crate::wasm_emulation::query::wasm::WasmQuerier;
use cosmwasm_std::CustomMsg;
use cosmwasm_std::Env;
use cosmwasm_vm::BackendResult;
use cosmwasm_vm::GasInfo;
use serde::de::DeserializeOwned;
use cosmwasm_std::Binary;
use cosmwasm_std::Coin;
use cosmwasm_std::SystemError;
use cosmwasm_std::from_json;
use cosmwasm_std::{ContractResult, SystemResult};
use cosmwasm_std::{CustomQuery, QueryRequest};
use cosmwasm_std::{FullDelegation, Validator};
use cosmwasm_std::Attribute;
use cosmwasm_std::QuerierResult;
use crate::wasm_emulation::input::QuerierStorage;
use crate::Contract;
use super::gas::GAS_COST_QUERY_ERROR;
#[derive(Clone)]
pub struct LocalForkedState<ExecC, QueryC> {
pub contracts: HashMap<usize, *mut dyn Contract<ExecC, QueryC>>,
pub env: Env,
}
#[derive(Clone)]
pub struct ForkState<ExecC, QueryC>
where
QueryC: CustomQuery + DeserializeOwned + 'static,
ExecC: CustomMsg + 'static,
{
pub remote: RemoteChannel,
pub local_state: LocalForkedState<ExecC, QueryC>,
pub querier_storage: QuerierStorage,
}
pub type QueryResultWithGas = (QuerierResult, GasInfo);
pub type MockQuerierCustomHandlerResult = SystemResult<ContractResult<Binary>>;
pub struct MockQuerier<
ExecC: CustomMsg + DeserializeOwned + 'static,
QueryC: CustomQuery + DeserializeOwned + 'static,
> {
bank: BankQuerier,
staking: StakingQuerier,
wasm: WasmQuerier<ExecC, QueryC>,
custom_handler: Box<dyn for<'a> Fn(&'a QueryC) -> QueryResultWithGas>,
remote: RemoteChannel,
}
impl<
ExecC: CustomMsg + DeserializeOwned + 'static,
QueryC: CustomQuery + DeserializeOwned + 'static,
> MockQuerier<ExecC, QueryC>
{
pub fn new(fork_state: ForkState<ExecC, QueryC>) -> Self {
MockQuerier {
bank: BankQuerier::new(
fork_state.remote.clone(),
fork_state.querier_storage.bank.storage.clone(),
),
staking: StakingQuerier::default(),
wasm: WasmQuerier::new(fork_state.clone()),
custom_handler: Box::from(|_: &_| -> QueryResultWithGas {
(
SystemResult::Err(SystemError::UnsupportedRequest {
kind: "custom".to_string(),
}),
GasInfo::free(),
)
}),
remote: fork_state.remote,
}
}
pub fn update_balance(
&mut self,
addr: impl Into<String>,
balance: Vec<Coin>,
) -> Option<Vec<Coin>> {
self.bank.update_balance(addr, balance)
}
pub fn update_staking(
&mut self,
denom: &str,
validators: &[Validator],
delegations: &[FullDelegation],
) {
self.staking = StakingQuerier::new(denom, validators, delegations);
}
pub fn with_custom_handler<CH>(mut self, handler: CH) -> Self
where
CH: Fn(&QueryC) -> QueryResultWithGas + 'static,
{
self.custom_handler = Box::from(handler);
self
}
}
impl<
ExecC: CustomMsg + DeserializeOwned + 'static,
QueryC: CustomQuery + DeserializeOwned + 'static,
> cosmwasm_vm::Querier for MockQuerier<ExecC, QueryC>
{
fn query_raw(
&self,
bin_request: &[u8],
_gas_limit: u64,
) -> BackendResult<SystemResult<ContractResult<Binary>>> {
let request: QueryRequest<QueryC> = match from_json(bin_request) {
Ok(v) => v,
Err(e) => {
return (
Ok(SystemResult::Err(SystemError::InvalidRequest {
error: format!("Parsing query request: {}", e),
request: bin_request.into(),
})),
GasInfo::with_externally_used(GAS_COST_QUERY_ERROR),
)
}
};
let result = self.handle_query(&request);
(Ok(result.0), result.1)
}
}
impl<
ExecC: CustomMsg + DeserializeOwned + 'static,
QueryC: CustomQuery + DeserializeOwned + 'static,
> cosmwasm_std::Querier for MockQuerier<ExecC, QueryC>
{
fn raw_query(&self, bin_request: &[u8]) -> SystemResult<ContractResult<Binary>> {
let request: QueryRequest<QueryC> = match from_json(bin_request) {
Ok(v) => v,
Err(e) => {
return SystemResult::Err(SystemError::InvalidRequest {
error: format!("Parsing query request: {}", e),
request: bin_request.into(),
})
}
};
let result = self.handle_query(&request);
result.0
}
}
impl<
ExecC: CustomMsg + DeserializeOwned + 'static,
QueryC: CustomQuery + DeserializeOwned + 'static,
> MockQuerier<ExecC, QueryC>
{
pub fn handle_query(&self, request: &QueryRequest<QueryC>) -> QueryResultWithGas {
match &request {
QueryRequest::Bank(bank_query) => self.bank.query(bank_query),
QueryRequest::Custom(custom_query) => (*self.custom_handler)(custom_query),
QueryRequest::Staking(staking_query) => self.staking.query(staking_query),
QueryRequest::Wasm(msg) => self.wasm.query(self.remote.clone(), msg),
#[allow(deprecated)]
QueryRequest::Stargate { .. } => (
SystemResult::Err(SystemError::UnsupportedRequest {
kind: "Stargate".to_string(),
}),
GasInfo::with_externally_used(GAS_COST_QUERY_ERROR),
),
QueryRequest::Grpc(_req) => (
SystemResult::Err(SystemError::UnsupportedRequest {
kind: "Stargate".to_string(),
}),
GasInfo::with_externally_used(GAS_COST_QUERY_ERROR),
),
&_ => panic!("Query Type Not implemented"),
}
}
}
pub fn digit_sum(input: &[u8]) -> usize {
input.iter().fold(0, |sum, val| sum + (*val as usize))
}
pub fn mock_wasmd_attr(key: impl Into<String>, value: impl Into<String>) -> Attribute {
Attribute {
key: key.into(),
value: value.into(),
}
}