solscript_codegen/

test_gen.rs

//! Test File Generator
//!
//! Generates Anchor test files for the program.

use crate::ir::*;
use crate::CodegenError;

/// Test file generator
pub struct TestGenerator {
    program_name: String,
}

impl Default for TestGenerator {
    fn default() -> Self {
        Self::new()
    }
}

impl TestGenerator {
    pub fn new() -> Self {
        Self {
            program_name: String::new(),
        }
    }

    /// Generate the test file
    pub fn generate(&mut self, ir: &SolanaProgram) -> Result<String, CodegenError> {
        self.program_name = to_snake_case(&ir.name);
        let class_name = to_camel_case(&ir.name);

        let mut output = String::new();

        // Header and imports
        output.push_str(&self.generate_header(&class_name));
        output.push('\n');

        // Test setup
        output.push_str(&self.generate_setup(&class_name));
        output.push('\n');

        // Generate test cases for each instruction
        for instruction in &ir.instructions {
            output.push_str(&self.generate_test_case(instruction, &class_name)?);
            output.push('\n');
        }

        // Close the describe block
        output.push_str("});\n");

        Ok(output)
    }

    fn generate_header(&self, class_name: &str) -> String {
        format!(
            r#"/**
 * Generated by SolScript compiler
 * Anchor tests for {} program
 */

import * as anchor from "@coral-xyz/anchor";
import {{ Program }} from "@coral-xyz/anchor";
import {{ Keypair, PublicKey, SystemProgram }} from "@solana/web3.js";
import {{ assert }} from "chai";

// Import the generated types
// import {{ {} }} from "../target/types/{}";

describe("{}", () => {{
  // Configure the client to use the local cluster
  const provider = anchor.AnchorProvider.env();
  anchor.setProvider(provider);

  // Get the program
  // const program = anchor.workspace.{} as Program<{}>;

  // Test accounts
  let stateAccount: Keypair;
  let authority: Keypair;
"#,
            class_name, class_name, self.program_name, class_name, class_name, class_name
        )
    }

    fn generate_setup(&self, _class_name: &str) -> String {
        r#"
  before(async () => {
    stateAccount = Keypair.generate();
    authority = provider.wallet.payer;
  });

  // Helper to get state PDA
  const getStatePDA = (): [PublicKey, number] => {
    return PublicKey.findProgramAddressSync(
      [Buffer.from("state")],
      // program.programId
      SystemProgram.programId // placeholder
    );
  };
"#
        .to_string()
    }

    fn generate_test_case(
        &self,
        instruction: &Instruction,
        _class_name: &str,
    ) -> Result<String, CodegenError> {
        let test_name = to_snake_case(&instruction.name).replace('_', " ");
        let method_name = to_camel_case_lower(&instruction.name);

        // Generate mock argument values based on parameter types
        let mut args = Vec::new();
        for param in &instruction.params {
            let mock_value = self.generate_mock_value(&param.ty, &param.name);
            args.push(format!(
                "const {} = {};",
                to_camel_case_lower(&param.name),
                mock_value
            ));
        }
        let args_setup = args.join("\n    ");

        // Generate the instruction call arguments
        let call_args: Vec<String> = instruction
            .params
            .iter()
            .map(|p| to_camel_case_lower(&p.name))
            .collect();
        let call_args_str = call_args.join(", ");

        // Determine if this is the initialize/constructor
        let is_init = instruction.name.to_lowercase() == "initialize"
            || instruction.name.to_lowercase() == "init"
            || instruction.name.to_lowercase() == "constructor";

        let accounts_str = if is_init {
            r#"//       state: stateAccount.publicKey,
    //       signer: authority.publicKey,
    //       systemProgram: SystemProgram.programId,"#
        } else {
            r#"//       state: stateAccount.publicKey,
    //       signer: authority.publicKey,"#
        };

        let signers_str = if is_init {
            "[stateAccount, authority]"
        } else {
            "[authority]"
        };

        Ok(format!(
            r#"
  it("should {}", async () => {{
    {}

    // Uncomment when program is available:
    // const tx = await program.methods
    //   .{}({})
    //   .accounts({{
    {}
    //   }})
    //   .signers({})
    //   .rpc();
    // console.log("Transaction signature:", tx);

    // Add assertions based on expected behavior
    // const state = await program.account.{}State.fetch(stateAccount.publicKey);
    // assert.ok(state, "State should exist");
  }});
"#,
            test_name,
            args_setup,
            method_name,
            call_args_str,
            accounts_str,
            signers_str,
            to_camel_case_lower(&self.program_name)
        ))
    }

    fn generate_mock_value(&self, ty: &SolanaType, name: &str) -> String {
        match ty {
            SolanaType::U8 => "42".to_string(),
            SolanaType::U16 => "1000".to_string(),
            SolanaType::U32 => "100000".to_string(),
            SolanaType::U64 | SolanaType::U128 => "new anchor.BN(1000000)".to_string(),
            SolanaType::I8 | SolanaType::I16 | SolanaType::I32 => "42".to_string(),
            SolanaType::I64 | SolanaType::I128 => "new anchor.BN(1000000)".to_string(),
            SolanaType::Bool => "true".to_string(),
            SolanaType::String => format!("\"test_{}\"", name),
            SolanaType::Pubkey | SolanaType::Signer => "Keypair.generate().publicKey".to_string(),
            SolanaType::Bytes => "Buffer.from(\"test\")".to_string(),
            SolanaType::FixedBytes(n) => format!("new Uint8Array({})", n),
            SolanaType::Array(inner, size) => {
                let inner_val = self.generate_mock_value(inner, name);
                format!("Array({}).fill({})", size, inner_val)
            }
            SolanaType::Vec(inner) => {
                let inner_val = self.generate_mock_value(inner, name);
                format!("[{}]", inner_val)
            }
            SolanaType::Option(_) => "null".to_string(),
            SolanaType::Mapping(_, _) => "new Map()".to_string(),
            SolanaType::Custom(type_name) => {
                // Generate a placeholder object for custom types
                // In real usage, this should be filled with actual field values
                format!("{{ /* {} instance - fill in fields */ }}", type_name)
            }
        }
    }
}

// Helper functions for case conversion
fn to_camel_case(s: &str) -> String {
    let mut result = String::new();
    let mut capitalize_next = true;

    for c in s.chars() {
        if c == '_' || c == '-' {
            capitalize_next = true;
        } else if capitalize_next {
            result.push(c.to_ascii_uppercase());
            capitalize_next = false;
        } else {
            result.push(c);
        }
    }

    result
}

fn to_camel_case_lower(s: &str) -> String {
    let camel = to_camel_case(s);
    let mut chars = camel.chars();
    match chars.next() {
        None => String::new(),
        Some(c) => c.to_ascii_lowercase().to_string() + chars.as_str(),
    }
}

fn to_snake_case(s: &str) -> String {
    let mut result = String::new();
    for (i, c) in s.chars().enumerate() {
        if c.is_uppercase() && i > 0 {
            result.push('_');
            result.push(c.to_ascii_lowercase());
        } else {
            result.push(c.to_ascii_lowercase());
        }
    }
    result
}