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