✨ Set custom + some additional docs updates

Ackee-Blockchain · Oct 31, 2024 · 5da8541 · 5da8541
1 parent ec9b597
commit 5da8541
Show file tree

Hide file tree

Showing 8 changed files with 155 additions and 137 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,6 +10,10 @@ incremented upon a breaking change and the patch version will be incremented for
 
 ## [dev] - Unreleased
 
+**Changed**
+
+- improve manipulations with AccountsStorages in get_accounts() function ([219](https://github.com/Ackee-Blockchain/trident/pull/219))
+
 **Added**
 
 - add/ add support for Clock sysvar manipulations with the client(i.e. warp to slot/epoch and forward in time) ([217](https://github.com/Ackee-Blockchain/trident/pull/217))

diff --git a/crates/client/src/source_code_generators/fuzz_instructions_generator.rs b/crates/client/src/source_code_generators/fuzz_instructions_generator.rs
@@ -49,8 +49,8 @@ pub fn generate_source_code(idls: &[Idl]) -> String {
 
         #(#all_instructions_ixops_impls)*
 
-        /// Use AccountsStorage<T> where T can be one of:
-        /// Keypair, PdaStore, TokenStore, MintStore, ProgramStore
+        /// Check supported AccountsStorages at
+        /// https://ackee.xyz/trident/docs/latest/features/account-storages/
         #[derive(Default)]
         pub struct FuzzAccounts {
             #(#all_fuzz_accounts),*
@@ -219,7 +219,7 @@ fn get_instruction_inputs(
                 /// Custom data types must derive `Debug` and `Arbitrary`.
                 /// To do this, redefine the type in the fuzz test and implement the `From` trait
                 /// to convert it into the type defined in the program.
-                /// For more details, see: https://ackee.xyz/trident/docs/dev/features/arbitrary-data/#custom-data-types
+                /// For more details, see: https://ackee.xyz/trident/docs/latest/features/fuzz-instructions/#custom-data-types
                 #[derive(Arbitrary, Debug)]
                 pub struct #instruction_data_name {
                      #(pub #parameters),*
@@ -309,7 +309,7 @@ fn get_instruction_ixops(
                     /// Definition of the Instruction data.
                     /// Use randomly generated data from the fuzzer using `self.data.arg_name`
                     /// or customize the data as needed.
-                    /// For more details, visit: https://ackee.xyz/trident/docs/dev/features/fuzz-instructions/#get-data
+                    /// For more details, visit: https://ackee.xyz/trident/docs/latest/features/fuzz-instructions/#get-data
                     fn get_data(
                         &self,
                         _client: &mut impl FuzzClient,
@@ -322,10 +322,11 @@ fn get_instruction_ixops(
                     }
 
                     /// Definition of of the accounts required by the Instruction.
-                    /// To utilize accounts stored in `FuzzAccounts`, use `fuzz_accounts.account_name.get_or_create_account()`.
+                    /// To utilize accounts stored in `FuzzAccounts`, use
+                    /// `fuzz_accounts.account_name.get_or_create_account()`.
                     /// If no signers are required, leave the vector empty.
                     /// For AccountMetas use <program>::accounts::<corresponding_metas>
-                    /// For more details, see: https://ackee.xyz/trident/docs/dev/features/fuzz-instructions/#get-accounts
+                    /// For more details, see: https://ackee.xyz/trident/docs/latest/features/fuzz-instructions/#get-accounts
                     fn get_accounts(
                     &self,
                     client: &mut impl FuzzClient,

diff --git a/crates/client/src/source_code_generators/test_fuzz_generator.rs b/crates/client/src/source_code_generators/test_fuzz_generator.rs
@@ -29,7 +29,7 @@ pub fn generate_source_code(idl_instructions: &[Idl]) -> String {
         ///     Ok(vec![init])
         /// }
         /// ```
-        /// For more details, see: https://ackee.xyz/trident/docs/dev/features/instructions-sequences/#instructions-sequences
+        /// For more details, see: https://ackee.xyz/trident/docs/latest/features/instructions-sequences/#instructions-sequences
         impl FuzzDataBuilder<FuzzInstruction> for MyFuzzData {}
 
         /// `fn fuzz_iteration` runs during every fuzzing iteration.

diff --git a/crates/client/tests/expected_source_codes/expected_fuzz_instructions.rs b/crates/client/tests/expected_source_codes/expected_fuzz_instructions.rs
@@ -23,7 +23,7 @@ pub struct InitializeIxDummy2Accounts {
 /// To do this, redefine the type in the fuzz test and implement the `From`
 /// trait
 /// to convert it into the type defined in the program.
-/// For more details, see: https://ackee.xyz/trident/docs/dev/features/arbitrary-data/#custom-data-types
+/// For more details, see: https://ackee.xyz/trident/docs/latest/features/fuzz-instructions/#custom-data-types
 #[derive(Arbitrary, Debug)]
 pub struct InitializeIxDummy2Data {
     pub _var1: bool,
@@ -73,7 +73,7 @@ pub struct InitializeIxDummyExampleAccounts {
 /// To do this, redefine the type in the fuzz test and implement the `From`
 /// trait
 /// to convert it into the type defined in the program.
-/// For more details, see: https://ackee.xyz/trident/docs/dev/features/arbitrary-data/#custom-data-types
+/// For more details, see: https://ackee.xyz/trident/docs/latest/features/fuzz-instructions/#custom-data-types
 #[derive(Arbitrary, Debug)]
 pub struct InitializeIxDummyExampleData {
     pub _var1: bool,
@@ -111,7 +111,7 @@ impl<'info> IxOps<'info> for InitializeIxDummy2 {
     /// Definition of the Instruction data.
     /// Use randomly generated data from the fuzzer using `self.data.arg_name`
     /// or customize the data as needed.
-    /// For more details, visit: https://ackee.xyz/trident/docs/dev/features/fuzz-instructions/#get-data
+    /// For more details, visit: https://ackee.xyz/trident/docs/latest/features/fuzz-instructions/#get-data
     fn get_data(
         &self,
         _client: &mut impl FuzzClient,
@@ -148,7 +148,7 @@ impl<'info> IxOps<'info> for InitializeIxDummy2 {
     /// `fuzz_accounts.account_name.get_or_create_account()`.
     /// If no signers are required, leave the vector empty.
     /// For AccountMetas use <program>::accounts::<corresponding_metas>
-    /// For more details, see: https://ackee.xyz/trident/docs/dev/features/fuzz-instructions/#get-accounts
+    /// For more details, see: https://ackee.xyz/trident/docs/latest/features/fuzz-instructions/#get-accounts
     fn get_accounts(
         &self,
         client: &mut impl FuzzClient,
@@ -172,7 +172,7 @@ impl<'info> IxOps<'info> for InitializeIxDummyExample {
     /// Definition of the Instruction data.
     /// Use randomly generated data from the fuzzer using `self.data.arg_name`
     /// or customize the data as needed.
-    /// For more details, visit: https://ackee.xyz/trident/docs/dev/features/fuzz-instructions/#get-data
+    /// For more details, visit: https://ackee.xyz/trident/docs/latest/features/fuzz-instructions/#get-data
     fn get_data(
         &self,
         _client: &mut impl FuzzClient,
@@ -209,7 +209,7 @@ impl<'info> IxOps<'info> for InitializeIxDummyExample {
     /// `fuzz_accounts.account_name.get_or_create_account()`.
     /// If no signers are required, leave the vector empty.
     /// For AccountMetas use <program>::accounts::<corresponding_metas>
-    /// For more details, see: https://ackee.xyz/trident/docs/dev/features/fuzz-instructions/#get-accounts
+    /// For more details, see: https://ackee.xyz/trident/docs/latest/features/fuzz-instructions/#get-accounts
     fn get_accounts(
         &self,
         client: &mut impl FuzzClient,
@@ -220,8 +220,8 @@ impl<'info> IxOps<'info> for InitializeIxDummyExample {
         Ok((signers, acc_meta))
     }
 }
-/// Use AccountsStorage<T> where T can be one of:
-/// Keypair, PdaStore, TokenStore, MintStore, ProgramStore
+/// Check supported AccountsStorages at
+/// https://ackee.xyz/trident/docs/latest/features/account-storages/
 #[derive(Default)]
 pub struct FuzzAccounts {
     signer_dummy_2: AccountsStorage<todo!()>,

diff --git a/crates/client/tests/expected_source_codes/expected_test_fuzz.rs b/crates/client/tests/expected_source_codes/expected_test_fuzz.rs
@@ -18,7 +18,7 @@ struct MyFuzzData;
 ///     Ok(vec![init])
 /// }
 /// ```
-/// For more details, see: https://ackee.xyz/trident/docs/dev/features/instructions-sequences/#instructions-sequences
+/// For more details, see: https://ackee.xyz/trident/docs/latest/features/instructions-sequences/#instructions-sequences
 impl FuzzDataBuilder<FuzzInstruction> for MyFuzzData {}
 /// `fn fuzz_iteration` runs during every fuzzing iteration.
 /// Modification is not required.

diff --git a/crates/fuzz/src/accounts_storage.rs b/crates/fuzz/src/accounts_storage.rs
@@ -2,6 +2,7 @@
 use std::collections::HashMap;
 
 use solana_sdk::{
+    account::AccountSharedData,
     clock::{Clock, Epoch},
     pubkey::Pubkey,
     signature::Keypair,
@@ -29,7 +30,7 @@ pub struct MintStore {
 }
 
 pub struct ProgramStore {
-    pub pubkey: u8,
+    pub pubkey: Pubkey,
 }
 pub struct VoteAccountStore {
     pub pubkey: Pubkey,
@@ -53,9 +54,17 @@ impl<T> AccountsStorage<T> {
         }
     }
 
-    /// Returns a mutable reference to the underlying HashMap that stores accounts with IDs as keys
-    fn storage(&mut self) -> &mut HashMap<AccountId, T> {
-        &mut self.accounts
+    pub fn set_custom(
+        &mut self,
+        account_id: AccountId,
+        client: &mut impl FuzzClient,
+        address: Pubkey,
+        account: AccountSharedData,
+    ) where
+        T: From<Pubkey>,
+    {
+        client.set_account_custom(&address, &account);
+        self.accounts.insert(account_id, T::from(address));
     }
 }
 
@@ -86,6 +95,63 @@ impl AccountsStorage<Keypair> {
     }
 }
 
+impl AccountsStorage<ProgramStore> {
+    pub fn get_or_create_account(
+        &mut self,
+        account_id: AccountId,
+        _client: &mut impl FuzzClient,
+        program_id: Pubkey,
+    ) -> Pubkey {
+        let program_id = self
+            .accounts
+            .entry(account_id)
+            .or_insert_with(|| ProgramStore { pubkey: program_id });
+        program_id.pubkey
+    }
+    pub fn get(&self, account_id: AccountId) -> Pubkey {
+        match self.accounts.get(&account_id) {
+            Some(v) => v.pubkey,
+            None => Pubkey::new_unique(),
+        }
+    }
+}
+impl From<Pubkey> for TokenStore {
+    fn from(pubkey: Pubkey) -> Self {
+        TokenStore { pubkey }
+    }
+}
+impl From<Pubkey> for MintStore {
+    fn from(pubkey: Pubkey) -> Self {
+        MintStore { pubkey }
+    }
+}
+
+impl From<Pubkey> for ProgramStore {
+    fn from(pubkey: Pubkey) -> Self {
+        ProgramStore { pubkey }
+    }
+}
+
+impl From<Pubkey> for VoteAccountStore {
+    fn from(pubkey: Pubkey) -> Self {
+        VoteAccountStore { pubkey }
+    }
+}
+
+impl From<Pubkey> for StakeAccountStore {
+    fn from(pubkey: Pubkey) -> Self {
+        StakeAccountStore { pubkey }
+    }
+}
+impl From<Pubkey> for PdaStore {
+    fn from(pubkey: Pubkey) -> Self {
+        PdaStore {
+            pubkey,
+            seeds: Vec::new(), // Note: This creates empty seeds
+        }
+    }
+}
+
 impl AccountsStorage<TokenStore> {
     #[allow(clippy::too_many_arguments)]
     pub fn get_or_create_account(
@@ -140,7 +206,7 @@ impl AccountsStorage<MintStore> {
     pub fn get(&self, account_id: AccountId) -> Pubkey {
         match self.accounts.get(&account_id) {
             Some(v) => v.pubkey,
-            None => Pubkey::default(),
+            None => Pubkey::new_unique(),
         }
     }
 }
@@ -164,15 +230,15 @@ impl AccountsStorage<PdaStore> {
                     self.accounts.insert(account_id, pda_store);
                     key
                 } else {
-                    Pubkey::default()
+                    Pubkey::new_unique()
                 }
             }
         }
     }
     pub fn get(&self, account_id: AccountId) -> Pubkey {
         match self.accounts.get(&account_id) {
             Some(v) => v.pubkey,
-            None => Pubkey::default(),
+            None => Pubkey::new_unique(),
         }
     }
 }
@@ -204,7 +270,7 @@ impl AccountsStorage<VoteAccountStore> {
     pub fn get(&self, account_id: AccountId) -> Pubkey {
         match self.accounts.get(&account_id) {
             Some(v) => v.pubkey,
-            None => Pubkey::default(),
+            None => Pubkey::new_unique(),
         }
     }
 }
@@ -254,7 +320,7 @@ impl AccountsStorage<StakeAccountStore> {
     pub fn get(&self, account_id: AccountId) -> Pubkey {
         match self.accounts.get(&account_id) {
             Some(v) => v.pubkey,
-            None => Pubkey::default(),
+            None => Pubkey::new_unique(),
         }
     }
 }
diff --git a/documentation/docs/features/account-storages.md b/documentation/docs/features/account-storages.md
@@ -22,6 +22,9 @@ Instead, Trident generates random **AccountIDs** which are indexes to **Account
     - PDA
     - Token Account
     - Program account
+    - Stake account
+    - Vote account
+
 
     Then use the corresponding AccountsStorage.