abstract_testing/
mock_querier.rs

1use std::{collections::HashMap, ops::Deref};
2
3use abstract_std::{
4    native_addrs,
5    objects::{
6        gov_type::GovernanceDetails, ownership::Ownership,
7        storage_namespaces::OWNERSHIP_STORAGE_KEY,
8    },
9};
10use cosmwasm_std::{
11    testing::MockApi, Addr, Binary, CodeInfoResponse, ContractInfoResponse, ContractResult, Empty,
12    QuerierWrapper, SystemResult, WasmQuery,
13};
14use cw2::{ContractVersion, CONTRACT};
15use cw_storage_plus::{Item, Map, PrimaryKey};
16use serde::{de::DeserializeOwned, Serialize};
17
18use crate::prelude::*;
19
20type BinaryQueryResult = Result<Binary, String>;
21type FallbackHandler = dyn for<'a> Fn(&'a Addr, &'a Binary) -> BinaryQueryResult;
22type SmartHandler = dyn for<'a> Fn(&'a Binary) -> BinaryQueryResult;
23type RawHandler = dyn for<'a> Fn(&'a str) -> BinaryQueryResult;
24
25/// [`MockQuerierBuilder`] is a helper to build a [`MockQuerier`].
26/// Usage:
27///
28/// ```
29/// use cosmwasm_std::{from_json, to_json_binary};
30/// use abstract_testing::MockQuerierBuilder;
31/// use cosmwasm_std::testing::{MockQuerier, MockApi};
32/// use abstract_sdk::mock_module::MockModuleExecuteMsg;
33///
34/// let api = MockApi::default();
35/// let contract_address = api.addr_make("contract_address");
36/// let querier = MockQuerierBuilder::default().with_smart_handler(&contract_address, |msg| {
37///    // handle the message
38///     let res = match from_json::<MockModuleExecuteMsg>(msg).unwrap() {
39///         // handle the message
40///        _ => panic!("unexpected message"),
41///    };
42///
43///   Ok(to_json_binary(&msg).unwrap())
44/// }).build();
45/// ```
46pub struct MockQuerierBuilder {
47    base: MockQuerier,
48    fallback_raw_handler: Box<FallbackHandler>,
49    fallback_smart_handler: Box<FallbackHandler>,
50    smart_handlers: HashMap<Addr, Box<SmartHandler>>,
51    raw_handlers: HashMap<Addr, Box<RawHandler>>,
52    raw_mappings: HashMap<Addr, HashMap<Binary, Binary>>,
53    contract_admin: HashMap<Addr, Addr>,
54    // Used for Address generation
55    pub api: MockApi,
56}
57
58impl Default for MockQuerierBuilder {
59    /// Create a default
60    fn default() -> Self {
61        Self::new(MockApi::default())
62    }
63}
64
65impl MockQuerierBuilder {
66    pub fn new(api: MockApi) -> Self {
67        let raw_fallback: fn(&Addr, &Binary) -> BinaryQueryResult = |addr, key| {
68            let str_key = std::str::from_utf8(key.as_slice()).unwrap();
69            Err(format!(
70                "No raw query handler for {addr:?} with key {str_key:?}"
71            ))
72        };
73        let smart_fallback: fn(&Addr, &Binary) -> BinaryQueryResult = |addr, key| {
74            let str_key = std::str::from_utf8(key.as_slice()).unwrap();
75            Err(format!(
76                "unexpected smart-query on contract: {addr:?} {str_key:?}"
77            ))
78        };
79
80        Self {
81            base: MockQuerier::default(),
82            fallback_raw_handler: Box::from(raw_fallback),
83            fallback_smart_handler: Box::from(smart_fallback),
84            smart_handlers: HashMap::default(),
85            raw_handlers: HashMap::default(),
86            raw_mappings: HashMap::default(),
87            contract_admin: HashMap::default(),
88            api,
89        }
90    }
91}
92
93pub fn map_key<'a, K, V>(map: &Map<K, V>, key: K) -> String
94where
95    V: Serialize + DeserializeOwned,
96    K: PrimaryKey<'a>,
97{
98    String::from_utf8(raw_map_key(map, key)).unwrap()
99}
100
101pub fn raw_map_key<'a, K, V>(map: &Map<K, V>, key: K) -> Vec<u8>
102where
103    V: Serialize + DeserializeOwned,
104    K: PrimaryKey<'a>,
105{
106    map.key(key).deref().to_vec()
107}
108
109impl MockQuerierBuilder {
110    pub fn with_fallback_smart_handler<SH>(mut self, handler: SH) -> Self
111    where
112        SH: 'static + Fn(&Addr, &Binary) -> BinaryQueryResult,
113    {
114        self.fallback_smart_handler = Box::new(handler);
115        self
116    }
117
118    pub fn with_fallback_raw_handler<RH>(mut self, handler: RH) -> Self
119    where
120        RH: 'static + Fn(&Addr, &Binary) -> BinaryQueryResult,
121    {
122        self.fallback_raw_handler = Box::new(handler);
123        self
124    }
125
126    /// Add a smart query contract handler to the mock querier. The handler will be called when the
127    /// contract address is queried with the given message.
128    /// Usage:
129    /// ```rust
130    /// use cosmwasm_std::{from_json, to_json_binary};
131    /// use abstract_testing::MockQuerierBuilder;
132    /// use cosmwasm_std::testing::{MockQuerier, MockApi};
133    /// use abstract_sdk::mock_module::{MockModuleQueryMsg, MockModuleQueryResponse};
134    ///
135    /// let api = MockApi::default();
136    /// let contract_address = api.addr_make("contract_address");
137    /// let querier = MockQuerierBuilder::default().with_smart_handler(&contract_address, |msg| {
138    ///    // handle the message
139    ///     let res = match from_json::<MockModuleQueryMsg>(msg).unwrap() {
140    ///         // handle the message
141    ///         MockModuleQueryMsg =>
142    ///                         return to_json_binary(&MockModuleQueryResponse {}).map_err(|e| e.to_string())
143    ///    };
144    /// }).build();
145    ///
146    /// ```
147    pub fn with_smart_handler<SH>(mut self, contract: &Addr, handler: SH) -> Self
148    where
149        SH: 'static + Fn(&Binary) -> BinaryQueryResult,
150    {
151        self.smart_handlers
152            .insert(contract.clone(), Box::new(handler));
153        self
154    }
155
156    /// Add a raw query contract handler to the mock querier. The handler will be called when the
157    /// contract address is queried with the given message.
158    /// Usage:
159    ///
160    /// ```rust
161    /// use cosmwasm_std::{from_json, to_json_binary};
162    /// use abstract_testing::MockQuerierBuilder;
163    /// use cosmwasm_std::testing::{MockQuerier, MockApi};
164    /// use abstract_sdk::mock_module::{MockModuleQueryMsg, MockModuleQueryResponse};
165    ///
166    /// let api = MockApi::default();
167    /// let contract_address = api.addr_make("contract1");
168    /// let querier = MockQuerierBuilder::default().with_raw_handler(&contract_address, |key: &str| {
169    ///     // Example: Let's say, in the raw storage, the key "the key" maps to the value "the value"
170    ///     match key {
171    ///         "the key" => to_json_binary("the value").map_err(|e| e.to_string()),
172    ///         _ => to_json_binary("").map_err(|e| e.to_string())
173    ///     }
174    /// }).build();
175    /// ```
176    pub fn with_raw_handler<RH>(mut self, contract: &Addr, handler: RH) -> Self
177    where
178        RH: 'static + Fn(&str) -> BinaryQueryResult,
179    {
180        self.raw_handlers
181            .insert(contract.clone(), Box::new(handler));
182        self
183    }
184
185    fn insert_contract_key_value(&mut self, contract: &Addr, key: Vec<u8>, value: Binary) {
186        let raw_map = self.raw_mappings.entry(contract.clone()).or_default();
187        raw_map.insert(Binary::new(key), value);
188    }
189
190    /// Add a map entry to the querier for the given contract.
191    /// ```rust
192    /// use cw_storage_plus::Map;
193    /// use cosmwasm_std::testing::MockApi;
194    /// use abstract_testing::MockQuerierBuilder;
195    ///
196    /// let api = MockApi::default();
197    /// let contract_address = api.addr_make("contract1");
198    ///
199    /// const MAP: Map<String, String> = Map::new("map");
200    ///
201    /// MockQuerierBuilder::default()
202    ///     .with_contract_map_entry(
203    ///     &contract_address,
204    ///     MAP,
205    ///     ("key".to_string(), "value".to_string())
206    /// );
207    pub fn with_contract_map_entry<'a, K, V>(
208        self,
209        contract: &Addr,
210        cw_map: Map<K, V>,
211        entry: (K, V),
212    ) -> Self
213    where
214        K: PrimaryKey<'a>,
215        V: Serialize + DeserializeOwned,
216    {
217        self.with_contract_map_entries(contract, cw_map, vec![entry])
218    }
219
220    pub fn with_contract_map_entries<'a, K, V>(
221        mut self,
222        contract: &Addr,
223        cw_map: Map<K, V>,
224        entries: Vec<(K, V)>,
225    ) -> Self
226    where
227        K: PrimaryKey<'a>,
228        V: Serialize + DeserializeOwned,
229    {
230        for (key, value) in entries {
231            self.insert_contract_key_value(
232                contract,
233                raw_map_key(&cw_map, key),
234                to_json_binary(&value).unwrap(),
235            );
236        }
237
238        self
239    }
240
241    /// Add an empty map key to the querier for the given contract.
242    /// This is useful when you want the item to exist, but not have a value.
243    pub fn with_contract_map_key<'a, K, V>(
244        mut self,
245        contract: &Addr,
246        cw_map: Map<K, V>,
247        key: K,
248    ) -> Self
249    where
250        K: PrimaryKey<'a>,
251        V: Serialize + DeserializeOwned,
252    {
253        self.insert_contract_key_value(contract, raw_map_key(&cw_map, key), Binary::default());
254
255        self
256    }
257
258    /// Add an empty item key to the querier for the given contract.
259    /// This is useful when you want the item to exist, but not have a value.
260    pub fn with_empty_contract_item<T>(mut self, contract: &Addr, cw_item: Item<T>) -> Self
261    where
262        T: Serialize + DeserializeOwned,
263    {
264        self.insert_contract_key_value(contract, cw_item.as_slice().to_vec(), Binary::default());
265
266        self
267    }
268
269    /// Include a contract item in the mock querier.
270    /// ```rust
271    /// use cw_storage_plus::Item;
272    /// use cosmwasm_std::testing::MockApi;
273    /// use abstract_testing::MockQuerierBuilder;
274    ///
275    /// let api = MockApi::default();
276    /// let contract_address = api.addr_make("contract1");
277    ///
278    /// const ITEM: Item<String> = Item::new("item");
279    ///
280    /// MockQuerierBuilder::default()
281    ///     .with_contract_item(
282    ///     &contract_address,
283    ///     ITEM,
284    ///     &"value".to_string(),
285    /// );
286    /// ```
287    pub fn with_contract_item<T>(mut self, contract: &Addr, cw_item: Item<T>, value: &T) -> Self
288    where
289        T: Serialize + DeserializeOwned,
290    {
291        self.insert_contract_key_value(
292            contract,
293            cw_item.as_slice().to_vec(),
294            to_json_binary(value).unwrap(),
295        );
296
297        self
298    }
299
300    /// Add a specific version of the contract to the mock querier.
301    /// ```rust
302    /// use abstract_testing::MockQuerierBuilder;
303    /// use cosmwasm_std::testing::MockApi;
304    ///
305    /// let api = MockApi::default();
306    /// let contract_address = api.addr_make("contract1");
307    ///
308    /// MockQuerierBuilder::default()
309    ///    .with_contract_version(&contract_address, "contract1", "v1.0.0");
310    /// ```
311    pub fn with_contract_version(
312        self,
313        contract: &Addr,
314        name: impl Into<String>,
315        version: impl Into<String>,
316    ) -> Self {
317        self.with_contract_item(
318            contract,
319            CONTRACT,
320            &ContractVersion {
321                contract: name.into(),
322                version: version.into(),
323            },
324        )
325    }
326    /// set the SDK-level contract admin for a contract.
327    pub fn with_contract_admin(mut self, contract: &Addr, admin: &Addr) -> Self {
328        self.contract_admin.insert(contract.clone(), admin.clone());
329        self
330    }
331
332    /// Build the [`MockQuerier`].
333    pub fn build(mut self) -> MockQuerier {
334        self.base.update_wasm(move |wasm| {
335            let res = match wasm {
336                WasmQuery::Raw { contract_addr, key } => {
337                    let str_key = std::str::from_utf8(key.as_slice()).unwrap();
338                    let addr = Addr::unchecked(contract_addr);
339
340                    // First check for raw mappings
341                    if let Some(raw_map) = self.raw_mappings.get(&addr) {
342                        if let Some(value) = raw_map.get(key) {
343                            return SystemResult::Ok(ContractResult::Ok(value.clone()));
344                        }
345                    }
346
347                    // Then check the handlers
348                    let raw_handler = self.raw_handlers.get(&addr);
349
350                    match raw_handler {
351                        Some(handler) => (*handler)(str_key),
352                        None => (*self.fallback_raw_handler)(&addr, key),
353                    }
354                }
355                WasmQuery::Smart { contract_addr, msg } => {
356                    let addr = Addr::unchecked(contract_addr);
357                    let contract_handler = self.smart_handlers.get(&addr);
358
359                    match contract_handler {
360                        Some(handler) => (*handler)(msg),
361                        None => (*self.fallback_smart_handler)(&addr, msg),
362                    }
363                }
364                WasmQuery::ContractInfo { contract_addr } => {
365                    let addr = Addr::unchecked(contract_addr);
366                    let creator = self.api.addr_make(crate::OWNER);
367                    let info = ContractInfoResponse::new(
368                        1,
369                        creator,
370                        self.contract_admin.get(&addr).map(Addr::unchecked),
371                        false,
372                        None,
373                    );
374                    Ok(to_json_binary(&info).unwrap())
375                }
376                WasmQuery::CodeInfo { code_id } => {
377                    let creator = self.api.addr_make(crate::OWNER);
378                    let checksum = native_addrs::BLOB_CHECKSUM;
379
380                    let code_info = CodeInfoResponse::new(*code_id, creator, checksum.into());
381                    Ok(to_json_binary(&code_info).unwrap())
382                }
383                unexpected => panic!("Unexpected query: {unexpected:?}"),
384            };
385
386            match res {
387                Ok(res) => SystemResult::Ok(ContractResult::Ok(res)),
388                Err(e) => SystemResult::Ok(ContractResult::Err(e)),
389            }
390        });
391        self.base
392    }
393}
394
395pub trait MockQuerierOwnership {
396    /// Add the [`cw_gov_ownable::Ownership`] to the querier.
397    fn with_owner(self, contract: &Addr, owner: Option<&Addr>) -> Self;
398}
399
400impl MockQuerierOwnership for MockQuerierBuilder {
401    fn with_owner(mut self, contract: &Addr, owner: Option<&Addr>) -> Self {
402        let owner = if let Some(owner) = owner {
403            GovernanceDetails::Monarchy {
404                monarch: owner.clone(),
405            }
406        } else {
407            GovernanceDetails::Renounced {}
408        };
409        self = self.with_contract_item(
410            contract,
411            Item::new(OWNERSHIP_STORAGE_KEY),
412            &Ownership {
413                owner,
414                pending_owner: None,
415                pending_expiry: None,
416            },
417        );
418        self
419    }
420}
421
422pub fn wrap_querier(querier: &MockQuerier) -> QuerierWrapper<'_, Empty> {
423    QuerierWrapper::<Empty>::new(querier)
424}
425
426#[cfg(test)]
427mod tests {
428    use abstract_std::{
429        account::state::{ACCOUNT_ID, ACCOUNT_MODULES},
430        objects::ABSTRACT_ACCOUNT_ID,
431        registry::state::ACCOUNT_ADDRESSES,
432    };
433
434    use super::*;
435    use cosmwasm_std::testing::mock_dependencies;
436
437    mod account {
438
439        use abstract_std::registry::Account;
440
441        use crate::abstract_mock_querier_builder;
442
443        use super::*;
444
445        #[test]
446        fn should_return_admin_account_address() {
447            let mut deps = mock_dependencies();
448            deps.querier = abstract_mock_querier(deps.api);
449            let abstr = AbstractMockAddrs::new(deps.api);
450
451            let actual = ACCOUNT_ADDRESSES.query(
452                &wrap_querier(&deps.querier),
453                abstr.registry,
454                &ABSTRACT_ACCOUNT_ID,
455            );
456
457            let expected = abstr.account;
458
459            assert_eq!(actual, Ok(Some(expected)));
460        }
461
462        #[test]
463        fn should_return_account_address() {
464            let mut deps = mock_dependencies();
465            let account = Account::new(deps.api.addr_make("my_account"));
466            deps.querier = abstract_mock_querier_builder(deps.api)
467                .account(&account, TEST_ACCOUNT_ID)
468                .build();
469            let abstr = AbstractMockAddrs::new(deps.api);
470
471            let actual = ACCOUNT_ADDRESSES.query(
472                &wrap_querier(&deps.querier),
473                abstr.registry,
474                &TEST_ACCOUNT_ID,
475            );
476
477            assert_eq!(actual, Ok(Some(account)));
478        }
479    }
480
481    mod queries {
482        use super::*;
483
484        use abstract_sdk::mock_module::{MockModuleQueryMsg, MockModuleQueryResponse};
485        use cosmwasm_std::QueryRequest;
486
487        #[test]
488        fn smart_query() {
489            let api = MockApi::default();
490            // ## ANCHOR: smart_query
491            let contract_address = api.addr_make("contract_address");
492            let querier = MockQuerierBuilder::default()
493                .with_smart_handler(&contract_address, |msg| {
494                    // handle the message
495                    let MockModuleQueryMsg {} = from_json::<MockModuleQueryMsg>(msg).unwrap();
496                    to_json_binary(&MockModuleQueryResponse {}).map_err(|e| e.to_string())
497                })
498                .build();
499            // ## ANCHOR_END: smart_query
500
501            let resp_bin = querier
502                .handle_query(&QueryRequest::Wasm(WasmQuery::Smart {
503                    contract_addr: contract_address.to_string(),
504                    msg: to_json_binary(&MockModuleQueryMsg {}).unwrap(),
505                }))
506                .unwrap()
507                .unwrap();
508            let resp: MockModuleQueryResponse = from_json(resp_bin).unwrap();
509
510            assert_eq!(resp, MockModuleQueryResponse {});
511        }
512
513        #[test]
514        fn raw_query() {
515            let api = MockApi::default();
516            // ## ANCHOR: raw_query
517            let contract_address = api.addr_make("contract_address");
518            let querier = MockQuerierBuilder::default()
519                .with_raw_handler(&contract_address, |key: &str| {
520                    // Example: Let's say, in the raw storage, the key "the_key" maps to the value "the_value"
521                    match key {
522                        "the_key" => to_json_binary("the_value").map_err(|e| e.to_string()),
523                        _ => to_json_binary("").map_err(|e| e.to_string()),
524                    }
525                })
526                .build();
527            // ## ANCHOR_END: raw_query
528
529            let resp_bin = querier
530                .handle_query(&QueryRequest::Wasm(WasmQuery::Raw {
531                    contract_addr: contract_address.to_string(),
532                    key: Binary::from("the_key".joined_key()),
533                }))
534                .unwrap()
535                .unwrap();
536            let resp: String = from_json(resp_bin).unwrap();
537
538            assert_eq!(resp, "the_value");
539        }
540    }
541
542    mod account_id {
543        use crate::abstract_mock_querier_builder;
544
545        use super::*;
546
547        #[test]
548        fn should_return_admin_acct_id() {
549            let mut deps = mock_dependencies();
550            deps.querier = abstract_mock_querier(deps.api);
551            let root_account = admin_account(deps.api);
552
553            let actual =
554                ACCOUNT_ID.query(&wrap_querier(&deps.querier), root_account.addr().clone());
555
556            assert_eq!(actual, Ok(ABSTRACT_ACCOUNT_ID));
557        }
558
559        #[test]
560        fn should_return_test_acct_id() {
561            let mut deps = mock_dependencies();
562            let test_base = test_account(deps.api);
563            deps.querier = abstract_mock_querier_builder(deps.api)
564                .account(&test_base, TEST_ACCOUNT_ID)
565                .build();
566
567            let actual = ACCOUNT_ID.query(&wrap_querier(&deps.querier), test_base.into_addr());
568
569            assert_eq!(actual, Ok(TEST_ACCOUNT_ID));
570        }
571    }
572
573    mod account_modules {
574        use super::*;
575
576        #[test]
577        fn should_return_test_module_address_for_test_module() {
578            let mut deps = mock_dependencies();
579            deps.querier = abstract_mock_querier(deps.api);
580            let abstr = AbstractMockAddrs::new(deps.api);
581
582            let actual = ACCOUNT_MODULES.query(
583                &wrap_querier(&deps.querier),
584                abstr.account.into_addr(),
585                TEST_MODULE_ID,
586            );
587
588            assert_eq!(actual, Ok(Some(abstr.module_address)));
589        }
590
591        // #[test]
592        // fn should_return_none_for_unknown_module() {
593        //     let mut deps = mock_dependencies();
594        //     deps.querier = querier();
595        //
596        //     let actual = ACCOUNT_MODULES.query(
597        //         &wrap_querier(&deps.querier),
598        //         Addr::unchecked(TEST_ACCOUNT),
599        //         "unknown_module",
600        //     );
601        //
602        //     assert_that!(actual).is_ok().is_none();
603        // }
604    }
605}
abstract_testing/mock_querier.rs

abstract_testing/
mock_querier.rs
