litesvm_utils/
assertions.rs

1//! Assertion helpers for testing account states
2//!
3//! This module provides convenient assertion methods for verifying
4//! account states in tests.
5
6use litesvm::LiteSVM;
7use solana_program::pubkey::Pubkey;
8use litesvm_token::spl_token;
9use solana_program_pack::Pack;
10
11/// Assertion helper methods for LiteSVM
12pub trait AssertionHelpers {
13    /// Assert that an account is closed (doesn't exist or has 0 lamports and 0 data)
14    ///
15    /// # Example
16    /// ```no_run
17    /// # use litesvm_utils::AssertionHelpers;
18    /// # use litesvm::LiteSVM;
19    /// # use solana_program::pubkey::Pubkey;
20    /// # let svm = LiteSVM::new();
21    /// # let account = Pubkey::new_unique();
22    /// svm.assert_account_closed(&account);
23    /// ```
24    fn assert_account_closed(&self, pubkey: &Pubkey);
25
26    /// Assert that an account exists
27    ///
28    /// # Example
29    /// ```no_run
30    /// # use litesvm_utils::AssertionHelpers;
31    /// # use litesvm::LiteSVM;
32    /// # use solana_program::pubkey::Pubkey;
33    /// # let svm = LiteSVM::new();
34    /// # let account = Pubkey::new_unique();
35    /// svm.assert_account_exists(&account);
36    /// ```
37    fn assert_account_exists(&self, pubkey: &Pubkey);
38
39    /// Assert token account balance
40    ///
41    /// # Example
42    /// ```no_run
43    /// # use litesvm_utils::AssertionHelpers;
44    /// # use litesvm::LiteSVM;
45    /// # use solana_program::pubkey::Pubkey;
46    /// # let svm = LiteSVM::new();
47    /// # let token_account = Pubkey::new_unique();
48    /// svm.assert_token_balance(&token_account, 1_000_000_000); // 1 token with 9 decimals
49    /// ```
50    fn assert_token_balance(&self, token_account: &Pubkey, expected: u64);
51
52    /// Assert SOL balance
53    ///
54    /// # Example
55    /// ```no_run
56    /// # use litesvm_utils::AssertionHelpers;
57    /// # use litesvm::LiteSVM;
58    /// # use solana_program::pubkey::Pubkey;
59    /// # let svm = LiteSVM::new();
60    /// # let account = Pubkey::new_unique();
61    /// svm.assert_sol_balance(&account, 1_000_000_000); // 1 SOL
62    /// ```
63    fn assert_sol_balance(&self, pubkey: &Pubkey, expected: u64);
64
65    /// Assert token mint supply
66    ///
67    /// # Example
68    /// ```no_run
69    /// # use litesvm_utils::AssertionHelpers;
70    /// # use litesvm::LiteSVM;
71    /// # use solana_program::pubkey::Pubkey;
72    /// # let svm = LiteSVM::new();
73    /// # let mint = Pubkey::new_unique();
74    /// svm.assert_mint_supply(&mint, 1_000_000_000);
75    /// ```
76    fn assert_mint_supply(&self, mint: &Pubkey, expected: u64);
77
78    /// Assert that an account is owned by a specific program
79    ///
80    /// # Example
81    /// ```no_run
82    /// # use litesvm_utils::AssertionHelpers;
83    /// # use litesvm::LiteSVM;
84    /// # use solana_program::pubkey::Pubkey;
85    /// # let svm = LiteSVM::new();
86    /// # let account = Pubkey::new_unique();
87    /// # let owner = Pubkey::new_unique();
88    /// svm.assert_account_owner(&account, &owner);
89    /// ```
90    fn assert_account_owner(&self, account: &Pubkey, expected_owner: &Pubkey);
91
92    /// Assert that an account has a specific data length
93    ///
94    /// # Example
95    /// ```no_run
96    /// # use litesvm_utils::AssertionHelpers;
97    /// # use litesvm::LiteSVM;
98    /// # use solana_program::pubkey::Pubkey;
99    /// # let svm = LiteSVM::new();
100    /// # let account = Pubkey::new_unique();
101    /// svm.assert_account_data_len(&account, 100);
102    /// ```
103    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        // Should not panic for non-existent account
207        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        // Should not panic for existing account
216        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        // Should panic for non-existent account
226        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        // Mint tokens
239        let amount = 1_000_000;
240        svm.mint_to(&mint.pubkey(), &token_account, &authority, amount)
241            .unwrap();
242
243        // Should not panic with correct balance
244        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        // Mint 1000 tokens
258        svm.mint_to(&mint.pubkey(), &token_account, &authority, 1000)
259            .unwrap();
260
261        // Should panic when expecting wrong balance
262        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        // Should not panic with correct balance
272        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        // Should panic when expecting wrong balance
282        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        // Should not panic when expecting 0 for non-existent account
291        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        // Mint tokens
304        let amount = 5_000_000;
305        svm.mint_to(&mint.pubkey(), &token_account, &authority, amount)
306            .unwrap();
307
308        // Should not panic with correct supply
309        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        // Mint 100 tokens
323        svm.mint_to(&mint.pubkey(), &token_account, &authority, 100)
324            .unwrap();
325
326        // Should panic when expecting wrong supply
327        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        // New mint should have 0 supply
337        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        // Mint should be owned by token program
347        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        // Should panic when expecting wrong owner
358        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        // Mint account data is 82 bytes
369        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        // Should panic when expecting wrong length
380        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        // Token account data is 165 bytes
391        svm.assert_account_data_len(&token_account.pubkey(), 165);
392    }
393}