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