program-tests/* refactor to include TestIndexerExtensions

Lightprotocol · Jan 9, 2025 · a20ca37 · a20ca37
1 parent f08e48c
commit a20ca37
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 10 deletions.
diff --git a/program-tests/utils/src/assert_compressed_tx.rs b/program-tests/utils/src/assert_compressed_tx.rs
@@ -2,7 +2,6 @@ use account_compression::{state::QueueAccount, StateMerkleTreeAccount};
 use anchor_lang::Discriminator;
 use forester_utils::{
     get_concurrent_merkle_tree, get_hash_set,
-    indexer::{Indexer, StateMerkleTreeAccounts},
     AccountZeroCopy,
 };
 use light_batched_merkle_tree::{
@@ -19,8 +18,10 @@ use light_system_program::sdk::{
 use num_bigint::BigUint;
 use num_traits::FromBytes;
 use solana_sdk::{account::ReadableAccount, pubkey::Pubkey};
+use light_client::indexer::{Indexer, StateMerkleTreeAccounts};
+use light_program_test::indexer::TestIndexerExtensions;
 
-pub struct AssertCompressedTransactionInputs<'a, R: RpcConnection, I: Indexer<R>> {
+pub struct AssertCompressedTransactionInputs<'a, R: RpcConnection, I: Indexer<R> + TestIndexerExtensions<R>> {
     pub rpc: &'a mut R,
     pub test_indexer: &'a mut I,
     pub output_compressed_accounts: &'a [CompressedAccount],
@@ -48,7 +49,7 @@ pub struct AssertCompressedTransactionInputs<'a, R: RpcConnection, I: Indexer<R>
 /// 5. Merkle tree was updated correctly
 /// 6. TODO: Fees have been paid (after fee refactor)
 /// 7. Check compression amount was transferred
-pub async fn assert_compressed_transaction<R: RpcConnection, I: Indexer<R>>(
+pub async fn assert_compressed_transaction<R: RpcConnection, I: Indexer<R> + TestIndexerExtensions<R>>(
     input: AssertCompressedTransactionInputs<'_, R, I>,
 ) {
     // CHECK 1
@@ -316,7 +317,7 @@ pub struct MerkleTreeTestSnapShot {
 /// Asserts:
 /// 1. The root has been updated
 /// 2. The next index has been updated
-pub async fn assert_merkle_tree_after_tx<R: RpcConnection, I: Indexer<R>>(
+pub async fn assert_merkle_tree_after_tx<R: RpcConnection, I: Indexer<R> + TestIndexerExtensions<R>>(
     rpc: &mut R,
     snapshots: &[MerkleTreeTestSnapShot],
     test_indexer: &mut I,

diff --git a/program-tests/utils/src/assert_token_tx.rs b/program-tests/utils/src/assert_token_tx.rs
@@ -6,6 +6,7 @@ use light_system_program::sdk::{
 };
 use solana_sdk::{program_pack::Pack, pubkey::Pubkey};
 use light_client::indexer::Indexer;
+use light_program_test::indexer::TestIndexerExtensions;
 use light_sdk::token::TokenDataWithMerkleContext;
 use crate::assert_compressed_tx::{
     assert_merkle_tree_after_tx, assert_nullifiers_exist_in_hash_sets,
@@ -21,7 +22,7 @@ use crate::assert_compressed_tx::{
 /// 6. Check compression amount was transferred (outside of this function)
 /// No addresses in token transactions
 #[allow(clippy::too_many_arguments)]
-pub async fn assert_transfer<R: RpcConnection, I: Indexer<R>>(
+pub async fn assert_transfer<R: RpcConnection, I: Indexer<R> + TestIndexerExtensions<R>>(
     context: &mut R,
     test_indexer: &mut I,
     out_compressed_accounts: &[TokenTransferOutputData],
@@ -78,7 +79,7 @@ pub async fn assert_transfer<R: RpcConnection, I: Indexer<R>>(
     );
 }
 
-pub fn assert_compressed_token_accounts<R: RpcConnection, I: Indexer<R>>(
+pub fn assert_compressed_token_accounts<R: RpcConnection, I: Indexer<R> + TestIndexerExtensions<R>>(
     test_indexer: &mut I,
     out_compressed_accounts: &[TokenTransferOutputData],
     lamports: Option<Vec<Option<u64>>>,
@@ -188,7 +189,7 @@ pub fn assert_compressed_token_accounts<R: RpcConnection, I: Indexer<R>>(
 }
 
 #[allow(clippy::too_many_arguments)]
-pub async fn assert_mint_to<'a, R: RpcConnection, I: Indexer<R>>(
+pub async fn assert_mint_to<'a, R: RpcConnection, I: Indexer<R> + TestIndexerExtensions<R>>(
     rpc: &mut R,
     test_indexer: &'a mut I,
     recipients: &[Pubkey],

diff --git a/program-tests/utils/src/e2e_test_env.rs b/program-tests/utils/src/e2e_test_env.rs
@@ -145,6 +145,7 @@ use crate::{
     },
     test_forester::{empty_address_queue_test, nullify_compressed_accounts},
 };
+use crate::conversions::sdk_to_program_compressed_account_with_merkle_context;
 
 pub struct User {
     pub keypair: Keypair,
@@ -2199,7 +2200,10 @@ where
     ) -> Vec<CompressedAccountWithMerkleContext> {
         let input_compressed_accounts = self
             .indexer
-            .get_compressed_accounts_by_owner(&self.users[user_index].keypair.pubkey());
+            .get_compressed_accounts_with_merkle_context_by_owner(&self.users[user_index].keypair.pubkey())
+            .into_iter()
+            .map(sdk_to_program_compressed_account_with_merkle_context)
+            .collect::<Vec<_>>();
         if input_compressed_accounts.is_empty() {
             return vec![];
         }
@@ -2236,7 +2240,10 @@ where
         &self,
         pubkey: &Pubkey,
     ) -> Vec<CompressedAccountWithMerkleContext> {
-        self.indexer.get_compressed_accounts_by_owner(pubkey)
+        self.indexer.get_compressed_accounts_with_merkle_context_by_owner(pubkey)
+            .into_iter()
+            .map(sdk_to_program_compressed_account_with_merkle_context)
+            .collect()
     }
 
     pub fn get_merkle_tree_pubkeys(&mut self, num: u64) -> Vec<Pubkey> {

diff --git a/program-tests/utils/src/spl.rs b/program-tests/utils/src/spl.rs
@@ -1566,7 +1566,7 @@ pub async fn freeze_or_thaw_test<R: RpcConnection, const FREEZE: bool, I: Indexe
         created_output_accounts
             .iter()
             .map(|x| x.compressed_account.clone())
-            .sdk_to_program_compressed_account_with_merkle_context()
+            .map(sdk_to_program_compressed_account_with_merkle_context)
             .collect::<Vec<_>>()
             .as_slice(),
         change_lamports,