trident_client/
test_generator.rs

1use crate::commander::Commander;
2use crate::construct_path;
3use crate::error::Error;
4use fehler::throws;
5use std::path::Path;
6use std::path::PathBuf;
7use trident_idl_spec::Idl;
8use trident_template::GeneratedFiles;
9use trident_template::TridentTemplates;
10
11pub struct TestGenerator {
12    pub(crate) root: PathBuf,
13    pub(crate) skip_build: bool,
14    pub(crate) anchor_idls: Vec<Idl>,
15    pub(crate) template_engine: TridentTemplates,
16    pub(crate) generated_files: Option<GeneratedFiles>,
17}
18
19impl TestGenerator {
20    #[throws]
21    pub fn new_with_root(root: &str, skip_build: bool) -> Self {
22        Self {
23            root: Path::new(&root).to_path_buf(),
24            skip_build,
25            anchor_idls: Vec::default(),
26            template_engine: TridentTemplates::new()?,
27            generated_files: None,
28        }
29    }
30
31    #[throws]
32    pub async fn initialize(&mut self, program_name: Option<String>, test_name: Option<String>) {
33        if !self.skip_build {
34            Commander::build_anchor_project(&self.root, program_name.clone()).await?;
35        }
36
37        self.load_programs_idl(program_name.clone())?;
38        self.create_template().await?;
39        self.add_new_fuzz_test(&test_name).await?;
40        self.create_trident_toml().await?;
41        self.create_vscode_settings().await?;
42    }
43
44    #[throws]
45    pub async fn add_fuzz_test(&mut self, program_name: Option<String>, test_name: Option<String>) {
46        if !self.skip_build {
47            Commander::build_anchor_project(&self.root, program_name.clone()).await?;
48        }
49
50        self.load_programs_idl(program_name.clone())?;
51
52        self.create_template().await?;
53
54        self.add_new_fuzz_test(&test_name).await?;
55    }
56
57    #[throws]
58    pub async fn refresh_fuzz_test(
59        &mut self,
60        fuzz_test_name: String,
61        program_name: Option<String>,
62    ) {
63        if !self.skip_build {
64            Commander::build_anchor_project(&self.root, program_name.clone()).await?;
65        }
66
67        self.load_programs_idl(program_name)?;
68        self.create_template().await?;
69        self.refresh_types_file(&fuzz_test_name).await?;
70    }
71
72    #[throws]
73    async fn create_template(&mut self) {
74        let current_package_version = env!("CARGO_PKG_VERSION");
75
76        // Generate templates using Tera
77        let output = self
78            .template_engine
79            .generate(&self.anchor_idls, current_package_version)?;
80
81        // Store the generated output
82        self.generated_files = Some(output);
83    }
84
85    #[throws]
86    fn load_programs_idl(&mut self, program_name: Option<String>) {
87        let target_path = construct_path!(self.root, "target/idl/");
88
89        // TODO consider optionally excluding packages
90        self.anchor_idls = crate::idl_loader::load_idls(target_path, program_name)?;
91    }
92
93    pub(crate) fn get_test_fuzz(&self) -> String {
94        if let Some(ref output) = self.generated_files {
95            output.test_fuzz.clone()
96        } else {
97            String::new()
98        }
99    }
100
101    pub(crate) fn get_types(&self) -> String {
102        if let Some(ref output) = self.generated_files {
103            output.types.clone()
104        } else {
105            String::new()
106        }
107    }
108
109    pub(crate) fn get_fuzz_accounts(&self) -> String {
110        if let Some(ref output) = self.generated_files {
111            output.fuzz_accounts.clone()
112        } else {
113            String::new()
114        }
115    }
116
117    pub(crate) fn get_trident_toml(&self) -> String {
118        if let Some(ref output) = self.generated_files {
119            output.trident_toml.clone()
120        } else {
121            String::new()
122        }
123    }
124
125    pub(crate) fn get_cargo_fuzz_toml(&self) -> String {
126        if let Some(ref output) = self.generated_files {
127            output.cargo_fuzz_toml.clone()
128        } else {
129            String::new()
130        }
131    }
132}