settable Patch in the testgen

ETCDEVTeam · Oct 23, 2018 · ba69245 · ba69245
1 parent 24f72a0
commit ba69245
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 20 deletions.
diff --git a/jsontests/jsontests-derive/src/attr.rs b/jsontests/jsontests-derive/src/attr.rs
@@ -11,6 +11,7 @@ pub struct Config {
     pub test_with: ExternalRef,
     pub bench_with: Option<ExternalRef>,
     pub criterion_config: Option<ExternalRef>,
+    pub patch: Option<ExternalRef>,
     pub skip: bool,
     pub should_panic: bool,
 }
@@ -47,11 +48,12 @@ pub fn extract_attrs(ast: &syn::DeriveInput) -> Result<Config, Error> {
                 #[directory = \"../tests/testset\"]\n\
                 #[test_with = \"test::test_function\"]\n\
                 #[bench_wuth = \"test::bench_function\"] (Optional)\n\
+                #[patch = \"CustomTestPatch\" (Optional)\n\
                 #[skip] (Optional)\n\
                 #[should_panic] (Optional)\n\
                 struct TestSet;";
 
-    if ast.attrs.len() < 2 || ast.attrs.len() > 5 {
+    if ast.attrs.len() < 2 || ast.attrs.len() > 6 {
         bail!(ERROR_MSG);
     }
 
@@ -63,6 +65,7 @@ pub fn extract_attrs(ast: &syn::DeriveInput) -> Result<Config, Error> {
                     "test_with" => Config { test_with: ExternalRef::from(value.clone()), ..config },
                     "bench_with" => Config { bench_with: Some(ExternalRef::from(value.clone())), ..config },
                     "criterion_config" => Config { criterion_config: Some(ExternalRef::from(value.clone())), ..config },
+                    "patch" => Config { patch: Some(ExternalRef::from(value.clone())), ..config },
                     _ => panic!("{}", ERROR_MSG),
                 }
             },

diff --git a/jsontests/jsontests-derive/src/lib.rs b/jsontests/jsontests-derive/src/lib.rs
@@ -21,7 +21,7 @@ use itertools::Itertools;
 use syn::Ident;
 use proc_macro::TokenStream;
 
-#[proc_macro_derive(JsonTests, attributes(directory, test_with, bench_with, criterion_config, skip, should_panic))]
+#[proc_macro_derive(JsonTests, attributes(directory, test_with, bench_with, criterion_config, skip, should_panic, patch))]
 pub fn json_tests(input: TokenStream) -> TokenStream {
     // Construct a string representation of the type definition
     let s = input.to_string();
@@ -110,6 +110,7 @@ fn generate_test(config: &Config, test_name: &Ident, data: &str, tokens: &mut qu
     let test_func_path = &config.test_with.path;
     let test_func_name = &config.test_with.name;
     let test_name_str = test_name.as_ref();
+    let (patch_name, patch_path) = derive_patch(config);
 
     tokens.append(quote!{#[test]});
     if config.should_panic {
@@ -119,8 +120,9 @@ fn generate_test(config: &Config, test_name: &Ident, data: &str, tokens: &mut qu
     tokens.append(quote! {
         fn #test_name() {
             use #test_func_path;
+            use #patch_path;
             let data = #data;
-            #test_func_name(#test_name_str, data);
+            #test_func_name::<#patch_name>(#test_name_str, data);
         }
     });
 }
@@ -137,11 +139,14 @@ fn generate_bench(config: &Config, test_name: &Ident, data: &str, tokens: &mut q
     let bench_name = format!("bench_{}", test_name.as_ref());
     let bench_ident = Ident::from(bench_name.as_ref());
 
+    let (patch_name, patch_path) = derive_patch(config);
+
     tokens.append(quote! {
         pub fn #bench_ident(c: &mut Criterion) {
             use #bench_func_path;
+            use #patch_path;
             let data = #data;
-            #bench_func_name(c, #bench_name, data);
+            #bench_func_name::<#patch_name>(c, #bench_name, data);
         }
     });
 
@@ -166,3 +171,14 @@ fn generate_criterion_macros(config: &Config, benches: &[Ident], tokens: &mut qu
         tokens.append(template.as_ref().replace("TARGETS", &benches));
     }
 }
+
+fn derive_patch(config: &Config) -> (Ident, Ident) {
+    if let Some(patch) = config.patch.as_ref() {
+        (patch.name.clone(), patch.path.clone())
+    } else {
+        (
+            Ident::from("VMTestPatch"),
+            Ident::from("sputnikvm::VMTestPatch")
+        )
+    }
+}
diff --git a/jsontests/jsontests-derive/src/tests.rs b/jsontests/jsontests-derive/src/tests.rs
@@ -1,4 +1,5 @@
 use std::path::Path;
+use std::iter;
 use std::fs::{self, File, DirEntry};
 use failure::Error;
 
@@ -25,8 +26,14 @@ pub fn read_tests_from_dir<P: AsRef<Path>>(dir_path: P) -> Result<impl Iterator<
     Ok(iter)
 }
 
-pub fn tests_iterator_from_direntry(file: &DirEntry) -> Result<impl Iterator<Item=Test>, Error> {
+pub fn tests_iterator_from_direntry(file: &DirEntry) -> Result<Box<dyn Iterator<Item=Test>>, Error> {
     let path = file.path().to_owned();
+
+    // Skip non-json files
+    if !path.extension().map(|e| e == "json").unwrap_or(false) {
+        return Ok(Box::new(iter::empty()))
+    }
+
     let file = File::open(&path)?;
     let tests: Value = json::from_reader(file)?;
 
@@ -42,5 +49,5 @@ pub fn tests_iterator_from_direntry(file: &DirEntry) -> Result<impl Iterator<Ite
         data
     });
 
-    Ok(iter)
+    Ok(Box::new(iter))
 }
diff --git a/jsontests/src/lib.rs b/jsontests/src/lib.rs
@@ -19,9 +19,9 @@ use std::sync::{Arc, Mutex};
 use bigint::{Gas, M256, U256, H256, Address};
 use hexutil::*;
 use sputnikvm::errors::RequireError;
-use sputnikvm::{VM, SeqContextVM, Context, VMStatus, VMTestPatch};
+use sputnikvm::{VM, SeqContextVM, Context, VMStatus, Patch};
 
-pub fn fire_with_block(machine: &mut SeqContextVM<VMTestPatch>, block: &JSONBlock) {
+pub fn fire_with_block<P: Patch>(machine: &mut SeqContextVM<P>, block: &JSONBlock) {
     loop {
         match machine.fire() {
             Err(RequireError::Account(address)) => {
@@ -60,7 +60,7 @@ pub fn fire_with_block(machine: &mut SeqContextVM<VMTestPatch>, block: &JSONBloc
     }
 }
 
-pub fn apply_to_block(machine: &SeqContextVM<VMTestPatch>, block: &mut JSONBlock) {
+pub fn apply_to_block<P: Patch>(machine: &SeqContextVM<P>, block: &mut JSONBlock) {
     for account in machine.accounts() {
         let account = (*account).clone();
         block.apply_account(account);
@@ -71,16 +71,16 @@ pub fn apply_to_block(machine: &SeqContextVM<VMTestPatch>, block: &mut JSONBlock
     }
 }
 
-pub fn create_machine(v: &Value, block: &JSONBlock) -> SeqContextVM<VMTestPatch> {
+pub fn create_machine<P: Patch>(v: &Value, block: &JSONBlock) -> SeqContextVM<P> {
     let transaction = create_context(v);
 
-    SeqContextVM::<VMTestPatch>::new(transaction, block.block_header())
+    SeqContextVM::<P>::new(transaction, block.block_header())
 }
 
 // TODO: consider refactoring
 #[cfg_attr(feature = "cargo-clippy", allow(cyclomatic_complexity))]
 #[cfg_attr(feature = "cargo-clippy", allow(collapsible_if))]
-pub fn test_machine(v: &Value, machine: &SeqContextVM<VMTestPatch>, block: &JSONBlock, history: &[Context], debug: bool) -> bool {
+pub fn test_machine<P: Patch>(v: &Value, machine: &SeqContextVM<P>, block: &JSONBlock, history: &[Context], debug: bool) -> bool {
     let callcreates = &v["callcreates"];
 
     if callcreates.as_array().is_some() {
@@ -261,13 +261,13 @@ fn is_ok(status: &VMStatus) -> bool {
     }
 }
 
-pub fn test_transaction(_name: &str, v: &Value, debug: bool) -> Result<bool, VMStatus> {
+pub fn test_transaction<P: Patch>(_name: &str, v: &Value, debug: bool) -> Result<bool, VMStatus> {
     let _ = env_logger::try_init();
 
     let mut block = create_block(v);
     let history: Arc<Mutex<Vec<Context>>> = Arc::new(Mutex::new(Vec::new()));
     let history_closure = history.clone();
-    let mut machine = create_machine(v, &block);
+    let mut machine = create_machine::<P>(v, &block);
     machine.add_context_history_hook(move |context| {
         history_closure.lock().unwrap().push(context.clone());
     });
@@ -298,13 +298,13 @@ pub fn test_transaction(_name: &str, v: &Value, debug: bool) -> Result<bool, VMS
 
 use criterion::Criterion;
 
-pub fn bench_transaction(_name: &str, v: Value, c: &mut Criterion) {
+pub fn bench_transaction<P: Patch>(_name: &str, v: Value, c: &mut Criterion) {
     c.bench_function(_name, move |b| {
         b.iter_with_large_setup(|| {
             let block = create_block(&v);
             let history: Arc<Mutex<Vec<Context>>> = Arc::new(Mutex::new(Vec::new()));
             let history_closure = history.clone();
-            let mut machine = create_machine(&v, &block);
+            let mut machine = create_machine::<P>(&v, &block);
             machine.add_context_history_hook(move |context| {
                 history_closure.lock().unwrap().push(context.clone());
             });

diff --git a/jsontests/src/util.rs b/jsontests/src/util.rs
@@ -2,15 +2,16 @@ use serde_json::Value;
 use serde_json as json;
 use test_transaction;
 use bench_transaction;
+use sputnikvm::Patch;
 
-pub fn run_test(name: &str, test: &str) {
+pub fn run_test<P: Patch>(name: &str, test: &str) {
     let test: Value = json::from_str(test).unwrap();
-    assert_eq!(test_transaction(name, &test, true), Ok(true));
+    assert_eq!(test_transaction::<P>(name, &test, true), Ok(true));
 }
 
 use criterion::Criterion;
 
-pub fn run_bench(c: &mut Criterion, name: &'static str, test: &str) {
+pub fn run_bench<P: Patch>(c: &mut Criterion, name: &'static str, test: &str) {
     let test: Value = json::from_str(test).unwrap();
-    bench_transaction(name, test, c);
+    bench_transaction::<P>(name, test, c);
 }