multiversx_sc_scenario/scenario/model/account_data/
account.rs

1use crate::{
2    scenario::model::{
3        AddressValue, BigUintValue, BytesKey, BytesValue, Esdt, EsdtObject, U64Value,
4    },
5    scenario_format::{
6        interpret_trait::{InterpretableFrom, InterpreterContext, IntoRaw},
7        serde_raw::AccountRaw,
8    },
9};
10
11use std::collections::BTreeMap;
12
13#[derive(Debug, Default, Clone)]
14pub struct Account {
15    pub comment: Option<String>,
16    pub nonce: Option<U64Value>,
17    pub balance: Option<BigUintValue>,
18    pub esdt: BTreeMap<BytesKey, Esdt>,
19    pub username: Option<BytesValue>,
20    pub storage: BTreeMap<BytesKey, BytesValue>,
21    pub code: Option<BytesValue>,
22    pub code_metadata: Option<BytesValue>,
23    pub owner: Option<AddressValue>,
24    pub developer_rewards: Option<BigUintValue>,
25}
26
27impl Account {
28    pub fn new() -> Self {
29        Self::default()
30    }
31
32    pub fn nonce<V>(mut self, nonce: V) -> Self
33    where
34        U64Value: From<V>,
35    {
36        self.nonce = Some(U64Value::from(nonce));
37        self
38    }
39
40    pub fn balance<V>(mut self, balance_expr: V) -> Self
41    where
42        BigUintValue: From<V>,
43    {
44        self.balance = Some(BigUintValue::from(balance_expr));
45        self
46    }
47
48    pub fn esdt_balance<K, V>(mut self, token_id_expr: K, balance_expr: V) -> Self
49    where
50        BytesKey: From<K>,
51        BigUintValue: From<V>,
52    {
53        let token_id = BytesKey::from(token_id_expr);
54        let esdt_data_ref = self.get_esdt_data_or_create(&token_id);
55        esdt_data_ref.set_balance(0u64, balance_expr);
56
57        self
58    }
59
60    pub fn esdt_nft_balance<K, N, V, T>(
61        mut self,
62        token_id_expr: K,
63        nonce_expr: N,
64        balance_expr: V,
65        opt_attributes_expr: Option<T>,
66    ) -> Self
67    where
68        N: Clone,
69        BytesKey: From<K>,
70        U64Value: From<N>,
71        BigUintValue: From<V>,
72        BytesValue: From<T>,
73    {
74        let token_id = BytesKey::from(token_id_expr);
75
76        let esdt_obj_ref = self
77            .get_esdt_data_or_create(&token_id)
78            .get_mut_esdt_object();
79        esdt_obj_ref.set_balance(nonce_expr.clone(), balance_expr);
80
81        if let Some(attributes_expr) = opt_attributes_expr {
82            esdt_obj_ref.set_token_attributes(nonce_expr, attributes_expr);
83        }
84
85        self
86    }
87
88    #[allow(clippy::too_many_arguments)]
89    pub fn esdt_nft_all_properties<K, N, V, T, A>(
90        mut self,
91        token_id_expr: K,
92        nonce_expr: N,
93        balance_expr: V,
94        opt_attributes_expr: Option<T>,
95        royalties_expr: N,
96        creator_expr: Option<A>,
97        hash_expr: Option<T>,
98        uris_expr: Vec<T>,
99    ) -> Self
100    where
101        BytesKey: From<K>,
102        U64Value: From<N>,
103        BigUintValue: From<V>,
104        BytesValue: From<T>,
105        AddressValue: From<A>,
106    {
107        let token_id = BytesKey::from(token_id_expr);
108
109        let esdt_obj_ref = self
110            .get_esdt_data_or_create(&token_id)
111            .get_mut_esdt_object();
112
113        esdt_obj_ref.set_token_all_properties(
114            nonce_expr,
115            balance_expr,
116            opt_attributes_expr,
117            royalties_expr,
118            creator_expr,
119            hash_expr,
120            uris_expr,
121        );
122
123        self
124    }
125
126    pub fn esdt_nft_last_nonce<K, N>(mut self, token_id_expr: K, last_nonce_expr: N) -> Self
127    where
128        BytesKey: From<K>,
129        U64Value: From<N>,
130    {
131        let token_id = BytesKey::from(token_id_expr);
132
133        let esdt_obj_ref = self
134            .get_esdt_data_or_create(&token_id)
135            .get_mut_esdt_object();
136        esdt_obj_ref.set_last_nonce(last_nonce_expr);
137
138        self
139    }
140
141    // TODO: Find a better way to pass roles
142    pub fn esdt_roles<K>(mut self, token_id_expr: K, roles: Vec<String>) -> Self
143    where
144        BytesKey: From<K>,
145    {
146        let token_id = BytesKey::from(token_id_expr);
147
148        let esdt_obj_ref = self
149            .get_esdt_data_or_create(&token_id)
150            .get_mut_esdt_object();
151        esdt_obj_ref.set_roles(roles);
152
153        self
154    }
155
156    fn get_esdt_data_or_create(&mut self, token_id: &BytesKey) -> &mut Esdt {
157        if !self.esdt.contains_key(token_id) {
158            let _ = self
159                .esdt
160                .insert(token_id.clone(), Esdt::Full(EsdtObject::default()));
161        }
162
163        self.esdt.get_mut(token_id).unwrap()
164    }
165
166    pub fn code<V>(mut self, code_expr: V) -> Self
167    where
168        BytesValue: From<V>,
169    {
170        self.code = Some(BytesValue::from(code_expr));
171        self
172    }
173
174    pub fn owner<V>(mut self, owner_expr: V) -> Self
175    where
176        AddressValue: From<V>,
177    {
178        self.owner = Some(AddressValue::from(owner_expr));
179        self
180    }
181}
182
183impl InterpretableFrom<AccountRaw> for Account {
184    fn interpret_from(from: AccountRaw, context: &InterpreterContext) -> Self {
185        Account {
186            comment: from.comment,
187            nonce: from.nonce.map(|n| U64Value::interpret_from(n, context)),
188            balance: from
189                .balance
190                .map(|b| BigUintValue::interpret_from(b, context)),
191            esdt: from
192                .esdt
193                .into_iter()
194                .map(|(k, v)| {
195                    (
196                        BytesKey::interpret_from(k, context),
197                        Esdt::interpret_from(v, context),
198                    )
199                })
200                .collect(),
201            username: from
202                .username
203                .map(|c| BytesValue::interpret_from(c, context)),
204            storage: from
205                .storage
206                .into_iter()
207                .map(|(k, v)| {
208                    (
209                        BytesKey::interpret_from(k, context),
210                        BytesValue::interpret_from(v, context),
211                    )
212                })
213                .collect(),
214            code: from.code.map(|c| BytesValue::interpret_from(c, context)),
215            code_metadata: from
216                .code_metadata
217                .map(|c| BytesValue::interpret_from(c, context)),
218            owner: from.owner.map(|v| AddressValue::interpret_from(v, context)),
219            developer_rewards: from
220                .developer_rewards
221                .map(|b| BigUintValue::interpret_from(b, context)),
222        }
223    }
224}
225
226impl IntoRaw<AccountRaw> for Account {
227    fn into_raw(self) -> AccountRaw {
228        AccountRaw {
229            comment: self.comment,
230            nonce: self.nonce.map(|n| n.original),
231            balance: self.balance.map(|n| n.original),
232            esdt: self
233                .esdt
234                .into_iter()
235                .map(|(k, v)| (k.original, v.into_raw()))
236                .collect(),
237            username: self.username.map(|n| n.original),
238            storage: self
239                .storage
240                .into_iter()
241                .map(|(key, value)| (key.original, value.original))
242                .collect(),
243            code: self.code.map(|n| n.original),
244            code_metadata: self.code_metadata.map(|n| n.original),
245            owner: self.owner.map(|n| n.original),
246            developer_rewards: self.developer_rewards.map(|n| n.original),
247        }
248    }
249}