light_program_test/accounts/
initialize.rs

1use account_compression::{utils::constants::GROUP_AUTHORITY_SEED, GroupAuthority};
2use light_client::{
3    indexer::{AddressMerkleTreeAccounts, StateMerkleTreeAccounts},
4    rpc::{Rpc, RpcError},
5};
6use light_registry::{
7    account_compression_cpi::sdk::get_registered_program_pda,
8    sdk::{
9        create_initialize_governance_authority_instruction,
10        create_initialize_group_authority_instruction, create_update_protocol_config_instruction,
11    },
12    utils::{get_cpi_authority_pda, get_forester_pda, get_protocol_config_pda_address},
13};
14use solana_sdk::{
15    pubkey::Pubkey,
16    signature::{Keypair, Signer},
17};
18
19#[cfg(feature = "v2")]
20use super::{
21    address_tree_v2::create_batch_address_merkle_tree,
22    state_tree_v2::create_batched_state_merkle_tree,
23};
24#[cfg(feature = "devenv")]
25use crate::accounts::register_program::register_program_with_registry_program;
26use crate::{
27    accounts::{
28        address_tree::create_address_merkle_tree_and_queue_account,
29        state_tree::create_state_merkle_tree_and_queue_account,
30        test_accounts::{ProtocolAccounts, StateMerkleTreeAccountsV2, TestAccounts},
31        test_keypairs::*,
32    },
33    program_test::TestRpc,
34    ProgramTestConfig,
35};
36
37#[allow(clippy::too_many_arguments)]
38pub async fn initialize_accounts<R: Rpc + TestRpc>(
39    context: &mut R,
40    config: &ProgramTestConfig,
41    keypairs: &TestKeypairs,
42) -> Result<TestAccounts, RpcError> {
43    let ProgramTestConfig {
44        protocol_config,
45        v2_state_tree_config,
46        v2_address_tree_config,
47        skip_register_programs,
48        skip_second_v1_tree,
49        v1_state_tree_config,
50        v1_nullifier_queue_config,
51        v1_address_tree_config,
52        v1_address_queue_config,
53        ..
54    } = config;
55    let _v2_address_tree_config = v2_address_tree_config;
56    let _v2_state_tree_config = v2_state_tree_config;
57    let _skip_register_programs = skip_register_programs;
58    let cpi_authority_pda = get_cpi_authority_pda();
59    let protocol_config_pda = get_protocol_config_pda_address();
60    let instruction = create_initialize_governance_authority_instruction(
61        keypairs.governance_authority.pubkey(),
62        keypairs.governance_authority.pubkey(),
63        *protocol_config,
64    );
65    let update_instruction = create_update_protocol_config_instruction(
66        keypairs.governance_authority.pubkey(),
67        Some(keypairs.governance_authority.pubkey()),
68        None,
69    );
70    context
71        .create_and_send_transaction(
72            &[instruction, update_instruction],
73            &keypairs.governance_authority.pubkey(),
74            &[&keypairs.governance_authority],
75        )
76        .await?;
77
78    let group_pda = initialize_new_group(
79        &keypairs.group_pda_seed,
80        &keypairs.governance_authority,
81        context,
82        cpi_authority_pda.0,
83    )
84    .await?;
85
86    let gov_authority = context
87        .get_anchor_account::<GroupAuthority>(&protocol_config_pda.0)
88        .await?
89        .ok_or(RpcError::AccountDoesNotExist(
90            protocol_config_pda.0.to_string(),
91        ))?;
92    assert_eq!(
93        gov_authority.authority,
94        keypairs.governance_authority.pubkey()
95    );
96    if gov_authority.authority != keypairs.governance_authority.pubkey() {
97        return Err(RpcError::CustomError(
98            "Invalid governance authority.".to_string(),
99        ));
100    }
101
102    #[cfg(feature = "devenv")]
103    if !_skip_register_programs {
104        register_program_with_registry_program(
105            context,
106            &keypairs.governance_authority,
107            &group_pda,
108            &keypairs.system_program,
109        )
110        .await?;
111        register_program_with_registry_program(
112            context,
113            &keypairs.governance_authority,
114            &group_pda,
115            &keypairs.registry_program,
116        )
117        .await?;
118    }
119    let merkle_tree_pubkey = keypairs.state_merkle_tree.pubkey();
120    let nullifier_queue_pubkey = keypairs.nullifier_queue.pubkey();
121    if !config.skip_v1_trees {
122        create_state_merkle_tree_and_queue_account(
123            &keypairs.governance_authority,
124            true,
125            context,
126            &keypairs.state_merkle_tree,
127            &keypairs.nullifier_queue,
128            Some(&keypairs.cpi_context_account),
129            None,
130            None,
131            1,
132            v1_state_tree_config,
133            v1_nullifier_queue_config,
134        )
135        .await?;
136
137        if !skip_second_v1_tree {
138            create_state_merkle_tree_and_queue_account(
139                &keypairs.governance_authority,
140                true,
141                context,
142                &keypairs.state_merkle_tree_2,
143                &keypairs.nullifier_queue_2,
144                Some(&keypairs.cpi_context_2),
145                None,
146                None,
147                2,
148                v1_state_tree_config,
149                v1_nullifier_queue_config,
150            )
151            .await?;
152        }
153
154        create_address_merkle_tree_and_queue_account(
155            &keypairs.governance_authority,
156            true,
157            context,
158            &keypairs.address_merkle_tree,
159            &keypairs.address_merkle_tree_queue,
160            None,
161            None,
162            v1_address_tree_config,
163            v1_address_queue_config,
164            0,
165        )
166        .await?;
167    }
168    #[cfg(feature = "v2")]
169    if let Some(v2_state_tree_config) = _v2_state_tree_config {
170        create_batched_state_merkle_tree(
171            &keypairs.governance_authority,
172            true,
173            context,
174            &keypairs.batched_state_merkle_tree,
175            &keypairs.batched_output_queue,
176            &keypairs.batched_cpi_context,
177            *v2_state_tree_config,
178        )
179        .await?;
180    }
181    #[cfg(feature = "v2")]
182    if let Some(params) = _v2_address_tree_config {
183        create_batch_address_merkle_tree(
184            context,
185            &keypairs.governance_authority,
186            &keypairs.batch_address_merkle_tree,
187            *params,
188        )
189        .await?;
190    }
191
192    let registered_system_program_pda =
193        get_registered_program_pda(&Pubkey::from(light_sdk::constants::LIGHT_SYSTEM_PROGRAM_ID));
194    let registered_registry_program_pda = get_registered_program_pda(&light_registry::ID);
195    Ok(TestAccounts {
196        protocol: ProtocolAccounts {
197            governance_authority: keypairs.governance_authority.insecure_clone(),
198            governance_authority_pda: protocol_config_pda.0,
199            group_pda,
200            forester: keypairs.forester.insecure_clone(),
201            registered_program_pda: registered_system_program_pda,
202            registered_registry_program_pda,
203            registered_forester_pda: get_forester_pda(&keypairs.forester.pubkey()).0,
204        },
205        v1_state_trees: vec![StateMerkleTreeAccounts {
206            merkle_tree: merkle_tree_pubkey,
207            nullifier_queue: nullifier_queue_pubkey,
208            cpi_context: keypairs.cpi_context_account.pubkey(),
209        }],
210        v1_address_trees: vec![AddressMerkleTreeAccounts {
211            merkle_tree: keypairs.address_merkle_tree.pubkey(),
212            queue: keypairs.address_merkle_tree_queue.pubkey(),
213        }],
214        v2_state_trees: vec![StateMerkleTreeAccountsV2 {
215            merkle_tree: keypairs.batched_state_merkle_tree.pubkey(),
216            output_queue: keypairs.batched_output_queue.pubkey(),
217            cpi_context: keypairs.batched_cpi_context.pubkey(),
218        }],
219        v2_address_trees: vec![keypairs.batch_address_merkle_tree.pubkey()],
220    })
221}
222
223pub fn get_group_pda(seed: Pubkey) -> Pubkey {
224    Pubkey::find_program_address(
225        &[GROUP_AUTHORITY_SEED, seed.to_bytes().as_slice()],
226        &account_compression::ID,
227    )
228    .0
229}
230
231pub async fn initialize_new_group<R: Rpc>(
232    group_seed_keypair: &Keypair,
233    payer: &Keypair,
234    context: &mut R,
235    authority: Pubkey,
236) -> Result<Pubkey, RpcError> {
237    let group_pda = Pubkey::find_program_address(
238        &[
239            GROUP_AUTHORITY_SEED,
240            group_seed_keypair.pubkey().to_bytes().as_slice(),
241        ],
242        &account_compression::ID,
243    )
244    .0;
245
246    let instruction = create_initialize_group_authority_instruction(
247        payer.pubkey(),
248        group_pda,
249        group_seed_keypair.pubkey(),
250        authority,
251    );
252
253    context
254        .create_and_send_transaction(
255            &[instruction],
256            &payer.pubkey(),
257            &[payer, group_seed_keypair],
258        )
259        .await?;
260    let group_authority = context
261        .get_anchor_account::<GroupAuthority>(&group_pda)
262        .await?
263        .ok_or(RpcError::CustomError(
264            "Group authority account does not exist.".to_string(),
265        ))?;
266    if group_authority.authority != authority {
267        return Err(RpcError::CustomError(
268            "Group authority account does not match the provided authority.".to_string(),
269        ));
270    }
271    if group_authority.seed != group_seed_keypair.pubkey() {
272        return Err(RpcError::CustomError(
273            "Group authority account does not match the provided seed.".to_string(),
274        ));
275    }
276    Ok(group_pda)
277}