1use litesvm::LiteSVM;
7use solana_program::pubkey::Pubkey;
8use litesvm_token::spl_token;
9use solana_program_pack::Pack;
10
11pub trait AssertionHelpers {
13 fn assert_account_closed(&self, pubkey: &Pubkey);
25
26 fn assert_account_exists(&self, pubkey: &Pubkey);
38
39 fn assert_token_balance(&self, token_account: &Pubkey, expected: u64);
51
52 fn assert_sol_balance(&self, pubkey: &Pubkey, expected: u64);
64
65 fn assert_mint_supply(&self, mint: &Pubkey, expected: u64);
77
78 fn assert_account_owner(&self, account: &Pubkey, expected_owner: &Pubkey);
91
92 fn assert_account_data_len(&self, account: &Pubkey, expected_len: usize);
104}
105
106impl AssertionHelpers for LiteSVM {
107 fn assert_account_closed(&self, pubkey: &Pubkey) {
108 let account = self.get_account(pubkey);
109 assert!(
110 account.is_none() || (account.as_ref().unwrap().lamports == 0 && account.as_ref().unwrap().data.is_empty()),
111 "Expected account {} to be closed, but it exists with {} lamports and {} bytes of data",
112 pubkey,
113 account.as_ref().map_or(0, |a| a.lamports),
114 account.as_ref().map_or(0, |a| a.data.len())
115 );
116 }
117
118 fn assert_account_exists(&self, pubkey: &Pubkey) {
119 let account = self.get_account(pubkey);
120 assert!(
121 account.is_some(),
122 "Expected account {} to exist, but it doesn't",
123 pubkey
124 );
125 }
126
127 fn assert_token_balance(&self, token_account: &Pubkey, expected: u64) {
128 let account = self
129 .get_account(token_account)
130 .unwrap_or_else(|| panic!("Token account {} not found", token_account));
131
132 let token_data = spl_token::state::Account::unpack(&account.data)
133 .unwrap_or_else(|_| panic!("Failed to unpack token account {}", token_account));
134
135 assert_eq!(
136 token_data.amount, expected,
137 "Token balance mismatch for account {}. Expected: {}, Actual: {}",
138 token_account, expected, token_data.amount
139 );
140 }
141
142 fn assert_sol_balance(&self, pubkey: &Pubkey, expected: u64) {
143 let account = self.get_account(pubkey);
144 let actual = account.map_or(0, |a| a.lamports);
145 assert_eq!(
146 actual, expected,
147 "SOL balance mismatch for account {}. Expected: {}, Actual: {}",
148 pubkey, expected, actual
149 );
150 }
151
152 fn assert_mint_supply(&self, mint: &Pubkey, expected: u64) {
153 let account = self
154 .get_account(mint)
155 .unwrap_or_else(|| panic!("Mint {} not found", mint));
156
157 let mint_data = spl_token::state::Mint::unpack(&account.data)
158 .unwrap_or_else(|_| panic!("Failed to unpack mint {}", mint));
159
160 assert_eq!(
161 mint_data.supply, expected,
162 "Mint supply mismatch for {}. Expected: {}, Actual: {}",
163 mint, expected, mint_data.supply
164 );
165 }
166
167 fn assert_account_owner(&self, account: &Pubkey, expected_owner: &Pubkey) {
168 let acc = self
169 .get_account(account)
170 .unwrap_or_else(|| panic!("Account {} not found", account));
171
172 assert_eq!(
173 &acc.owner, expected_owner,
174 "Account owner mismatch for {}. Expected: {}, Actual: {}",
175 account, expected_owner, acc.owner
176 );
177 }
178
179 fn assert_account_data_len(&self, account: &Pubkey, expected_len: usize) {
180 let acc = self
181 .get_account(account)
182 .unwrap_or_else(|| panic!("Account {} not found", account));
183
184 assert_eq!(
185 acc.data.len(),
186 expected_len,
187 "Account data length mismatch for {}. Expected: {}, Actual: {}",
188 account,
189 expected_len,
190 acc.data.len()
191 );
192 }
193}
194
195#[cfg(test)]
196mod tests {
197 use super::*;
198 use crate::test_helpers::TestHelpers;
199 use solana_sdk::signature::Signer;
200
201 #[test]
202 fn test_assert_account_closed_nonexistent() {
203 let svm = LiteSVM::new();
204 let nonexistent_account = Pubkey::new_unique();
205
206 svm.assert_account_closed(&nonexistent_account);
208 }
209
210 #[test]
211 fn test_assert_account_exists() {
212 let mut svm = LiteSVM::new();
213 let account = svm.create_funded_account(1_000_000_000).unwrap();
214
215 svm.assert_account_exists(&account.pubkey());
217 }
218
219 #[test]
220 #[should_panic(expected = "Expected account")]
221 fn test_assert_account_exists_fails() {
222 let svm = LiteSVM::new();
223 let nonexistent = Pubkey::new_unique();
224
225 svm.assert_account_exists(&nonexistent);
227 }
228
229 #[test]
230 fn test_assert_token_balance() {
231 let mut svm = LiteSVM::new();
232 let authority = svm.create_funded_account(10_000_000_000).unwrap();
233 let mint = svm.create_token_mint(&authority, 9).unwrap();
234 let token_account = svm
235 .create_associated_token_account(&mint.pubkey(), &authority)
236 .unwrap();
237
238 let amount = 1_000_000;
240 svm.mint_to(&mint.pubkey(), &token_account, &authority, amount)
241 .unwrap();
242
243 svm.assert_token_balance(&token_account, amount);
245 }
246
247 #[test]
248 #[should_panic(expected = "Token balance mismatch")]
249 fn test_assert_token_balance_fails() {
250 let mut svm = LiteSVM::new();
251 let authority = svm.create_funded_account(10_000_000_000).unwrap();
252 let mint = svm.create_token_mint(&authority, 9).unwrap();
253 let token_account = svm
254 .create_associated_token_account(&mint.pubkey(), &authority)
255 .unwrap();
256
257 svm.mint_to(&mint.pubkey(), &token_account, &authority, 1000)
259 .unwrap();
260
261 svm.assert_token_balance(&token_account, 2000);
263 }
264
265 #[test]
266 fn test_assert_sol_balance() {
267 let mut svm = LiteSVM::new();
268 let lamports = 5_000_000_000;
269 let account = svm.create_funded_account(lamports).unwrap();
270
271 svm.assert_sol_balance(&account.pubkey(), lamports);
273 }
274
275 #[test]
276 #[should_panic(expected = "SOL balance mismatch")]
277 fn test_assert_sol_balance_fails() {
278 let mut svm = LiteSVM::new();
279 let account = svm.create_funded_account(1_000_000_000).unwrap();
280
281 svm.assert_sol_balance(&account.pubkey(), 2_000_000_000);
283 }
284
285 #[test]
286 fn test_assert_sol_balance_zero_for_nonexistent() {
287 let svm = LiteSVM::new();
288 let nonexistent = Pubkey::new_unique();
289
290 svm.assert_sol_balance(&nonexistent, 0);
292 }
293
294 #[test]
295 fn test_assert_mint_supply() {
296 let mut svm = LiteSVM::new();
297 let authority = svm.create_funded_account(10_000_000_000).unwrap();
298 let mint = svm.create_token_mint(&authority, 9).unwrap();
299 let token_account = svm
300 .create_associated_token_account(&mint.pubkey(), &authority)
301 .unwrap();
302
303 let amount = 5_000_000;
305 svm.mint_to(&mint.pubkey(), &token_account, &authority, amount)
306 .unwrap();
307
308 svm.assert_mint_supply(&mint.pubkey(), amount);
310 }
311
312 #[test]
313 #[should_panic(expected = "Mint supply mismatch")]
314 fn test_assert_mint_supply_fails() {
315 let mut svm = LiteSVM::new();
316 let authority = svm.create_funded_account(10_000_000_000).unwrap();
317 let mint = svm.create_token_mint(&authority, 9).unwrap();
318 let token_account = svm
319 .create_associated_token_account(&mint.pubkey(), &authority)
320 .unwrap();
321
322 svm.mint_to(&mint.pubkey(), &token_account, &authority, 100)
324 .unwrap();
325
326 svm.assert_mint_supply(&mint.pubkey(), 200);
328 }
329
330 #[test]
331 fn test_assert_mint_supply_zero_for_new_mint() {
332 let mut svm = LiteSVM::new();
333 let authority = svm.create_funded_account(10_000_000_000).unwrap();
334 let mint = svm.create_token_mint(&authority, 9).unwrap();
335
336 svm.assert_mint_supply(&mint.pubkey(), 0);
338 }
339
340 #[test]
341 fn test_assert_account_owner() {
342 let mut svm = LiteSVM::new();
343 let owner = svm.create_funded_account(10_000_000_000).unwrap();
344 let mint = svm.create_token_mint(&owner, 9).unwrap();
345
346 svm.assert_account_owner(&mint.pubkey(), &spl_token::id());
348 }
349
350 #[test]
351 #[should_panic(expected = "Account owner mismatch")]
352 fn test_assert_account_owner_fails() {
353 let mut svm = LiteSVM::new();
354 let owner = svm.create_funded_account(10_000_000_000).unwrap();
355 let mint = svm.create_token_mint(&owner, 9).unwrap();
356
357 let wrong_owner = Pubkey::new_unique();
359 svm.assert_account_owner(&mint.pubkey(), &wrong_owner);
360 }
361
362 #[test]
363 fn test_assert_account_data_len() {
364 let mut svm = LiteSVM::new();
365 let owner = svm.create_funded_account(10_000_000_000).unwrap();
366 let mint = svm.create_token_mint(&owner, 9).unwrap();
367
368 svm.assert_account_data_len(&mint.pubkey(), 82);
370 }
371
372 #[test]
373 #[should_panic(expected = "Account data length mismatch")]
374 fn test_assert_account_data_len_fails() {
375 let mut svm = LiteSVM::new();
376 let owner = svm.create_funded_account(10_000_000_000).unwrap();
377 let mint = svm.create_token_mint(&owner, 9).unwrap();
378
379 svm.assert_account_data_len(&mint.pubkey(), 100);
381 }
382
383 #[test]
384 fn test_assert_account_data_len_token_account() {
385 let mut svm = LiteSVM::new();
386 let owner = svm.create_funded_account(10_000_000_000).unwrap();
387 let mint = svm.create_token_mint(&owner, 9).unwrap();
388 let token_account = svm.create_token_account(&mint.pubkey(), &owner).unwrap();
389
390 svm.assert_account_data_len(&token_account.pubkey(), 165);
392 }
393}