light_program_test/program_test/
light_program_test.rs

1use std::fmt::{self, Debug, Formatter};
2
3use account_compression::{AddressMerkleTreeAccount, QueueAccount};
4use light_client::{
5    indexer::{AddressMerkleTreeAccounts, StateMerkleTreeAccounts},
6    rpc::{merkle_tree::MerkleTreeExt, RpcError},
7};
8use light_prover_client::prover::{spawn_prover, ProverConfig};
9use litesvm::LiteSVM;
10use solana_account::WritableAccount;
11use solana_sdk::signature::{Keypair, Signer};
12
13use crate::{
14    accounts::{
15        initialize::initialize_accounts, test_accounts::TestAccounts, test_keypairs::TestKeypairs,
16    },
17    indexer::TestIndexer,
18    program_test::TestRpc,
19    utils::setup_light_programs::setup_light_programs,
20    ProgramTestConfig,
21};
22
23pub struct LightProgramTest {
24    pub config: ProgramTestConfig,
25    pub context: LiteSVM,
26    pub indexer: Option<TestIndexer>,
27    pub test_accounts: TestAccounts,
28    pub payer: Keypair,
29}
30
31impl LightProgramTest {
32    /// Creates ProgramTestContext with light protocol and additional programs.
33    ///
34    /// Programs:
35    /// 1. light program
36    /// 2. account_compression program
37    /// 3. light_compressed_token program
38    /// 4. light_system_program program
39    ///
40    /// Light Protocol accounts:
41    /// 5. creates and initializes governance authority
42    /// 6. creates and initializes group authority
43    /// 7. registers the light_system_program program with the group authority
44    /// 8. initializes Merkle tree owned by
45    /// Note:
46    /// - registers a forester
47    /// - advances to the active phase slot 2
48    /// - active phase doesn't end
49    pub async fn new(config: ProgramTestConfig) -> Result<LightProgramTest, RpcError> {
50        let mut context = setup_light_programs(config.additional_programs.clone())?;
51        let payer = Keypair::new();
52        context
53            .airdrop(&payer.pubkey(), 100_000_000_000_000)
54            .expect("Payer airdrop failed.");
55        let mut context = Self {
56            context,
57            indexer: None,
58            test_accounts: TestAccounts::get_program_test_test_accounts(),
59            payer,
60            config: config.clone(),
61        };
62        let keypairs = TestKeypairs::program_test_default();
63
64        context
65            .context
66            .airdrop(&keypairs.governance_authority.pubkey(), 100_000_000_000_000)
67            .expect("governance_authority airdrop failed.");
68        context
69            .context
70            .airdrop(&keypairs.forester.pubkey(), 10_000_000_000)
71            .expect("forester airdrop failed.");
72
73        if !config.skip_protocol_init {
74            let restore_logs = context.config.no_logs;
75            if context.config.skip_startup_logs {
76                context.config.no_logs = true;
77            }
78            initialize_accounts(&mut context, &config, &keypairs).await?;
79            if context.config.skip_startup_logs {
80                context.config.no_logs = restore_logs;
81            }
82            let batch_size = config
83                .v2_state_tree_config
84                .as_ref()
85                .map(|config| config.output_queue_batch_size as usize);
86            let test_accounts = context.test_accounts.clone();
87            context.add_indexer(&test_accounts, batch_size).await?;
88
89            // TODO: add the same for v2 trees once we have grinded a mainnet keypair.
90            // ensure that address tree pubkey is amt1Ayt45jfbdw5YSo7iz6WZxUmnZsQTYXy82hVwyC2
91            {
92                let address_mt = context.test_accounts.v1_address_trees[0].merkle_tree;
93                let address_queue_pubkey = context.test_accounts.v1_address_trees[0].queue;
94                let mut account = context
95                    .context
96                    .get_account(&keypairs.address_merkle_tree.pubkey())
97                    .unwrap();
98                let merkle_tree_account = bytemuck::from_bytes_mut::<AddressMerkleTreeAccount>(
99                    &mut account.data_as_mut_slice()[8..AddressMerkleTreeAccount::LEN],
100                );
101                merkle_tree_account.metadata.associated_queue = address_queue_pubkey.into();
102                context.set_account(address_mt, account);
103
104                let mut account = context
105                    .context
106                    .get_account(&keypairs.address_merkle_tree_queue.pubkey())
107                    .unwrap();
108                let queue_account = bytemuck::from_bytes_mut::<QueueAccount>(
109                    &mut account.data_as_mut_slice()[8..QueueAccount::LEN],
110                );
111                queue_account.metadata.associated_merkle_tree = address_mt.into();
112                context.set_account(address_queue_pubkey, account);
113            }
114        }
115        // Will always start a prover server.
116        #[cfg(feature = "devenv")]
117        let prover_config = if config.prover_config.is_none() {
118            Some(ProverConfig::default())
119        } else {
120            config.prover_config
121        };
122        #[cfg(not(feature = "devenv"))]
123        let prover_config = if config.with_prover && config.prover_config.is_none() {
124            Some(ProverConfig::default())
125        } else {
126            config.prover_config
127        };
128        if let Some(ref prover_config) = prover_config {
129            spawn_prover(prover_config.clone()).await;
130        }
131        Ok(context)
132    }
133
134    pub fn indexer(&self) -> Result<&TestIndexer, RpcError> {
135        self.indexer.as_ref().ok_or(RpcError::IndexerNotInitialized)
136    }
137
138    pub fn indexer_mut(&mut self) -> Result<&mut TestIndexer, RpcError> {
139        self.indexer.as_mut().ok_or(RpcError::IndexerNotInitialized)
140    }
141
142    pub fn test_accounts(&self) -> &TestAccounts {
143        &self.test_accounts
144    }
145
146    /// Get account pubkeys of one state Merkle tree.
147    pub fn get_state_merkle_tree_account(&self) -> StateMerkleTreeAccounts {
148        self.test_accounts.v1_state_trees[0]
149    }
150
151    pub fn get_address_merkle_tree(&self) -> AddressMerkleTreeAccounts {
152        self.test_accounts.v1_address_trees[0]
153    }
154
155    #[cfg(feature = "v2")]
156    pub fn get_address_merkle_tree_v2(&self) -> solana_sdk::pubkey::Pubkey {
157        self.test_accounts.v2_address_trees[0]
158    }
159
160    pub async fn add_indexer(
161        &mut self,
162        test_accounts: &TestAccounts,
163        batch_size: Option<usize>,
164    ) -> Result<(), RpcError> {
165        let indexer = TestIndexer::init_from_acounts(
166            &self.payer,
167            test_accounts,
168            batch_size.unwrap_or_default(),
169        )
170        .await;
171        self.indexer = Some(indexer);
172        Ok(())
173    }
174
175    pub fn clone_indexer(&self) -> Result<TestIndexer, RpcError> {
176        Ok((*self
177            .indexer
178            .as_ref()
179            .ok_or(RpcError::IndexerNotInitialized)?)
180        .clone())
181    }
182}
183
184impl MerkleTreeExt for LightProgramTest {}
185
186impl Debug for LightProgramTest {
187    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
188        f.debug_struct("LightProgramTest")
189            .field("context", &"ProgramTestContext")
190            .field("indexer", &self.indexer)
191            .field("test_accounts", &self.test_accounts)
192            .finish()
193    }
194}
light_program_test/program_test/light_program_test.rs

light_program_test/program_test/
light_program_test.rs
