apollo_cw_multi_test/
wasm.rs

1use std::collections::HashMap;
2use std::fmt;
3use std::ops::Deref;
4
5use cosmwasm_std::{
6    to_binary, Addr, Api, Attribute, BankMsg, Binary, BlockInfo, Coin, ContractInfo,
7    ContractInfoResponse, CustomQuery, Deps, DepsMut, Env, Event, MessageInfo, Order, Querier,
8    QuerierWrapper, Record, Reply, ReplyOn, Response, StdResult, Storage, SubMsg, SubMsgResponse,
9    SubMsgResult, TransactionInfo, WasmMsg, WasmQuery,
10};
11use prost::Message;
12use schemars::JsonSchema;
13use serde::de::DeserializeOwned;
14use serde::{Deserialize, Serialize};
15
16use cw_storage_plus::Map;
17
18use crate::app::{CosmosRouter, RouterQuerier};
19use crate::contracts::Contract;
20use crate::error::Error;
21use crate::executor::AppResponse;
22use crate::prefixed_storage::{prefixed, prefixed_read, PrefixedStorage, ReadonlyPrefixedStorage};
23use crate::transactions::transactional;
24use crate::{AddressGenerator, SimpleAddressGenerator};
25use cosmwasm_std::testing::mock_wasmd_attr;
26
27use anyhow::{bail, Context, Result as AnyResult};
28
29// Contract state is kept in Storage, separate from the contracts themselves
30const CONTRACTS: Map<&Addr, ContractData> = Map::new("contracts");
31
32pub const NAMESPACE_WASM: &[u8] = b"wasm";
33const CONTRACT_ATTR: &str = "_contract_addr";
34
35#[derive(Clone, std::fmt::Debug, PartialEq, Eq, JsonSchema)]
36pub struct WasmSudo {
37    pub contract_addr: Addr,
38    pub msg: Binary,
39}
40
41impl WasmSudo {
42    pub fn new<T: Serialize>(contract_addr: &Addr, msg: &T) -> StdResult<WasmSudo> {
43        Ok(WasmSudo {
44            contract_addr: contract_addr.clone(),
45            msg: to_binary(msg)?,
46        })
47    }
48}
49
50/// Contract Data includes information about contract, equivalent of `ContractInfo` in wasmd
51/// interface.
52#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, JsonSchema)]
53pub struct ContractData {
54    /// Identifier of stored contract code
55    pub code_id: usize,
56    /// Address of account who initially instantiated the contract
57    pub creator: Addr,
58    /// Optional address of account who can execute migrations
59    pub admin: Option<Addr>,
60    /// Metadata passed while contract instantiation
61    pub label: String,
62    /// Blockchain height in the moment of instantiating the contract
63    pub created: u64,
64}
65
66pub trait Wasm<ExecC, QueryC> {
67    /// Handles all WasmQuery requests
68    fn query(
69        &self,
70        api: &dyn Api,
71        storage: &dyn Storage,
72        querier: &dyn Querier,
73        block: &BlockInfo,
74        request: WasmQuery,
75    ) -> AnyResult<Binary>;
76
77    /// Handles all WasmMsg messages
78    fn execute(
79        &self,
80        api: &dyn Api,
81        storage: &mut dyn Storage,
82        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
83        block: &BlockInfo,
84        sender: Addr,
85        msg: WasmMsg,
86    ) -> AnyResult<AppResponse>;
87
88    /// Admin interface, cannot be called via CosmosMsg
89    fn sudo(
90        &self,
91        api: &dyn Api,
92        contract_addr: Addr,
93        storage: &mut dyn Storage,
94        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
95        block: &BlockInfo,
96        msg: Binary,
97    ) -> AnyResult<AppResponse>;
98}
99
100pub struct WasmKeeper<ExecC, QueryC> {
101    /// code is in-memory lookup that stands in for wasm code
102    /// this can only be edited on the WasmRouter, and just read in caches
103    codes: HashMap<usize, Box<dyn Contract<ExecC, QueryC>>>,
104    /// Just markers to make type elision fork when using it as `Wasm` trait
105    _p: std::marker::PhantomData<QueryC>,
106    generator: Box<dyn AddressGenerator>,
107}
108
109impl<ExecC, QueryC> Default for WasmKeeper<ExecC, QueryC> {
110    fn default() -> Self {
111        Self {
112            codes: HashMap::default(),
113            _p: std::marker::PhantomData,
114            generator: Box::new(SimpleAddressGenerator {}),
115        }
116    }
117}
118
119impl<ExecC, QueryC> Wasm<ExecC, QueryC> for WasmKeeper<ExecC, QueryC>
120where
121    ExecC: Clone + fmt::Debug + PartialEq + JsonSchema + DeserializeOwned + 'static,
122    QueryC: CustomQuery + DeserializeOwned + 'static,
123{
124    fn query(
125        &self,
126        api: &dyn Api,
127        storage: &dyn Storage,
128        querier: &dyn Querier,
129        block: &BlockInfo,
130        request: WasmQuery,
131    ) -> AnyResult<Binary> {
132        match request {
133            WasmQuery::Smart { contract_addr, msg } => {
134                let addr = api.addr_validate(&contract_addr)?;
135                self.query_smart(addr, api, storage, querier, block, msg.into())
136            }
137            WasmQuery::Raw { contract_addr, key } => {
138                let addr = api.addr_validate(&contract_addr)?;
139                Ok(self.query_raw(addr, storage, &key))
140            }
141            WasmQuery::ContractInfo { contract_addr } => {
142                let addr = api.addr_validate(&contract_addr)?;
143                let contract = self.load_contract(storage, &addr)?;
144                let mut res = ContractInfoResponse::default();
145                res.code_id = contract.code_id as u64;
146                res.creator = contract.creator.to_string();
147                res.admin = contract.admin.map(|x| x.into());
148                to_binary(&res).map_err(Into::into)
149            }
150            query => bail!(Error::UnsupportedWasmQuery(query)),
151        }
152    }
153
154    fn execute(
155        &self,
156        api: &dyn Api,
157        storage: &mut dyn Storage,
158        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
159        block: &BlockInfo,
160        sender: Addr,
161        msg: WasmMsg,
162    ) -> AnyResult<AppResponse> {
163        self.execute_wasm(api, storage, router, block, sender.clone(), msg.clone())
164            .context(format!(
165                "error executing WasmMsg:\nsender: {}\n{:?}",
166                sender, msg
167            ))
168    }
169
170    fn sudo(
171        &self,
172        api: &dyn Api,
173        contract: Addr,
174        storage: &mut dyn Storage,
175        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
176        block: &BlockInfo,
177        msg: Binary,
178    ) -> AnyResult<AppResponse> {
179        let custom_event = Event::new("sudo").add_attribute(CONTRACT_ATTR, &contract);
180
181        let res = self.call_sudo(contract.clone(), api, storage, router, block, msg.to_vec())?;
182        let (res, msgs) = self.build_app_response(&contract, custom_event, res);
183        self.process_response(api, router, storage, block, contract, res, msgs)
184    }
185}
186
187impl<ExecC, QueryC> WasmKeeper<ExecC, QueryC> {
188    pub fn store_code(&mut self, code: Box<dyn Contract<ExecC, QueryC>>) -> usize {
189        let idx = self.codes.len() + 1;
190        self.codes.insert(idx, code);
191        idx
192    }
193
194    pub fn load_contract(&self, storage: &dyn Storage, address: &Addr) -> AnyResult<ContractData> {
195        CONTRACTS
196            .load(&prefixed_read(storage, NAMESPACE_WASM), address)
197            .map_err(Into::into)
198    }
199
200    pub fn dump_wasm_raw(&self, storage: &dyn Storage, address: &Addr) -> Vec<Record> {
201        let storage = self.contract_storage_readonly(storage, address);
202        storage.range(None, None, Order::Ascending).collect()
203    }
204
205    fn contract_namespace(&self, contract: &Addr) -> Vec<u8> {
206        let mut name = b"contract_data/".to_vec();
207        name.extend_from_slice(contract.as_bytes());
208        name
209    }
210
211    fn contract_storage<'a>(
212        &self,
213        storage: &'a mut dyn Storage,
214        address: &Addr,
215    ) -> Box<dyn Storage + 'a> {
216        // We double-namespace this, once from global storage -> wasm_storage
217        // then from wasm_storage -> the contracts subspace
218        let namespace = self.contract_namespace(address);
219        let storage = PrefixedStorage::multilevel(storage, &[NAMESPACE_WASM, &namespace]);
220        Box::new(storage)
221    }
222
223    // fails RUNTIME if you try to write. please don't
224    fn contract_storage_readonly<'a>(
225        &self,
226        storage: &'a dyn Storage,
227        address: &Addr,
228    ) -> Box<dyn Storage + 'a> {
229        // We double-namespace this, once from global storage -> wasm_storage
230        // then from wasm_storage -> the contracts subspace
231        let namespace = self.contract_namespace(address);
232        let storage = ReadonlyPrefixedStorage::multilevel(storage, &[NAMESPACE_WASM, &namespace]);
233        Box::new(storage)
234    }
235
236    fn verify_attributes(attributes: &[Attribute]) -> AnyResult<()> {
237        for attr in attributes {
238            let key = attr.key.trim();
239            let val = attr.value.trim();
240
241            if key.is_empty() {
242                bail!(Error::empty_attribute_key(val));
243            }
244
245            if val.is_empty() {
246                bail!(Error::empty_attribute_value(key));
247            }
248
249            if key.starts_with('_') {
250                bail!(Error::reserved_attribute_key(key));
251            }
252        }
253
254        Ok(())
255    }
256
257    fn verify_response<T>(response: Response<T>) -> AnyResult<Response<T>>
258    where
259        T: Clone + fmt::Debug + PartialEq + JsonSchema,
260    {
261        Self::verify_attributes(&response.attributes)?;
262
263        for event in &response.events {
264            Self::verify_attributes(&event.attributes)?;
265            let ty = event.ty.trim();
266            if ty.len() < 2 {
267                bail!(Error::event_type_too_short(ty));
268            }
269        }
270
271        Ok(response)
272    }
273}
274
275impl<ExecC, QueryC> WasmKeeper<ExecC, QueryC>
276where
277    ExecC: Clone + fmt::Debug + PartialEq + JsonSchema + DeserializeOwned + 'static,
278    QueryC: CustomQuery + DeserializeOwned + 'static,
279{
280    pub fn new() -> Self {
281        Self::default()
282    }
283
284    pub fn new_with_custom_address_generator(generator: impl AddressGenerator + 'static) -> Self {
285        let default = Self::default();
286        Self {
287            codes: default.codes,
288            _p: default._p,
289            generator: Box::new(generator),
290        }
291    }
292
293    pub fn query_smart(
294        &self,
295        address: Addr,
296        api: &dyn Api,
297        storage: &dyn Storage,
298        querier: &dyn Querier,
299        block: &BlockInfo,
300        msg: Vec<u8>,
301    ) -> AnyResult<Binary> {
302        self.with_storage_readonly(
303            api,
304            storage,
305            querier,
306            block,
307            address,
308            |handler, deps, env| handler.query(deps, env, msg),
309        )
310    }
311
312    pub fn query_raw(&self, address: Addr, storage: &dyn Storage, key: &[u8]) -> Binary {
313        let storage = self.contract_storage_readonly(storage, &address);
314        let data = storage.get(key).unwrap_or_default();
315        data.into()
316    }
317
318    fn send<T>(
319        &self,
320        api: &dyn Api,
321        storage: &mut dyn Storage,
322        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
323        block: &BlockInfo,
324        sender: T,
325        recipient: String,
326        amount: &[Coin],
327    ) -> AnyResult<AppResponse>
328    where
329        T: Into<Addr>,
330    {
331        if !amount.is_empty() {
332            let msg: cosmwasm_std::CosmosMsg<ExecC> = BankMsg::Send {
333                to_address: recipient,
334                amount: amount.to_vec(),
335            }
336            .into();
337            let res = router.execute(api, storage, block, sender.into(), msg)?;
338            Ok(res)
339        } else {
340            Ok(AppResponse::default())
341        }
342    }
343
344    /// unified logic for UpdateAdmin and ClearAdmin messages
345    fn update_admin(
346        &self,
347        api: &dyn Api,
348        storage: &mut dyn Storage,
349        sender: Addr,
350        contract_addr: &str,
351        new_admin: Option<String>,
352    ) -> AnyResult<AppResponse> {
353        let contract_addr = api.addr_validate(contract_addr)?;
354        let admin = new_admin.map(|a| api.addr_validate(&a)).transpose()?;
355
356        // check admin status
357        let mut data = self.load_contract(storage, &contract_addr)?;
358        if data.admin != Some(sender) {
359            bail!("Only admin can update the contract admin: {:?}", data.admin);
360        }
361        // update admin field
362        data.admin = admin;
363        self.save_contract(storage, &contract_addr, &data)?;
364
365        // no custom event here
366        Ok(AppResponse {
367            data: None,
368            events: vec![],
369        })
370    }
371
372    // this returns the contract address as well, so we can properly resend the data
373    fn execute_wasm(
374        &self,
375        api: &dyn Api,
376        storage: &mut dyn Storage,
377        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
378        block: &BlockInfo,
379        sender: Addr,
380        wasm_msg: WasmMsg,
381    ) -> AnyResult<AppResponse> {
382        match wasm_msg {
383            WasmMsg::Execute {
384                contract_addr,
385                msg,
386                funds,
387            } => {
388                let contract_addr = api.addr_validate(&contract_addr)?;
389                // first move the cash
390                self.send(
391                    api,
392                    storage,
393                    router,
394                    block,
395                    sender.clone(),
396                    contract_addr.clone().into(),
397                    &funds,
398                )?;
399
400                // then call the contract
401                let info = MessageInfo { sender, funds };
402                let res = self.call_execute(
403                    api,
404                    storage,
405                    contract_addr.clone(),
406                    router,
407                    block,
408                    info,
409                    msg.to_vec(),
410                )?;
411
412                let custom_event =
413                    Event::new("execute").add_attribute(CONTRACT_ATTR, &contract_addr);
414
415                let (res, msgs) = self.build_app_response(&contract_addr, custom_event, res);
416                let mut res =
417                    self.process_response(api, router, storage, block, contract_addr, res, msgs)?;
418                res.data = execute_response(res.data);
419                Ok(res)
420            }
421            WasmMsg::Instantiate {
422                admin,
423                code_id,
424                msg,
425                funds,
426                label,
427            } => {
428                if label.is_empty() {
429                    bail!("Label is required on all contracts");
430                }
431
432                let contract_addr = self.register_contract(
433                    api,
434                    storage,
435                    code_id as usize,
436                    sender.clone(),
437                    admin.map(Addr::unchecked),
438                    label,
439                    block.height,
440                )?;
441
442                // move the cash
443                self.send(
444                    api,
445                    storage,
446                    router,
447                    block,
448                    sender.clone(),
449                    contract_addr.clone().into(),
450                    &funds,
451                )?;
452
453                // then call the contract
454                let info = MessageInfo { sender, funds };
455                let res = self.call_instantiate(
456                    contract_addr.clone(),
457                    api,
458                    storage,
459                    router,
460                    block,
461                    info,
462                    msg.to_vec(),
463                )?;
464
465                let custom_event = Event::new("instantiate")
466                    .add_attribute(CONTRACT_ATTR, &contract_addr)
467                    .add_attribute("code_id", code_id.to_string());
468
469                let (res, msgs) = self.build_app_response(&contract_addr, custom_event, res);
470                let mut res = self.process_response(
471                    api,
472                    router,
473                    storage,
474                    block,
475                    contract_addr.clone(),
476                    res,
477                    msgs,
478                )?;
479                res.data = Some(instantiate_response(res.data, &contract_addr));
480                Ok(res)
481            }
482            WasmMsg::Migrate {
483                contract_addr,
484                new_code_id,
485                msg,
486            } => {
487                let contract_addr = api.addr_validate(&contract_addr)?;
488
489                // check admin status and update the stored code_id
490                let new_code_id = new_code_id as usize;
491                if !self.codes.contains_key(&new_code_id) {
492                    bail!("Cannot migrate contract to unregistered code id");
493                }
494                let mut data = self.load_contract(storage, &contract_addr)?;
495                if data.admin != Some(sender) {
496                    bail!("Only admin can migrate contract: {:?}", data.admin);
497                }
498                data.code_id = new_code_id;
499                self.save_contract(storage, &contract_addr, &data)?;
500
501                // then call migrate
502                let res = self.call_migrate(
503                    contract_addr.clone(),
504                    api,
505                    storage,
506                    router,
507                    block,
508                    msg.to_vec(),
509                )?;
510
511                let custom_event = Event::new("migrate")
512                    .add_attribute(CONTRACT_ATTR, &contract_addr)
513                    .add_attribute("code_id", new_code_id.to_string());
514                let (res, msgs) = self.build_app_response(&contract_addr, custom_event, res);
515                let mut res =
516                    self.process_response(api, router, storage, block, contract_addr, res, msgs)?;
517                res.data = execute_response(res.data);
518                Ok(res)
519            }
520            WasmMsg::UpdateAdmin {
521                contract_addr,
522                admin,
523            } => self.update_admin(api, storage, sender, &contract_addr, Some(admin)),
524            WasmMsg::ClearAdmin { contract_addr } => {
525                self.update_admin(api, storage, sender, &contract_addr, None)
526            }
527            msg => bail!(Error::UnsupportedWasmMsg(msg)),
528        }
529    }
530
531    /// This will execute the given messages, making all changes to the local cache.
532    /// This *will* write some data to the cache if the message fails half-way through.
533    /// All sequential calls to RouterCache will be one atomic unit (all commit or all fail).
534    ///
535    /// For normal use cases, you can use Router::execute() or Router::execute_multi().
536    /// This is designed to be handled internally as part of larger process flows.
537    ///
538    /// The `data` on `AppResponse` is data returned from `reply` call, not from execution of
539    /// submessage itself. In case if `reply` is not called, no `data` is set.
540    fn execute_submsg(
541        &self,
542        api: &dyn Api,
543        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
544        storage: &mut dyn Storage,
545        block: &BlockInfo,
546        contract: Addr,
547        msg: SubMsg<ExecC>,
548    ) -> AnyResult<AppResponse> {
549        let SubMsg {
550            msg, id, reply_on, ..
551        } = msg;
552
553        // execute in cache
554        let res = transactional(storage, |write_cache, _| {
555            router.execute(api, write_cache, block, contract.clone(), msg)
556        });
557
558        // call reply if meaningful
559        if let Ok(mut r) = res {
560            if matches!(reply_on, ReplyOn::Always | ReplyOn::Success) {
561                let reply = Reply {
562                    id,
563                    result: SubMsgResult::Ok(SubMsgResponse {
564                        events: r.events.clone(),
565                        data: r.data,
566                    }),
567                };
568                // do reply and combine it with the original response
569                let reply_res = self._reply(api, router, storage, block, contract, reply)?;
570                // override data
571                r.data = reply_res.data;
572                // append the events
573                r.events.extend_from_slice(&reply_res.events);
574            } else {
575                // reply is not called, no data should be rerturned
576                r.data = None;
577            }
578
579            Ok(r)
580        } else if let Err(e) = res {
581            if matches!(reply_on, ReplyOn::Always | ReplyOn::Error) {
582                let reply = Reply {
583                    id,
584                    result: SubMsgResult::Err(e.to_string()),
585                };
586                self._reply(api, router, storage, block, contract, reply)
587            } else {
588                Err(e)
589            }
590        } else {
591            res
592        }
593    }
594
595    fn _reply(
596        &self,
597        api: &dyn Api,
598        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
599        storage: &mut dyn Storage,
600        block: &BlockInfo,
601        contract: Addr,
602        reply: Reply,
603    ) -> AnyResult<AppResponse> {
604        let ok_attr = if reply.result.is_ok() {
605            "handle_success"
606        } else {
607            "handle_failure"
608        };
609        let custom_event = Event::new("reply")
610            .add_attribute(CONTRACT_ATTR, &contract)
611            .add_attribute("mode", ok_attr);
612
613        let res = self.call_reply(contract.clone(), api, storage, router, block, reply)?;
614        let (res, msgs) = self.build_app_response(&contract, custom_event, res);
615        self.process_response(api, router, storage, block, contract, res, msgs)
616    }
617
618    // this captures all the events and data from the contract call.
619    // it does not handle the messages
620    fn build_app_response(
621        &self,
622        contract: &Addr,
623        custom_event: Event, // entry-point specific custom event added by x/wasm
624        response: Response<ExecC>,
625    ) -> (AppResponse, Vec<SubMsg<ExecC>>) {
626        let Response {
627            messages,
628            attributes,
629            events,
630            data,
631            ..
632        } = response;
633
634        // always add custom event
635        let mut app_events = Vec::with_capacity(2 + events.len());
636        app_events.push(custom_event);
637
638        // we only emit the `wasm` event if some attributes are specified
639        if !attributes.is_empty() {
640            // turn attributes into event and place it first
641            let wasm_event = Event::new("wasm")
642                .add_attribute(CONTRACT_ATTR, contract)
643                .add_attributes(attributes);
644            app_events.push(wasm_event);
645        }
646
647        // These need to get `wasm-` prefix to match the wasmd semantics (custom wasm messages cannot
648        // fake system level event types, like transfer from the bank module)
649        let wasm_events = events.into_iter().map(|mut ev| {
650            ev.ty = format!("wasm-{}", ev.ty);
651            ev.attributes
652                .insert(0, mock_wasmd_attr(CONTRACT_ATTR, contract));
653            ev
654        });
655        app_events.extend(wasm_events);
656
657        let app = AppResponse {
658            events: app_events,
659            data,
660        };
661        (app, messages)
662    }
663
664    fn process_response(
665        &self,
666        api: &dyn Api,
667        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
668        storage: &mut dyn Storage,
669        block: &BlockInfo,
670        contract: Addr,
671        response: AppResponse,
672        messages: Vec<SubMsg<ExecC>>,
673    ) -> AnyResult<AppResponse> {
674        let AppResponse { mut events, data } = response;
675
676        // recurse in all messages
677        let data = messages.into_iter().try_fold(data, |data, resend| {
678            let subres =
679                self.execute_submsg(api, router, storage, block, contract.clone(), resend)?;
680            events.extend_from_slice(&subres.events);
681            Ok::<_, anyhow::Error>(subres.data.or(data))
682        })?;
683
684        Ok(AppResponse { events, data })
685    }
686
687    /// This just creates an address and empty storage instance, returning the new address
688    /// You must call init after this to set up the contract properly.
689    /// These are separated into two steps to have cleaner return values.
690    pub fn register_contract(
691        &self,
692        api: &dyn Api,
693        storage: &mut dyn Storage,
694        code_id: usize,
695        creator: Addr,
696        admin: impl Into<Option<Addr>>,
697        label: String,
698        created: u64,
699    ) -> AnyResult<Addr> {
700        if !self.codes.contains_key(&code_id) {
701            bail!("Cannot init contract with unregistered code id");
702        }
703        // generate a new contract address
704        let instance_id = self.instance_count(storage) as u64;
705        let addr = self
706            .generator
707            .contract_address(api, storage, code_id as u64, instance_id)?;
708
709        let info = ContractData {
710            code_id,
711            creator,
712            admin: admin.into(),
713            label,
714            created,
715        };
716        self.save_contract(storage, &addr, &info)?;
717        Ok(addr)
718    }
719
720    pub fn call_execute(
721        &self,
722        api: &dyn Api,
723        storage: &mut dyn Storage,
724        address: Addr,
725        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
726        block: &BlockInfo,
727        info: MessageInfo,
728        msg: Vec<u8>,
729    ) -> AnyResult<Response<ExecC>> {
730        Self::verify_response(self.with_storage(
731            api,
732            storage,
733            router,
734            block,
735            address,
736            |contract, deps, env| contract.execute(deps, env, info, msg),
737        )?)
738    }
739
740    pub fn call_instantiate(
741        &self,
742        address: Addr,
743        api: &dyn Api,
744        storage: &mut dyn Storage,
745        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
746        block: &BlockInfo,
747        info: MessageInfo,
748        msg: Vec<u8>,
749    ) -> AnyResult<Response<ExecC>> {
750        Self::verify_response(self.with_storage(
751            api,
752            storage,
753            router,
754            block,
755            address,
756            |contract, deps, env| contract.instantiate(deps, env, info, msg),
757        )?)
758    }
759
760    pub fn call_reply(
761        &self,
762        address: Addr,
763        api: &dyn Api,
764        storage: &mut dyn Storage,
765        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
766        block: &BlockInfo,
767        reply: Reply,
768    ) -> AnyResult<Response<ExecC>> {
769        Self::verify_response(self.with_storage(
770            api,
771            storage,
772            router,
773            block,
774            address,
775            |contract, deps, env| contract.reply(deps, env, reply),
776        )?)
777    }
778
779    pub fn call_sudo(
780        &self,
781        address: Addr,
782        api: &dyn Api,
783        storage: &mut dyn Storage,
784        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
785        block: &BlockInfo,
786        msg: Vec<u8>,
787    ) -> AnyResult<Response<ExecC>> {
788        Self::verify_response(self.with_storage(
789            api,
790            storage,
791            router,
792            block,
793            address,
794            |contract, deps, env| contract.sudo(deps, env, msg),
795        )?)
796    }
797
798    pub fn call_migrate(
799        &self,
800        address: Addr,
801        api: &dyn Api,
802        storage: &mut dyn Storage,
803        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
804        block: &BlockInfo,
805        msg: Vec<u8>,
806    ) -> AnyResult<Response<ExecC>> {
807        Self::verify_response(self.with_storage(
808            api,
809            storage,
810            router,
811            block,
812            address,
813            |contract, deps, env| contract.migrate(deps, env, msg),
814        )?)
815    }
816
817    fn get_env<T: Into<Addr>>(&self, address: T, block: &BlockInfo) -> Env {
818        Env {
819            block: block.clone(),
820            contract: ContractInfo {
821                address: address.into(),
822            },
823            transaction: Some(TransactionInfo { index: 0 }),
824        }
825    }
826
827    fn with_storage_readonly<F, T>(
828        &self,
829        api: &dyn Api,
830        storage: &dyn Storage,
831        querier: &dyn Querier,
832        block: &BlockInfo,
833        address: Addr,
834        action: F,
835    ) -> AnyResult<T>
836    where
837        F: FnOnce(&Box<dyn Contract<ExecC, QueryC>>, Deps<QueryC>, Env) -> AnyResult<T>,
838    {
839        let contract = self.load_contract(storage, &address)?;
840        let handler = self
841            .codes
842            .get(&contract.code_id)
843            .ok_or(Error::UnregisteredCodeId(contract.code_id))?;
844        let storage = self.contract_storage_readonly(storage, &address);
845        let env = self.get_env(address, block);
846
847        let deps = Deps {
848            storage: storage.as_ref(),
849            api: api.deref(),
850            querier: QuerierWrapper::new(querier),
851        };
852        action(handler, deps, env)
853    }
854
855    fn with_storage<F, T>(
856        &self,
857        api: &dyn Api,
858        storage: &mut dyn Storage,
859        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
860        block: &BlockInfo,
861        address: Addr,
862        action: F,
863    ) -> AnyResult<T>
864    where
865        F: FnOnce(&Box<dyn Contract<ExecC, QueryC>>, DepsMut<QueryC>, Env) -> AnyResult<T>,
866        ExecC: DeserializeOwned,
867    {
868        let contract = self.load_contract(storage, &address)?;
869        let handler = self
870            .codes
871            .get(&contract.code_id)
872            .ok_or(Error::UnregisteredCodeId(contract.code_id))?;
873
874        // We don't actually need a transaction here, as it is already embedded in a transactional.
875        // execute_submsg or App.execute_multi.
876        // However, we need to get write and read access to the same storage in two different objects,
877        // and this is the only way I know how to do so.
878        transactional(storage, |write_cache, read_store| {
879            let mut contract_storage = self.contract_storage(write_cache, &address);
880            let querier = RouterQuerier::new(router, api, read_store, block);
881            let env = self.get_env(address, block);
882
883            let deps = DepsMut {
884                storage: contract_storage.as_mut(),
885                api: api.deref(),
886                querier: QuerierWrapper::new(&querier),
887            };
888            action(handler, deps, env)
889        })
890    }
891
892    pub fn save_contract(
893        &self,
894        storage: &mut dyn Storage,
895        address: &Addr,
896        contract: &ContractData,
897    ) -> AnyResult<()> {
898        CONTRACTS
899            .save(&mut prefixed(storage, NAMESPACE_WASM), address, contract)
900            .map_err(Into::into)
901    }
902
903    /// Returns the number of all contract instances.
904    fn instance_count(&self, storage: &dyn Storage) -> usize {
905        CONTRACTS
906            .range_raw(
907                &prefixed_read(storage, NAMESPACE_WASM),
908                None,
909                None,
910                Order::Ascending,
911            )
912            .count()
913    }
914}
915
916// TODO: replace with code in utils
917
918#[derive(Clone, PartialEq, Message)]
919struct InstantiateResponse {
920    #[prost(string, tag = "1")]
921    pub address: ::prost::alloc::string::String,
922    #[prost(bytes, tag = "2")]
923    pub data: ::prost::alloc::vec::Vec<u8>,
924}
925
926// TODO: encode helpers in utils
927fn instantiate_response(data: Option<Binary>, contact_address: &Addr) -> Binary {
928    let data = data.unwrap_or_default().to_vec();
929    let init_data = InstantiateResponse {
930        address: contact_address.into(),
931        data,
932    };
933    let mut new_data = Vec::<u8>::with_capacity(init_data.encoded_len());
934    // the data must encode successfully
935    init_data.encode(&mut new_data).unwrap();
936    new_data.into()
937}
938
939#[derive(Clone, PartialEq, Message)]
940struct ExecuteResponse {
941    #[prost(bytes, tag = "1")]
942    pub data: ::prost::alloc::vec::Vec<u8>,
943}
944
945// empty return if no data present in original
946fn execute_response(data: Option<Binary>) -> Option<Binary> {
947    data.map(|d| {
948        let exec_data = ExecuteResponse { data: d.to_vec() };
949        let mut new_data = Vec::<u8>::with_capacity(exec_data.encoded_len());
950        // the data must encode successfully
951        exec_data.encode(&mut new_data).unwrap();
952        new_data.into()
953    })
954}
955
956#[cfg(test)]
957mod test {
958    use cosmwasm_std::testing::{mock_env, mock_info, MockApi, MockQuerier, MockStorage};
959    use cosmwasm_std::{
960        coin, from_slice, to_vec, BankMsg, CanonicalAddr, Coin, CosmosMsg, Empty, GovMsg, IbcMsg,
961        IbcQuery, StdError,
962    };
963
964    use crate::app::Router;
965    use crate::bank::BankKeeper;
966    use crate::module::FailingModule;
967    use crate::staking::{DistributionKeeper, StakeKeeper};
968    use crate::stargate::StargateKeeper;
969    use crate::test_helpers::contracts::{caller, error, payout};
970    use crate::test_helpers::EmptyMsg;
971    use crate::transactions::StorageTransaction;
972
973    use super::*;
974
975    /// Type alias for default build `Router` to make its reference in typical scenario
976    type BasicRouter<ExecC = Empty, QueryC = Empty> = Router<
977        BankKeeper,
978        FailingModule<ExecC, QueryC, Empty>,
979        WasmKeeper<ExecC, QueryC>,
980        StakeKeeper,
981        DistributionKeeper,
982        FailingModule<IbcMsg, IbcQuery, Empty>,
983        FailingModule<GovMsg, Empty, Empty>,
984        StargateKeeper<ExecC, QueryC>,
985    >;
986
987    fn mock_router() -> BasicRouter {
988        Router {
989            wasm: WasmKeeper::new(),
990            bank: BankKeeper::new(),
991            custom: FailingModule::new(),
992            staking: StakeKeeper::new(),
993            distribution: DistributionKeeper::new(),
994            ibc: FailingModule::new(),
995            gov: FailingModule::new(),
996            stargate: StargateKeeper::new(),
997        }
998    }
999
1000    #[test]
1001    fn register_contract() {
1002        let api = MockApi::default();
1003        let mut wasm_storage = MockStorage::new();
1004        let mut keeper = WasmKeeper::new();
1005        let block = mock_env().block;
1006        let code_id = keeper.store_code(error::contract(false));
1007
1008        transactional(&mut wasm_storage, |cache, _| {
1009            // cannot register contract with unregistered codeId
1010            keeper.register_contract(
1011                &api,
1012                cache,
1013                code_id + 1,
1014                Addr::unchecked("foobar"),
1015                Addr::unchecked("admin"),
1016                "label".to_owned(),
1017                1000,
1018            )
1019        })
1020        .unwrap_err();
1021
1022        let contract_addr = transactional(&mut wasm_storage, |cache, _| {
1023            // we can register a new instance of this code
1024            keeper.register_contract(
1025                &api,
1026                cache,
1027                code_id,
1028                Addr::unchecked("foobar"),
1029                Addr::unchecked("admin"),
1030                "label".to_owned(),
1031                1000,
1032            )
1033        })
1034        .unwrap();
1035
1036        // verify contract data are as expected
1037        let contract_data = keeper.load_contract(&wasm_storage, &contract_addr).unwrap();
1038
1039        assert_eq!(
1040            contract_data,
1041            ContractData {
1042                code_id,
1043                creator: Addr::unchecked("foobar"),
1044                admin: Some(Addr::unchecked("admin")),
1045                label: "label".to_owned(),
1046                created: 1000,
1047            }
1048        );
1049
1050        let err = transactional(&mut wasm_storage, |cache, _| {
1051            // now, we call this contract and see the error message from the contract
1052            let info = mock_info("foobar", &[]);
1053            keeper.call_instantiate(
1054                contract_addr.clone(),
1055                &api,
1056                cache,
1057                &mock_router(),
1058                &block,
1059                info,
1060                b"{}".to_vec(),
1061            )
1062        })
1063        .unwrap_err();
1064
1065        // StdError from contract_error auto-converted to string
1066        assert_eq!(
1067            StdError::generic_err("Init failed"),
1068            err.downcast().unwrap()
1069        );
1070
1071        let err = transactional(&mut wasm_storage, |cache, _| {
1072            // and the error for calling an unregistered contract
1073            let info = mock_info("foobar", &[]);
1074            keeper.call_instantiate(
1075                Addr::unchecked("unregistered"),
1076                &api,
1077                cache,
1078                &mock_router(),
1079                &block,
1080                info,
1081                b"{}".to_vec(),
1082            )
1083        })
1084        .unwrap_err();
1085
1086        // Default error message from router when not found
1087        assert_eq!(
1088            StdError::not_found("apollo_cw_multi_test::wasm::ContractData"),
1089            err.downcast().unwrap()
1090        );
1091    }
1092
1093    #[test]
1094    fn query_contract_into() {
1095        let api = MockApi::default();
1096        let mut keeper = WasmKeeper::<Empty, Empty>::new();
1097        let block = mock_env().block;
1098        let code_id = keeper.store_code(payout::contract());
1099
1100        let mut wasm_storage = MockStorage::new();
1101
1102        let contract_addr = keeper
1103            .register_contract(
1104                &api,
1105                &mut wasm_storage,
1106                code_id,
1107                Addr::unchecked("foobar"),
1108                Addr::unchecked("admin"),
1109                "label".to_owned(),
1110                1000,
1111            )
1112            .unwrap();
1113
1114        let querier: MockQuerier<Empty> = MockQuerier::new(&[]);
1115        let query = WasmQuery::ContractInfo {
1116            contract_addr: contract_addr.to_string(),
1117        };
1118        let info = keeper
1119            .query(&api, &wasm_storage, &querier, &block, query)
1120            .unwrap();
1121
1122        let mut expected = ContractInfoResponse::default();
1123        expected.code_id = code_id as u64;
1124        expected.creator = "foobar".to_string();
1125        expected.admin = Some("admin".to_owned());
1126        assert_eq!(expected, from_slice(&info).unwrap());
1127    }
1128
1129    #[test]
1130    fn can_dump_raw_wasm_state() {
1131        let api = MockApi::default();
1132        let mut keeper = WasmKeeper::<Empty, Empty>::new();
1133        let block = mock_env().block;
1134        let code_id = keeper.store_code(payout::contract());
1135
1136        let mut wasm_storage = MockStorage::new();
1137
1138        let contract_addr = keeper
1139            .register_contract(
1140                &api,
1141                &mut wasm_storage,
1142                code_id,
1143                Addr::unchecked("foobar"),
1144                Addr::unchecked("admin"),
1145                "label".to_owned(),
1146                1000,
1147            )
1148            .unwrap();
1149
1150        // make a contract with state
1151        let payout = coin(1500, "mlg");
1152        let msg = payout::InstantiateMessage {
1153            payout: payout.clone(),
1154        };
1155        keeper
1156            .call_instantiate(
1157                contract_addr.clone(),
1158                &api,
1159                &mut wasm_storage,
1160                &mock_router(),
1161                &block,
1162                mock_info("foobar", &[]),
1163                to_vec(&msg).unwrap(),
1164            )
1165            .unwrap();
1166
1167        // dump state
1168        let state = keeper.dump_wasm_raw(&wasm_storage, &contract_addr);
1169        assert_eq!(state.len(), 2);
1170        // check contents
1171        let (k, v) = &state[0];
1172        assert_eq!(k.as_slice(), b"count");
1173        let count: u32 = from_slice(v).unwrap();
1174        assert_eq!(count, 1);
1175        let (k, v) = &state[1];
1176        assert_eq!(k.as_slice(), b"payout");
1177        let stored_pay: payout::InstantiateMessage = from_slice(v).unwrap();
1178        assert_eq!(stored_pay.payout, payout);
1179    }
1180
1181    #[test]
1182    fn contract_send_coins() {
1183        let api = MockApi::default();
1184        let mut keeper = WasmKeeper::new();
1185        let block = mock_env().block;
1186        let code_id = keeper.store_code(payout::contract());
1187
1188        let mut wasm_storage = MockStorage::new();
1189        let mut cache = StorageTransaction::new(&wasm_storage);
1190
1191        let contract_addr = keeper
1192            .register_contract(
1193                &api,
1194                &mut cache,
1195                code_id,
1196                Addr::unchecked("foobar"),
1197                None,
1198                "label".to_owned(),
1199                1000,
1200            )
1201            .unwrap();
1202
1203        let payout = coin(100, "TGD");
1204
1205        // init the contract
1206        let info = mock_info("foobar", &[]);
1207        let init_msg = to_vec(&payout::InstantiateMessage {
1208            payout: payout.clone(),
1209        })
1210        .unwrap();
1211        let res = keeper
1212            .call_instantiate(
1213                contract_addr.clone(),
1214                &api,
1215                &mut cache,
1216                &mock_router(),
1217                &block,
1218                info,
1219                init_msg,
1220            )
1221            .unwrap();
1222        assert_eq!(0, res.messages.len());
1223
1224        // execute the contract
1225        let info = mock_info("foobar", &[]);
1226        let res = keeper
1227            .call_execute(
1228                &api,
1229                &mut cache,
1230                contract_addr.clone(),
1231                &mock_router(),
1232                &block,
1233                info,
1234                b"{}".to_vec(),
1235            )
1236            .unwrap();
1237        assert_eq!(1, res.messages.len());
1238        match &res.messages[0].msg {
1239            CosmosMsg::Bank(BankMsg::Send { to_address, amount }) => {
1240                assert_eq!(to_address.as_str(), "foobar");
1241                assert_eq!(amount.as_slice(), &[payout.clone()]);
1242            }
1243            m => panic!("Unexpected message {:?}", m),
1244        }
1245
1246        // and flush before query
1247        cache.prepare().commit(&mut wasm_storage);
1248
1249        // query the contract
1250        let query = to_vec(&payout::QueryMsg::Payout {}).unwrap();
1251        let querier: MockQuerier<Empty> = MockQuerier::new(&[]);
1252        let data = keeper
1253            .query_smart(contract_addr, &api, &wasm_storage, &querier, &block, query)
1254            .unwrap();
1255        let res: payout::InstantiateMessage = from_slice(&data).unwrap();
1256        assert_eq!(res.payout, payout);
1257    }
1258
1259    fn assert_payout(
1260        router: &WasmKeeper<Empty, Empty>,
1261        storage: &mut dyn Storage,
1262        contract_addr: &Addr,
1263        payout: &Coin,
1264    ) {
1265        let api = MockApi::default();
1266        let info = mock_info("silly", &[]);
1267        let res = router
1268            .call_execute(
1269                &api,
1270                storage,
1271                contract_addr.clone(),
1272                &mock_router(),
1273                &mock_env().block,
1274                info,
1275                b"{}".to_vec(),
1276            )
1277            .unwrap();
1278        assert_eq!(1, res.messages.len());
1279        match &res.messages[0].msg {
1280            CosmosMsg::Bank(BankMsg::Send { to_address, amount }) => {
1281                assert_eq!(to_address.as_str(), "silly");
1282                assert_eq!(amount.as_slice(), &[payout.clone()]);
1283            }
1284            m => panic!("Unexpected message {:?}", m),
1285        }
1286    }
1287
1288    fn assert_no_contract(storage: &dyn Storage, contract_addr: &Addr) {
1289        let contract = CONTRACTS.may_load(storage, contract_addr).unwrap();
1290        assert!(contract.is_none(), "{:?}", contract_addr);
1291    }
1292
1293    #[test]
1294    fn multi_level_wasm_cache() {
1295        let api = MockApi::default();
1296        let mut keeper = WasmKeeper::new();
1297        let block = mock_env().block;
1298        let code_id = keeper.store_code(payout::contract());
1299
1300        let mut wasm_storage = MockStorage::new();
1301
1302        let payout1 = coin(100, "TGD");
1303
1304        // set contract 1 and commit (on router)
1305        let contract1 = transactional(&mut wasm_storage, |cache, _| {
1306            let contract = keeper
1307                .register_contract(
1308                    &api,
1309                    cache,
1310                    code_id,
1311                    Addr::unchecked("foobar"),
1312                    None,
1313                    "".to_string(),
1314                    1000,
1315                )
1316                .unwrap();
1317            let info = mock_info("foobar", &[]);
1318            let init_msg = to_vec(&payout::InstantiateMessage {
1319                payout: payout1.clone(),
1320            })
1321            .unwrap();
1322            keeper
1323                .call_instantiate(
1324                    contract.clone(),
1325                    &api,
1326                    cache,
1327                    &mock_router(),
1328                    &block,
1329                    info,
1330                    init_msg,
1331                )
1332                .unwrap();
1333
1334            Ok(contract)
1335        })
1336        .unwrap();
1337
1338        let payout2 = coin(50, "BTC");
1339        let payout3 = coin(1234, "ATOM");
1340
1341        // create a new cache and check we can use contract 1
1342        let (contract2, contract3) = transactional(&mut wasm_storage, |cache, wasm_reader| {
1343            assert_payout(&keeper, cache, &contract1, &payout1);
1344
1345            // create contract 2 and use it
1346            let contract2 = keeper
1347                .register_contract(
1348                    &api,
1349                    cache,
1350                    code_id,
1351                    Addr::unchecked("foobar"),
1352                    None,
1353                    "".to_owned(),
1354                    1000,
1355                )
1356                .unwrap();
1357            let info = mock_info("foobar", &[]);
1358            let init_msg = to_vec(&payout::InstantiateMessage {
1359                payout: payout2.clone(),
1360            })
1361            .unwrap();
1362            let _res = keeper
1363                .call_instantiate(
1364                    contract2.clone(),
1365                    &api,
1366                    cache,
1367                    &mock_router(),
1368                    &block,
1369                    info,
1370                    init_msg,
1371                )
1372                .unwrap();
1373            assert_payout(&keeper, cache, &contract2, &payout2);
1374
1375            // create a level2 cache and check we can use contract 1 and contract 2
1376            let contract3 = transactional(cache, |cache2, read| {
1377                assert_payout(&keeper, cache2, &contract1, &payout1);
1378                assert_payout(&keeper, cache2, &contract2, &payout2);
1379
1380                // create a contract on level 2
1381                let contract3 = keeper
1382                    .register_contract(
1383                        &api,
1384                        cache2,
1385                        code_id,
1386                        Addr::unchecked("foobar"),
1387                        None,
1388                        "".to_owned(),
1389                        1000,
1390                    )
1391                    .unwrap();
1392                let info = mock_info("johnny", &[]);
1393                let init_msg = to_vec(&payout::InstantiateMessage {
1394                    payout: payout3.clone(),
1395                })
1396                .unwrap();
1397                let _res = keeper
1398                    .call_instantiate(
1399                        contract3.clone(),
1400                        &api,
1401                        cache2,
1402                        &mock_router(),
1403                        &block,
1404                        info,
1405                        init_msg,
1406                    )
1407                    .unwrap();
1408                assert_payout(&keeper, cache2, &contract3, &payout3);
1409
1410                // ensure first cache still doesn't see this contract
1411                assert_no_contract(read, &contract3);
1412                Ok(contract3)
1413            })
1414            .unwrap();
1415
1416            // after applying transaction, all contracts present on cache
1417            assert_payout(&keeper, cache, &contract1, &payout1);
1418            assert_payout(&keeper, cache, &contract2, &payout2);
1419            assert_payout(&keeper, cache, &contract3, &payout3);
1420
1421            // but not yet the root router
1422            assert_no_contract(wasm_reader, &contract1);
1423            assert_no_contract(wasm_reader, &contract2);
1424            assert_no_contract(wasm_reader, &contract3);
1425
1426            Ok((contract2, contract3))
1427        })
1428        .unwrap();
1429
1430        // ensure that it is now applied to the router
1431        assert_payout(&keeper, &mut wasm_storage, &contract1, &payout1);
1432        assert_payout(&keeper, &mut wasm_storage, &contract2, &payout2);
1433        assert_payout(&keeper, &mut wasm_storage, &contract3, &payout3);
1434    }
1435
1436    fn assert_admin(
1437        storage: &dyn Storage,
1438        keeper: &WasmKeeper<Empty, Empty>,
1439        contract_addr: &impl ToString,
1440        admin: Option<Addr>,
1441    ) {
1442        let api = MockApi::default();
1443        let querier: MockQuerier<Empty> = MockQuerier::new(&[]);
1444        // query
1445        let data = keeper
1446            .query(
1447                &api,
1448                storage,
1449                &querier,
1450                &mock_env().block,
1451                WasmQuery::ContractInfo {
1452                    contract_addr: contract_addr.to_string(),
1453                },
1454            )
1455            .unwrap();
1456        let res: ContractInfoResponse = from_slice(&data).unwrap();
1457        assert_eq!(res.admin, admin.as_ref().map(Addr::to_string));
1458    }
1459
1460    #[test]
1461    fn update_clear_admin_works() {
1462        let api = MockApi::default();
1463        let mut keeper = WasmKeeper::new();
1464        let block = mock_env().block;
1465        let code_id = keeper.store_code(caller::contract());
1466
1467        let mut wasm_storage = MockStorage::new();
1468
1469        let admin: Addr = Addr::unchecked("admin");
1470        let new_admin: Addr = Addr::unchecked("new_admin");
1471        let normal_user: Addr = Addr::unchecked("normal_user");
1472
1473        let contract_addr = keeper
1474            .register_contract(
1475                &api,
1476                &mut wasm_storage,
1477                code_id,
1478                Addr::unchecked("creator"),
1479                admin.clone(),
1480                "label".to_owned(),
1481                1000,
1482            )
1483            .unwrap();
1484
1485        // init the contract
1486        let info = mock_info("admin", &[]);
1487        let init_msg = to_vec(&EmptyMsg {}).unwrap();
1488        let res = keeper
1489            .call_instantiate(
1490                contract_addr.clone(),
1491                &api,
1492                &mut wasm_storage,
1493                &mock_router(),
1494                &block,
1495                info,
1496                init_msg,
1497            )
1498            .unwrap();
1499        assert_eq!(0, res.messages.len());
1500
1501        assert_admin(&wasm_storage, &keeper, &contract_addr, Some(admin.clone()));
1502
1503        // non-admin should not be allowed to become admin on their own
1504        keeper
1505            .execute_wasm(
1506                &api,
1507                &mut wasm_storage,
1508                &mock_router(),
1509                &block,
1510                normal_user.clone(),
1511                WasmMsg::UpdateAdmin {
1512                    contract_addr: contract_addr.to_string(),
1513                    admin: normal_user.to_string(),
1514                },
1515            )
1516            .unwrap_err();
1517
1518        // should still be admin
1519        assert_admin(&wasm_storage, &keeper, &contract_addr, Some(admin.clone()));
1520
1521        // admin should be allowed to transfers adminship
1522        let res = keeper
1523            .execute_wasm(
1524                &api,
1525                &mut wasm_storage,
1526                &mock_router(),
1527                &block,
1528                admin,
1529                WasmMsg::UpdateAdmin {
1530                    contract_addr: contract_addr.to_string(),
1531                    admin: new_admin.to_string(),
1532                },
1533            )
1534            .unwrap();
1535        assert_eq!(res.events.len(), 0);
1536
1537        // new_admin should now be admin
1538        assert_admin(
1539            &wasm_storage,
1540            &keeper,
1541            &contract_addr,
1542            Some(new_admin.clone()),
1543        );
1544
1545        // new_admin should now be able to clear to admin
1546        let res = keeper
1547            .execute_wasm(
1548                &api,
1549                &mut wasm_storage,
1550                &mock_router(),
1551                &block,
1552                new_admin,
1553                WasmMsg::ClearAdmin {
1554                    contract_addr: contract_addr.to_string(),
1555                },
1556            )
1557            .unwrap();
1558        assert_eq!(res.events.len(), 0);
1559
1560        // should have no admin now
1561        assert_admin(&wasm_storage, &keeper, &contract_addr, None);
1562    }
1563
1564    #[test]
1565    fn by_default_uses_simple_address_generator() {
1566        let api = MockApi::default();
1567        let mut keeper = WasmKeeper::<Empty, Empty>::new();
1568        let code_id = keeper.store_code(payout::contract());
1569
1570        let mut wasm_storage = MockStorage::new();
1571
1572        let contract_addr = keeper
1573            .register_contract(
1574                &api,
1575                &mut wasm_storage,
1576                code_id,
1577                Addr::unchecked("foobar"),
1578                Addr::unchecked("admin"),
1579                "label".to_owned(),
1580                1000,
1581            )
1582            .unwrap();
1583
1584        assert_eq!(
1585            "contract0", contract_addr,
1586            "default address generator returned incorrect address"
1587        )
1588    }
1589
1590    struct TestAddressGenerator {
1591        address: Addr,
1592        predictable_address: Addr,
1593    }
1594
1595    impl AddressGenerator for TestAddressGenerator {
1596        fn contract_address(
1597            &self,
1598            _api: &dyn Api,
1599            _storage: &mut dyn Storage,
1600            _code_id: u64,
1601            _instance_id: u64,
1602        ) -> AnyResult<Addr> {
1603            Ok(self.address.clone())
1604        }
1605
1606        fn predictable_contract_address(
1607            &self,
1608            _api: &dyn Api,
1609            _storage: &mut dyn Storage,
1610            _code_id: u64,
1611            _instance_id: u64,
1612            _checksum: &[u8],
1613            _creator: &CanonicalAddr,
1614            _salt: &[u8],
1615        ) -> AnyResult<Addr> {
1616            Ok(self.predictable_address.clone())
1617        }
1618    }
1619
1620    #[test]
1621    fn can_use_custom_address_generator() {
1622        let api = MockApi::default();
1623        let expected_addr = Addr::unchecked("new_test_addr");
1624        let expected_predictable_addr = api.addr_make("predictable_address");
1625        let mut keeper: WasmKeeper<Empty, Empty> =
1626            WasmKeeper::new_with_custom_address_generator(TestAddressGenerator {
1627                address: expected_addr.clone(),
1628                predictable_address: expected_predictable_addr.clone(),
1629            });
1630        let code_id = keeper.store_code(payout::contract());
1631
1632        let mut wasm_storage = MockStorage::new();
1633
1634        let contract_addr = keeper
1635            .register_contract(
1636                &api,
1637                &mut wasm_storage,
1638                code_id,
1639                Addr::unchecked("foobar"),
1640                Addr::unchecked("admin"),
1641                "label".to_owned(),
1642                1000,
1643            )
1644            .unwrap();
1645
1646        assert_eq!(
1647            expected_addr, contract_addr,
1648            "custom address generator returned incorrect address"
1649        )
1650    }
1651}