cw20_base/
enumerable.rs

1use cosmwasm_std::{Deps, Order, StdResult};
2use cw20::{
3    AllAccountsResponse, AllAllowancesResponse, AllSpenderAllowancesResponse, AllowanceInfo,
4    SpenderAllowanceInfo,
5};
6
7use crate::state::{ALLOWANCES, ALLOWANCES_SPENDER, BALANCES};
8use cw_storage_plus::Bound;
9
10// settings for pagination
11const MAX_LIMIT: u32 = 30;
12const DEFAULT_LIMIT: u32 = 10;
13
14pub fn query_owner_allowances(
15    deps: Deps,
16    owner: String,
17    start_after: Option<String>,
18    limit: Option<u32>,
19) -> StdResult<AllAllowancesResponse> {
20    let owner_addr = deps.api.addr_validate(&owner)?;
21    let limit = limit.unwrap_or(DEFAULT_LIMIT).min(MAX_LIMIT) as usize;
22    let start = start_after.map(|s| Bound::ExclusiveRaw(s.into_bytes()));
23
24    let allowances = ALLOWANCES
25        .prefix(&owner_addr)
26        .range(deps.storage, start, None, Order::Ascending)
27        .take(limit)
28        .map(|item| {
29            item.map(|(addr, allow)| AllowanceInfo {
30                spender: addr.into(),
31                allowance: allow.allowance,
32                expires: allow.expires,
33            })
34        })
35        .collect::<StdResult<_>>()?;
36    Ok(AllAllowancesResponse { allowances })
37}
38
39pub fn query_spender_allowances(
40    deps: Deps,
41    spender: String,
42    start_after: Option<String>,
43    limit: Option<u32>,
44) -> StdResult<AllSpenderAllowancesResponse> {
45    let spender_addr = deps.api.addr_validate(&spender)?;
46    let limit = limit.unwrap_or(DEFAULT_LIMIT).min(MAX_LIMIT) as usize;
47    let start = start_after.map(|s| Bound::ExclusiveRaw(s.into_bytes()));
48
49    let allowances = ALLOWANCES_SPENDER
50        .prefix(&spender_addr)
51        .range(deps.storage, start, None, Order::Ascending)
52        .take(limit)
53        .map(|item| {
54            item.map(|(addr, allow)| SpenderAllowanceInfo {
55                owner: addr.into(),
56                allowance: allow.allowance,
57                expires: allow.expires,
58            })
59        })
60        .collect::<StdResult<_>>()?;
61    Ok(AllSpenderAllowancesResponse { allowances })
62}
63
64pub fn query_all_accounts(
65    deps: Deps,
66    start_after: Option<String>,
67    limit: Option<u32>,
68) -> StdResult<AllAccountsResponse> {
69    let limit = limit.unwrap_or(DEFAULT_LIMIT).min(MAX_LIMIT) as usize;
70    let start = start_after.map(|s| Bound::ExclusiveRaw(s.into()));
71
72    let accounts = BALANCES
73        .keys(deps.storage, start, None, Order::Ascending)
74        .take(limit)
75        .map(|item| item.map(Into::into))
76        .collect::<StdResult<_>>()?;
77
78    Ok(AllAccountsResponse { accounts })
79}
80
81#[cfg(test)]
82mod tests {
83    use super::*;
84
85    use cosmwasm_std::testing::{mock_dependencies_with_balance, mock_env, mock_info};
86    use cosmwasm_std::{coins, from_json, DepsMut, Uint128};
87    use cw20::{Cw20Coin, Expiration, TokenInfoResponse};
88
89    use crate::contract::{execute, instantiate, query, query_token_info};
90    use crate::msg::{ExecuteMsg, InstantiateMsg, QueryMsg};
91
92    // this will set up the instantiation for other tests
93    fn do_instantiate(mut deps: DepsMut, addr: &str, amount: Uint128) -> TokenInfoResponse {
94        let instantiate_msg = InstantiateMsg {
95            name: "Auto Gen".to_string(),
96            symbol: "AUTO".to_string(),
97            decimals: 3,
98            initial_balances: vec![Cw20Coin {
99                address: addr.into(),
100                amount,
101            }],
102            mint: None,
103            marketing: None,
104        };
105        let info = mock_info("creator", &[]);
106        let env = mock_env();
107        instantiate(deps.branch(), env, info, instantiate_msg).unwrap();
108        query_token_info(deps.as_ref()).unwrap()
109    }
110
111    #[test]
112    fn query_all_owner_allowances_works() {
113        let mut deps = mock_dependencies_with_balance(&coins(2, "token"));
114
115        let owner = deps.api.addr_make("owner").to_string();
116        // these are in alphabetical order same than insert order
117        let spender1 = deps.api.addr_make("earlier").to_string();
118        let spender2 = deps.api.addr_make("later").to_string();
119
120        let info = mock_info(owner.as_ref(), &[]);
121        let env = mock_env();
122        do_instantiate(deps.as_mut(), &owner, Uint128::new(12340000));
123
124        // no allowance to start
125        let allowances = query_owner_allowances(deps.as_ref(), owner.clone(), None, None).unwrap();
126        assert_eq!(allowances.allowances, vec![]);
127
128        // set allowance with height expiration
129        let allow1 = Uint128::new(7777);
130        let expires = Expiration::AtHeight(123_456);
131        let msg = ExecuteMsg::IncreaseAllowance {
132            spender: spender1.clone(),
133            amount: allow1,
134            expires: Some(expires),
135        };
136        execute(deps.as_mut(), env.clone(), info.clone(), msg).unwrap();
137
138        // set allowance with no expiration
139        let allow2 = Uint128::new(54321);
140        let msg = ExecuteMsg::IncreaseAllowance {
141            spender: spender2.clone(),
142            amount: allow2,
143            expires: None,
144        };
145        execute(deps.as_mut(), env, info, msg).unwrap();
146
147        // query list gets 2
148        let allowances = query_owner_allowances(deps.as_ref(), owner.clone(), None, None).unwrap();
149        assert_eq!(allowances.allowances.len(), 2);
150
151        // first one is spender1 (order of CanonicalAddr uncorrelated with String)
152        let allowances =
153            query_owner_allowances(deps.as_ref(), owner.clone(), None, Some(1)).unwrap();
154        assert_eq!(allowances.allowances.len(), 1);
155        let allow = &allowances.allowances[0];
156        assert_eq!(&allow.spender, &spender1);
157        assert_eq!(&allow.expires, &expires);
158        assert_eq!(&allow.allowance, &allow1);
159
160        // next one is spender2
161        let allowances = query_owner_allowances(
162            deps.as_ref(),
163            owner,
164            Some(allow.spender.clone()),
165            Some(10000),
166        )
167        .unwrap();
168        assert_eq!(allowances.allowances.len(), 1);
169        let allow = &allowances.allowances[0];
170        assert_eq!(&allow.spender, &spender2);
171        assert_eq!(&allow.expires, &Expiration::Never {});
172        assert_eq!(&allow.allowance, &allow2);
173    }
174
175    #[test]
176    fn query_all_spender_allowances_works() {
177        let mut deps = mock_dependencies_with_balance(&coins(2, "token"));
178
179        let mut addresses = [
180            deps.api.addr_make("owner1").to_string(),
181            deps.api.addr_make("owner2").to_string(),
182            deps.api.addr_make("spender").to_string(),
183        ];
184        addresses.sort();
185
186        // these are in alphabetical order same than insert order
187        let [owner1, owner2, spender] = addresses;
188
189        let info = mock_info(owner1.as_ref(), &[]);
190        let env = mock_env();
191        do_instantiate(deps.as_mut(), &owner1, Uint128::new(12340000));
192
193        // no allowance to start
194        let allowances =
195            query_spender_allowances(deps.as_ref(), spender.clone(), None, None).unwrap();
196        assert_eq!(allowances.allowances, vec![]);
197
198        // set allowance with height expiration
199        let allow1 = Uint128::new(7777);
200        let expires = Expiration::AtHeight(123_456);
201        let msg = ExecuteMsg::IncreaseAllowance {
202            spender: spender.clone(),
203            amount: allow1,
204            expires: Some(expires),
205        };
206        execute(deps.as_mut(), env, info, msg).unwrap();
207
208        // set allowance with no expiration, from the other owner
209        let info = mock_info(owner2.as_ref(), &[]);
210        let env = mock_env();
211        do_instantiate(deps.as_mut(), &owner2, Uint128::new(12340000));
212
213        let allow2 = Uint128::new(54321);
214        let msg = ExecuteMsg::IncreaseAllowance {
215            spender: spender.clone(),
216            amount: allow2,
217            expires: None,
218        };
219        execute(deps.as_mut(), env.clone(), info, msg).unwrap();
220
221        // query list gets both
222        let msg = QueryMsg::AllSpenderAllowances {
223            spender: spender.clone(),
224            start_after: None,
225            limit: None,
226        };
227        let allowances: AllSpenderAllowancesResponse =
228            from_json(query(deps.as_ref(), env.clone(), msg).unwrap()).unwrap();
229        assert_eq!(allowances.allowances.len(), 2);
230
231        // one is owner1 (order of CanonicalAddr uncorrelated with String)
232        let msg = QueryMsg::AllSpenderAllowances {
233            spender: spender.clone(),
234            start_after: None,
235            limit: Some(1),
236        };
237        let allowances: AllSpenderAllowancesResponse =
238            from_json(query(deps.as_ref(), env.clone(), msg).unwrap()).unwrap();
239        assert_eq!(allowances.allowances.len(), 1);
240        let allow = &allowances.allowances[0];
241        assert_eq!(&allow.owner, &owner1);
242        assert_eq!(&allow.expires, &expires);
243        assert_eq!(&allow.allowance, &allow1);
244
245        // other one is owner2
246        let msg = QueryMsg::AllSpenderAllowances {
247            spender,
248            start_after: Some(owner1),
249            limit: Some(10000),
250        };
251        let allowances: AllSpenderAllowancesResponse =
252            from_json(query(deps.as_ref(), env, msg).unwrap()).unwrap();
253        assert_eq!(allowances.allowances.len(), 1);
254        let allow = &allowances.allowances[0];
255        assert_eq!(&allow.owner, &owner2);
256        assert_eq!(&allow.expires, &Expiration::Never {});
257        assert_eq!(&allow.allowance, &allow2);
258    }
259
260    #[test]
261    fn query_all_accounts_works() {
262        let mut deps = mock_dependencies_with_balance(&coins(2, "token"));
263
264        // insert order and lexicographical order are different
265        let acct1 = deps.api.addr_make("acct1").to_string();
266        let acct2 = deps.api.addr_make("zebra").to_string();
267        let acct3 = deps.api.addr_make("nice").to_string();
268        let acct4 = deps.api.addr_make("aaardvark").to_string();
269
270        let mut expected_order = [acct1.clone(), acct2.clone(), acct3.clone(), acct4.clone()];
271        expected_order.sort();
272
273        do_instantiate(deps.as_mut(), &acct1, Uint128::new(12340000));
274
275        // put money everywhere (to create balanaces)
276        let info = mock_info(acct1.as_ref(), &[]);
277        let env = mock_env();
278        execute(
279            deps.as_mut(),
280            env.clone(),
281            info.clone(),
282            ExecuteMsg::Transfer {
283                recipient: acct2,
284                amount: Uint128::new(222222),
285            },
286        )
287        .unwrap();
288        execute(
289            deps.as_mut(),
290            env.clone(),
291            info.clone(),
292            ExecuteMsg::Transfer {
293                recipient: acct3,
294                amount: Uint128::new(333333),
295            },
296        )
297        .unwrap();
298        execute(
299            deps.as_mut(),
300            env,
301            info,
302            ExecuteMsg::Transfer {
303                recipient: acct4,
304                amount: Uint128::new(444444),
305            },
306        )
307        .unwrap();
308
309        // make sure we get the proper results
310        let accounts = query_all_accounts(deps.as_ref(), None, None).unwrap();
311        assert_eq!(accounts.accounts, expected_order);
312
313        // let's do pagination
314        let accounts = query_all_accounts(deps.as_ref(), None, Some(2)).unwrap();
315        assert_eq!(accounts.accounts, expected_order[0..2].to_vec());
316
317        let accounts =
318            query_all_accounts(deps.as_ref(), Some(accounts.accounts[1].clone()), Some(1)).unwrap();
319        assert_eq!(accounts.accounts, expected_order[2..3].to_vec());
320
321        let accounts =
322            query_all_accounts(deps.as_ref(), Some(accounts.accounts[0].clone()), Some(777))
323                .unwrap();
324        assert_eq!(accounts.accounts, expected_order[3..].to_vec());
325    }
326}
cw20_base/enumerable.rs

cw20_base/
enumerable.rs
