apollo_cw_multi_test/
stargate.rs

1use std::collections::HashMap;
2
3use anyhow::bail;
4use cosmwasm_std::{Addr, Api, Binary, BlockInfo, CosmosMsg, CustomQuery, Empty, Querier, Storage};
5
6use crate::{AppResponse, CosmosRouter};
7
8// TODO: turn into extensions of Fn trait
9pub trait StargateQueryHandler {
10    fn stargate_query(
11        &self,
12        api: &dyn Api,
13        storage: &dyn Storage,
14        querier: &dyn Querier,
15        block: &BlockInfo,
16        request: StargateMsg,
17    ) -> anyhow::Result<Binary>;
18
19    fn register_queries(&'static self, keeper: &mut StargateKeeper<Empty, Empty>);
20}
21
22pub trait StargateMessageHandler<ExecC, QueryC: CustomQuery> {
23    fn execute(
24        &self,
25        api: &dyn Api,
26        storage: &mut dyn Storage,
27        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
28        block: &BlockInfo,
29        sender: Addr,
30        msg: StargateMsg,
31    ) -> anyhow::Result<AppResponse>;
32
33    fn register_msgs(&'static self, keeper: &mut StargateKeeper<Empty, Empty>);
34}
35
36pub struct StargateKeeper<ExecC, QueryC> {
37    messages: HashMap<String, Box<dyn StargateMessageHandler<ExecC, QueryC>>>,
38    queries: HashMap<String, Box<dyn StargateQueryHandler>>,
39}
40
41impl<'a, ExecC, QueryC> StargateKeeper<ExecC, QueryC> {
42    pub fn new() -> Self {
43        Self {
44            messages: HashMap::new(),
45            queries: HashMap::new(),
46        }
47    }
48
49    pub fn register_msg(
50        &mut self,
51        type_url: &str,
52        handler: Box<dyn StargateMessageHandler<ExecC, QueryC>>,
53    ) {
54        self.messages.insert(type_url.to_string(), handler);
55    }
56
57    pub fn register_query(&mut self, type_url: &str, handler: Box<dyn StargateQueryHandler>) {
58        self.queries.insert(type_url.to_string(), handler);
59    }
60}
61
62pub struct StargateMsg {
63    pub type_url: String,
64    pub value: Binary,
65}
66
67impl From<StargateMsg> for CosmosMsg {
68    fn from(msg: StargateMsg) -> Self {
69        CosmosMsg::Stargate {
70            type_url: msg.type_url,
71            value: msg.value,
72        }
73    }
74}
75
76pub trait Stargate<ExecC, QueryC> {
77    fn execute(
78        &self,
79        api: &dyn Api,
80        storage: &mut dyn Storage,
81        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
82        block: &BlockInfo,
83        sender: Addr,
84        msg: StargateMsg,
85    ) -> anyhow::Result<crate::AppResponse>;
86
87    fn sudo(
88        &self,
89        api: &dyn Api,
90        storage: &mut dyn Storage,
91        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
92        block: &cosmwasm_std::BlockInfo,
93        msg: Empty,
94    ) -> anyhow::Result<crate::AppResponse>;
95
96    fn query(
97        &self,
98        api: &dyn cosmwasm_std::Api,
99        storage: &dyn cosmwasm_std::Storage,
100        querier: &dyn cosmwasm_std::Querier,
101        block: &cosmwasm_std::BlockInfo,
102        request: StargateMsg,
103    ) -> anyhow::Result<cosmwasm_std::Binary>;
104}
105
106impl<'a, ExecC, QueryC: CustomQuery> Stargate<ExecC, QueryC> for StargateKeeper<ExecC, QueryC> {
107    fn execute(
108        &self,
109        api: &dyn Api,
110        storage: &mut dyn Storage,
111        router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
112        block: &BlockInfo,
113        sender: Addr,
114        msg: StargateMsg,
115    ) -> anyhow::Result<crate::AppResponse> {
116        match self.messages.get(&msg.type_url.to_string()) {
117            Some(handler) => handler.execute(api, storage, router, block, sender, msg),
118            None => bail!("Unsupported stargate message: {}", msg.type_url),
119        }
120    }
121
122    fn sudo(
123        &self,
124        _api: &dyn Api,
125        _storage: &mut dyn Storage,
126        _router: &dyn CosmosRouter<ExecC = ExecC, QueryC = QueryC>,
127        _block: &cosmwasm_std::BlockInfo,
128        _msg: Empty,
129    ) -> anyhow::Result<crate::AppResponse> {
130        bail!("StargateKeeper does not support sudo")
131    }
132
133    fn query(
134        &self,
135        api: &dyn cosmwasm_std::Api,
136        storage: &dyn cosmwasm_std::Storage,
137        querier: &dyn cosmwasm_std::Querier,
138        block: &cosmwasm_std::BlockInfo,
139        request: StargateMsg,
140    ) -> anyhow::Result<cosmwasm_std::Binary> {
141        match self.queries.get(&request.type_url.to_string()) {
142            Some(handler) => handler.stargate_query(api, storage, querier, block, request),
143            None => bail!("Unsupported stargate query: {}", request.type_url),
144        }
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use std::str::FromStr;
151
152    use anyhow::Ok;
153    use cosmwasm_std::{coin, from_binary, to_binary, Coin, CosmosMsg, Event, QueryRequest};
154    use osmosis_std::types::cosmos::bank::v1beta1::{
155        QueryAllBalancesRequest, QueryBalanceRequest, QuerySupplyOfRequest,
156    };
157
158    use crate::{BasicAppBuilder, Executor};
159
160    use super::*;
161
162    #[derive(Clone)]
163    struct FooHandler;
164    impl StargateMessageHandler<Empty, Empty> for FooHandler {
165        fn execute(
166            &self,
167            _api: &dyn Api,
168            _storage: &mut dyn Storage,
169            _router: &dyn CosmosRouter<ExecC = Empty, QueryC = Empty>,
170            _block: &BlockInfo,
171            _sender: Addr,
172            msg: StargateMsg,
173        ) -> anyhow::Result<AppResponse> {
174            let mut res = AppResponse::default();
175            let num: u64 = from_binary(&msg.value)?;
176            res.events
177                .push(Event::new("foo").add_attribute("bar", num.to_string()));
178            Ok(res)
179        }
180
181        fn register_msgs(&'static self, keeper: &mut StargateKeeper<Empty, Empty>) {
182            keeper.register_msg("foo", Box::new(self.clone()))
183        }
184    }
185    const FOO_HANDLER: FooHandler = FooHandler;
186
187    #[derive(Clone)]
188    struct FooQueryHandler;
189    impl StargateQueryHandler for FooQueryHandler {
190        fn stargate_query(
191            &self,
192            _api: &dyn Api,
193            _storage: &dyn Storage,
194            _querier: &dyn Querier,
195            _block: &BlockInfo,
196            msg: StargateMsg,
197        ) -> anyhow::Result<Binary> {
198            let num: u64 = from_binary(&msg.value)?;
199            let bin = to_binary(&format!("bar{:?}", num)).unwrap();
200            Ok(bin)
201        }
202
203        fn register_queries(&'static self, keeper: &mut StargateKeeper<Empty, Empty>) {
204            keeper.register_query("foo", Box::new(self.clone()))
205        }
206    }
207    const FOO_QUERY_HANDLER: FooQueryHandler = FooQueryHandler;
208
209    #[derive(Clone)]
210    struct BarHandler;
211    impl StargateMessageHandler<Empty, Empty> for BarHandler {
212        fn execute(
213            &self,
214            api: &dyn Api,
215            storage: &mut dyn Storage,
216            router: &dyn CosmosRouter<ExecC = Empty, QueryC = Empty>,
217            block: &BlockInfo,
218            _sender: Addr,
219            msg: StargateMsg,
220        ) -> anyhow::Result<AppResponse> {
221            let query_res = router.query(
222                api,
223                storage,
224                block,
225                QueryRequest::Stargate {
226                    path: "foo".to_string(),
227                    data: msg.value,
228                },
229            )?;
230
231            let mut res = AppResponse::default();
232            res.data = Some(query_res);
233
234            Ok(res)
235        }
236
237        fn register_msgs(&'static self, keeper: &mut StargateKeeper<Empty, Empty>) {
238            keeper.register_msg("bar", Box::new(self.clone()))
239        }
240    }
241    const BAR_HANDLER: BarHandler = BarHandler;
242
243    #[test]
244    fn new_stargate_keeper() {
245        StargateKeeper::<Empty, Empty>::new();
246    }
247
248    #[test]
249    fn register_and_call_stargate_msg() {
250        let mut stargate_keeper = StargateKeeper::new();
251        stargate_keeper.register_msg("foo", Box::new(FOO_HANDLER));
252
253        let app = BasicAppBuilder::<Empty, Empty>::new()
254            .with_stargate(stargate_keeper)
255            .build(|_, _, _| {});
256
257        let res = app
258            .execute(
259                Addr::unchecked("unchecked"),
260                CosmosMsg::Stargate {
261                    type_url: "foo".to_string(),
262                    value: to_binary(&1337u64).unwrap(),
263                },
264            )
265            .unwrap();
266
267        res.assert_event(&Event::new("foo").add_attribute("bar", "1337"));
268    }
269
270    #[test]
271    fn register_and_call_stargate_query() {
272        let mut stargate_keeper = StargateKeeper::new();
273        stargate_keeper.register_query("foo", Box::new(FOO_QUERY_HANDLER));
274
275        let app = BasicAppBuilder::<Empty, Empty>::new()
276            .with_stargate(stargate_keeper)
277            .build(|_, _, _| {});
278
279        let querier = app.wrap();
280
281        let res: String = querier
282            .query(&QueryRequest::Stargate {
283                path: "foo".to_string(),
284                data: to_binary(&1337u64).unwrap(),
285            })
286            .unwrap();
287
288        assert_eq!(res, "bar1337".to_string());
289    }
290
291    #[test]
292    fn query_inside_execution() {
293        let mut stargate_keeper = StargateKeeper::new();
294        stargate_keeper.register_msg("bar", Box::new(BAR_HANDLER));
295        stargate_keeper.register_query("foo", Box::new(FOO_QUERY_HANDLER));
296
297        let app = BasicAppBuilder::<Empty, Empty>::new()
298            .with_stargate(stargate_keeper)
299            .build(|_, _, _| {});
300
301        let res = app
302            .execute(
303                Addr::unchecked("unchecked"),
304                CosmosMsg::Stargate {
305                    type_url: "bar".to_string(),
306                    value: to_binary(&1337u64).unwrap(),
307                },
308            )
309            .unwrap();
310
311        let x: String = from_binary(&res.data.unwrap()).unwrap();
312        assert_eq!(x, "bar1337");
313    }
314
315    #[test]
316    fn query_bank_module_via_stargate() {
317        let stargate_keeper = StargateKeeper::new();
318
319        let owner = Addr::unchecked("owner");
320        let init_funds = vec![coin(20, "btc"), coin(100, "eth")];
321
322        let app = BasicAppBuilder::<Empty, Empty>::new()
323            .with_stargate(stargate_keeper)
324            .build(|router, _, storage| {
325                router
326                    .bank
327                    .init_balance(storage, &owner, init_funds.clone())
328                    .unwrap();
329            });
330
331        let querier = app.wrap();
332
333        // QueryAllBalancesRequest
334        let res = QueryAllBalancesRequest {
335            address: owner.to_string(),
336            pagination: None,
337        }
338        .query(&querier)
339        .unwrap();
340        let blances: Vec<Coin> = res
341            .balances
342            .into_iter()
343            .map(|c| Coin::new(u128::from_str(&c.amount).unwrap(), c.denom))
344            .collect();
345        assert_eq!(blances, init_funds);
346
347        // QueryBalanceRequest
348        let res = QueryBalanceRequest {
349            address: owner.to_string(),
350            denom: "eth".to_string(),
351        }
352        .query(&querier)
353        .unwrap();
354        let balance = res.balance.unwrap();
355        assert_eq!(balance.amount, init_funds[1].amount.to_string());
356        assert_eq!(balance.denom, init_funds[1].denom);
357
358        // QueryTotalSupplyRequest
359        let res = QuerySupplyOfRequest {
360            denom: "eth".to_string(),
361        };
362        let res = res.query(&querier).unwrap();
363        let supply = res.amount.unwrap();
364        assert_eq!(supply.amount, init_funds[1].amount.to_string());
365        assert_eq!(supply.denom, init_funds[1].denom);
366    }
367}
apollo_cw_multi_test/stargate.rs

apollo_cw_multi_test/
stargate.rs
