Skip to content

Commit

Permalink
Internal: use rstest for complex tests (#10)
Browse files Browse the repository at this point in the history 
Problem: we tend to repeat a lot of boilerplate in tests because normal
Rust tests are limited (ex: no fixtures).

Solution: use rstest.
  • Loading branch information
Olivier Desenfans authored Nov 20, 2023
1 parent 30df4d2 commit b933115
Show file tree
Hide file tree
Showing 3 changed files with 124 additions and 68 deletions.
1 change: 1 addition & 0 deletions Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,3 +10,4 @@ serde_json = "1.0.108"
tokio = { version = "1.34.0", features = ["macros", "process", "rt-multi-thread"] }
tonic = "0.10.2"
tonic-build = "0.10.2"
rstest = "0.18.2"
1 change: 1 addition & 0 deletions stone-prover/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,4 +12,5 @@ thiserror = "1.0.50"
tokio = { workspace = true }

[dev-dependencies]
rstest = { workspace = true }
test-toolkit = { path = "../test-toolkit" }
190 changes: 122 additions & 68 deletions stone-prover/src/prover.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -223,6 +223,7 @@ pub async fn run_prover_async(

#[cfg(test)]
mod test {
use rstest::{fixture, rstest};
use tempfile::NamedTempFile;

use test_toolkit::get_fixture_path;
Expand All @@ -238,114 +239,167 @@ mod test {
proof
}

/// Check that the Stone Prover command-line wrapper works.
#[test]
fn test_run_prover_from_command_line() {
let prover_config_file = get_fixture_path("fibonacci/cpu_air_prover_config.json");
let parameter_file = get_fixture_path("fibonacci/cpu_air_params.json");
/// All the files forming a complete prover test case.
struct ProverTestCase {
public_input_file: PathBuf,
prover_config_file: PathBuf,
prover_parameter_file: PathBuf,
memory_file: PathBuf,
trace_file: PathBuf,
proof_file: PathBuf,
}

#[fixture]
fn fibonacci() -> ProverTestCase {
let public_input_file = get_fixture_path("fibonacci/fibonacci_public_input.json");
let prover_config_file = get_fixture_path("fibonacci/cpu_air_prover_config.json");
let prover_parameter_file = get_fixture_path("fibonacci/cpu_air_params.json");
let memory_file = get_fixture_path("fibonacci/fibonacci_memory.bin");
let trace_file = get_fixture_path("fibonacci/fibonacci_trace.bin");
let proof_file = get_fixture_path("fibonacci/fibonacci_proof.json");

ProverTestCase {
public_input_file,
prover_config_file,
prover_parameter_file,
memory_file,
trace_file,
proof_file,
}
}

/// Test case files adapted to match the prover command line arguments.
struct ProverCliTestCase {
public_input_file: PathBuf,
private_input_file: NamedTempFile,
prover_config_file: PathBuf,
prover_parameter_file: PathBuf,
proof: Proof,
}

#[fixture]
fn prover_cli_test_case(#[from(fibonacci)] files: ProverTestCase) -> ProverCliTestCase {
// Generate the private input in a temporary file
let private_input_file =
NamedTempFile::new().expect("Creating temporary private input file failed");
let private_input = PrivateInput {
memory_path: memory_file,
trace_path: trace_file,
memory_path: files.memory_file.clone(),
trace_path: files.trace_file.clone(),
pedersen: vec![],
range_check: vec![],
ecdsa: vec![],
};

serde_json::to_writer(&private_input_file, &private_input)
.expect("Writing private input file failed");

let proof = read_proof_file(&files.proof_file);

ProverCliTestCase {
public_input_file: files.public_input_file,
private_input_file,
prover_config_file: files.prover_config_file,
prover_parameter_file: files.prover_parameter_file,
proof,
}
}

struct ParsedProverTestCase {
public_input: PublicInput,
memory: Vec<u8>,
trace: Vec<u8>,
prover_config: ProverConfig,
prover_parameters: ProverParameters,
proof: Proof,
}

#[fixture]
fn parsed_prover_test_case(#[from(fibonacci)] files: ProverTestCase) -> ParsedProverTestCase {
let public_input: PublicInput = read_json_from_file(files.public_input_file).unwrap();
let prover_config: ProverConfig = read_json_from_file(files.prover_config_file).unwrap();
let prover_parameters: ProverParameters =
read_json_from_file(files.prover_parameter_file).unwrap();
let memory = std::fs::read(files.memory_file).unwrap();
let trace = std::fs::read(files.trace_file).unwrap();

let proof = read_proof_file(&files.proof_file);

ParsedProverTestCase {
public_input,
memory,
trace,
prover_config,
prover_parameters,
proof,
}
}

#[fixture]
fn prover_in_path() {
// Add build dir to path for the duration of the test
let path = std::env::var("PATH").unwrap_or_default();
let build_dir = env!("OUT_DIR");
std::env::set_var("PATH", format!("{build_dir}:{path}"));
}

/// Check that the Stone Prover command-line wrapper works.
#[rstest]
fn test_run_prover_from_command_line(
prover_cli_test_case: ProverCliTestCase,
#[from(prover_in_path)] _path: (),
) {
// Add build dir to path for the duration of the test
let path = std::env::var("PATH").unwrap_or_default();
let build_dir = env!("OUT_DIR");
std::env::set_var("PATH", format!("{build_dir}:{path}"));

let output_file = NamedTempFile::new().expect("Creating output file failed");
run_prover_from_command_line(
&public_input_file,
private_input_file.path(),
&prover_config_file,
&parameter_file,
&prover_cli_test_case.public_input_file,
&prover_cli_test_case.private_input_file.path(),
&prover_cli_test_case.prover_config_file,
&prover_cli_test_case.prover_parameter_file,
output_file.path(),
)
.unwrap();

let proof = read_proof_file(output_file.path());
let expected_proof_file = get_fixture_path("fibonacci/fibonacci_proof.json");
let expected_proof = read_proof_file(expected_proof_file);

assert_eq!(proof.proof_hex, expected_proof.proof_hex);
assert_eq!(proof.proof_hex, prover_cli_test_case.proof.proof_hex);
}

#[test]
fn test_run_prover() {
let public_input_file = get_fixture_path("fibonacci/fibonacci_public_input.json");
let prover_config_file = get_fixture_path("fibonacci/cpu_air_prover_config.json");
let parameter_file = get_fixture_path("fibonacci/cpu_air_params.json");
let memory_file = get_fixture_path("fibonacci/fibonacci_memory.bin");
let trace_file = get_fixture_path("fibonacci/fibonacci_trace.bin");

let public_input: PublicInput = read_json_from_file(public_input_file).unwrap();
let prover_config: ProverConfig = read_json_from_file(prover_config_file).unwrap();
let prover_parameters: ProverParameters = read_json_from_file(parameter_file).unwrap();
let memory = std::fs::read(memory_file).unwrap();
let trace = std::fs::read(trace_file).unwrap();

// Add build dir to path for the duration of the test
let path = std::env::var("PATH").unwrap_or_default();
let build_dir = env!("OUT_DIR");
std::env::set_var("PATH", format!("{build_dir}:{path}"));

#[rstest]
fn test_run_prover(
parsed_prover_test_case: ParsedProverTestCase,
#[from(prover_in_path)] _path: (),
) {
let proof = run_prover(
&public_input,
&memory,
&trace,
&prover_config,
&prover_parameters,
&parsed_prover_test_case.public_input,
&parsed_prover_test_case.memory,
&parsed_prover_test_case.trace,
&parsed_prover_test_case.prover_config,
&parsed_prover_test_case.prover_parameters,
)
.unwrap();

let expected_proof_file = get_fixture_path("fibonacci/fibonacci_proof.json");
let expected_proof = read_proof_file(expected_proof_file);
assert_eq!(proof.proof_hex, expected_proof.proof_hex);
assert_eq!(proof.proof_hex, parsed_prover_test_case.proof.proof_hex);
}

#[rstest]
#[tokio::test]
async fn test_run_prover_async() {
let public_input_file = get_fixture_path("fibonacci/fibonacci_public_input.json");
let prover_config_file = get_fixture_path("fibonacci/cpu_air_prover_config.json");
let parameter_file = get_fixture_path("fibonacci/cpu_air_params.json");
let memory_file = get_fixture_path("fibonacci/fibonacci_memory.bin");
let trace_file = get_fixture_path("fibonacci/fibonacci_trace.bin");

let public_input: PublicInput = read_json_from_file(public_input_file).unwrap();
let prover_config: ProverConfig = read_json_from_file(prover_config_file).unwrap();
let prover_parameters: ProverParameters = read_json_from_file(parameter_file).unwrap();
let memory = std::fs::read(memory_file).unwrap();
let trace = std::fs::read(trace_file).unwrap();

// Add build dir to path for the duration of the test
let path = std::env::var("PATH").unwrap_or_default();
let build_dir = env!("OUT_DIR");
std::env::set_var("PATH", format!("{build_dir}:{path}"));

async fn test_run_prover_async(
parsed_prover_test_case: ParsedProverTestCase,
#[from(prover_in_path)] _path: (),
) {
let proof = run_prover_async(
&public_input,
&memory,
&trace,
&prover_config,
&prover_parameters,
&parsed_prover_test_case.public_input,
&parsed_prover_test_case.memory,
&parsed_prover_test_case.trace,
&parsed_prover_test_case.prover_config,
&parsed_prover_test_case.prover_parameters,
)
.await
.unwrap();

let expected_proof_file = get_fixture_path("fibonacci/fibonacci_proof.json");
let expected_proof = read_proof_file(expected_proof_file);
assert_eq!(proof.proof_hex, expected_proof.proof_hex);
assert_eq!(proof.proof_hex, parsed_prover_test_case.proof.proof_hex);
}
}

0 comments on commit b933115

Please sign in to comment.