1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

use crate::{AttributeQuerier, MarkerQuerier, MetadataQuerier, NameQuerier};
use cosmwasm_std::testing::{MockApi, MockQuerier, MockStorage, MOCK_CONTRACT_ADDR};
use cosmwasm_std::{
    from_slice, to_binary, Coin, OwnedDeps, Querier, QuerierResult, QueryRequest, StdResult,
    SystemError, SystemResult,
};
use provwasm_std::{Marker, ProvenanceQuery, ProvenanceQueryParams, Records, Scope, Sessions};
use std::marker::PhantomData;

/// A drop-in replacement for cosmwasm_std::testing::mock_dependencies that uses the mock
/// provenance querier.
pub fn mock_dependencies(
    contract_balance: &[Coin],
) -> OwnedDeps<MockStorage, MockApi, ProvenanceMockQuerier, ProvenanceQuery> {
    let base = MockQuerier::new(&[(MOCK_CONTRACT_ADDR, contract_balance)]);
    OwnedDeps {
        storage: MockStorage::default(),
        api: MockApi::default(),
        querier: ProvenanceMockQuerier::new(base),
        custom_query_type: PhantomData,
    }
}

/// Initializes the mock querier with the account balances provided. NOTE: contract balance must
/// be set in the balances slice passed if required.
pub fn mock_dependencies_with_balances(
    balances: &[(&str, &[Coin])],
) -> OwnedDeps<MockStorage, MockApi, ProvenanceMockQuerier, ProvenanceQuery> {
    let base = MockQuerier::new(balances);
    OwnedDeps {
        storage: MockStorage::default(),
        api: MockApi::default(),
        querier: ProvenanceMockQuerier::new(base),
        custom_query_type: PhantomData,
    }
}

pub struct ProvenanceMockQuerier {
    pub base: MockQuerier,
    name: NameQuerier,
    attribute: AttributeQuerier,
    marker: MarkerQuerier,
    metadata: MetadataQuerier,
}

impl Querier for ProvenanceMockQuerier {
    fn raw_query(&self, bin_request: &[u8]) -> QuerierResult {
        let request: StdResult<QueryRequest<ProvenanceQuery>> = from_slice(bin_request);
        match &request {
            Ok(provenance_query) => self.handle_query(provenance_query),
            _ => self.base.raw_query(bin_request),
        }
    }
}

impl ProvenanceMockQuerier {
    pub fn handle_query(&self, request: &QueryRequest<ProvenanceQuery>) -> QuerierResult {
        match request {
            QueryRequest::Custom(custom) => match &custom.params {
                ProvenanceQueryParams::Attribute(p) => self.attribute.query(p),
                ProvenanceQueryParams::Marker(p) => self.marker.query(p),
                ProvenanceQueryParams::Name(p) => self.name.query(p),
                ProvenanceQueryParams::Metadata(p) => self.metadata.query(p),
            },
            QueryRequest::Bank(q) => self.base.handle_query(&QueryRequest::Bank(q.clone())),
            #[cfg(feature = "staking")]
            QueryRequest::Staking(q) => self.base.handle_query(&QueryRequest::Staking(q.clone())),
            QueryRequest::Wasm(q) => self.base.handle_query(&QueryRequest::Wasm(q.clone())),
            // Note: As of 0.14, no mocks exist for stargate or ibc queries in base, so we just
            // return an error.
            _ => SystemResult::Err(SystemError::InvalidRequest {
                error: "invalid query request type".into(),
                request: to_binary(&request).unwrap(),
            }),
        }
    }
}

impl ProvenanceMockQuerier {
    pub fn new(base: MockQuerier) -> Self {
        ProvenanceMockQuerier {
            base,
            attribute: AttributeQuerier::default(),
            name: NameQuerier::default(),
            marker: MarkerQuerier::default(),
            metadata: MetadataQuerier::default(),
        }
    }

    pub fn with_names(&mut self, inputs: &[(&str, &str, bool)]) {
        self.name = NameQuerier::new(inputs);
    }

    pub fn with_attributes(&mut self, address: &str, inputs: &[(&str, &str, &str)]) {
        self.attribute = AttributeQuerier::new(address, inputs);
    }

    pub fn with_markers(&mut self, inputs: Vec<Marker>) {
        self.marker = MarkerQuerier::new(inputs);
    }

    pub fn with_scope(&mut self, scope: Scope) {
        self.metadata = MetadataQuerier::new(scope, None, None)
    }

    pub fn with_sessions(&mut self, scope: Scope, sessions: Sessions) {
        self.metadata = MetadataQuerier::new(scope, Some(sessions), None)
    }

    pub fn with_records(&mut self, scope: Scope, records: Records) {
        self.metadata = MetadataQuerier::new(scope, None, Some(records))
    }

    pub fn with_metadata(&mut self, scope: Scope, sessions: Sessions, records: Records) {
        self.metadata = MetadataQuerier::new(scope, Some(sessions), Some(records))
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use cosmwasm_std::{coin, from_binary, BalanceResponse, BankQuery};

    #[test]
    fn query_with_balances() {
        let amount = coin(12345, "denom");
        let deps = mock_dependencies_with_balances(&[("alice", &[amount.clone()])]);
        let bin = deps
            .querier
            .handle_query(
                &BankQuery::Balance {
                    address: "alice".into(),
                    denom: "denom".into(),
                }
                .into(),
            )
            .unwrap()
            .unwrap();

        let res: BalanceResponse = from_binary(&bin).unwrap();
        assert_eq!(res.amount, amount);
    }
}