multi-thread test runner

skyzh · skyzh · commit 250c46e99b3f · 2022-06-09T22:36:27.000+08:00
Signed-off-by: Alex Chi &lt;iskyzh@gmail.com&gt;
diff --git a/Cargo.toml b/Cargo.toml
@@ -13,13 +13,14 @@ keywords = ["sql", "database", "planner", "cli"]
 [dependencies]
 anyhow = "1"
 async-trait = "0.1"
+console = "0.15"
+futures-util = { version = "0.3", default-features = false, features = ["alloc"] }
 glob = "0.3"
 libtest-mimic = "0.4"
 serde = { version = "1.0", features = ["derive"] }
 serde_yaml = "0.8"
 similar = "2"
 tokio = { version = "1", features = ["rt", "fs"] }
-console = "0.15"
 
 [workspace]
 members = ["naivedb"]
diff --git a/README.md b/README.md
@@ -19,7 +19,7 @@ Here's an example of test description file:
     SELECT * FROM t1, t2 WHERE t1.v1 = t2.v2;
   desc: Test whether join works correctly.
   before: ["*sql1", "*sql2", "CREATE TABLE t3(v3 int);"]
-  test:
+  tasks:
   - logical
   - physical
 ```
@@ -34,7 +34,7 @@ Basically, it is like:
   - "*test_case_1"           # use *id to reference to another test case
   - "*test_case_2"
   - CREATE TABLE t2(v2 int); # or directly write a SQL here
-  test:                      # run logical and physical test for this case
+  tasks:                      # run logical and physical test for this case
   - logical
   - physical
 ```
diff --git a/naivedb/Cargo.toml b/naivedb/Cargo.toml
@@ -5,9 +5,9 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+anyhow = { version = "1", features = ["backtrace"] }
 async-trait = "0.1"
 sqlplannertest = { path = ".." }
-anyhow = "1"
 tokio = { version = "1", features = ["rt", "rt-multi-thread", "macros", "fs"] }
 
 [[test]]
diff --git a/naivedb/src/lib.rs b/naivedb/src/lib.rs
@@ -1,4 +1,4 @@
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 use async_trait::async_trait;
 
 #[derive(Default)]
@@ -12,7 +12,25 @@ impl NaiveDb {
 
 #[async_trait]
 impl sqlplannertest::PlannerTestRunner for NaiveDb {
-    async fn run(&mut self, _test_case: &sqlplannertest::ParsedTestCase) -> Result<String> {
-        Ok("hello, world!".to_string())
+    async fn run(&mut self, test_case: &sqlplannertest::ParsedTestCase) -> Result<String> {
+        use std::fmt::Write;
+        let mut result = String::new();
+        let r = &mut result;
+        for task in &test_case.tasks {
+            writeln!(r, "=== {}", task)?;
+            writeln!(r, "I'm a naive db, so I don't now how to process")?;
+            for before in &test_case.before_sql {
+                writeln!(r, "{}", before)?;
+            }
+            writeln!(r, "{}", test_case.sql)?;
+            writeln!(r)?;
+        }
+        if test_case.sql.contains("ERROR") {
+            return Err(anyhow!("Ooops, error!"));
+        }
+        if test_case.sql.contains("PANIC") {
+            panic!("I'm panic!");
+        }
+        Ok(result)
     }
 }
diff --git a/naivedb/tests/1.sql b/naivedb/tests/1.sql
diff --git a/naivedb/tests/2.sql b/naivedb/tests/2.sql
diff --git a/naivedb/tests/2.yml b/naivedb/tests/2.yml
diff --git a/naivedb/tests/_panic.yml b/naivedb/tests/_panic.yml
@@ -0,0 +1,6 @@
+- sql: |
+    PANIC
+  desc: Test whether panic works
+  tasks:
+  - logical
+  - physical
diff --git a/naivedb/tests/error.planner.sql b/naivedb/tests/error.planner.sql
@@ -0,0 +1,8 @@
+-- Test whether error works
+ERROR
+
+/*
+Error
+Ooops, error!
+*/
+
diff --git a/naivedb/tests/error.yml b/naivedb/tests/error.yml
@@ -0,0 +1,6 @@
+- sql: |
+    ERROR
+  desc: Test whether error works
+  tasks:
+  - logical
+  - physical
diff --git a/naivedb/tests/ok.planner.sql b/naivedb/tests/ok.planner.sql
@@ -0,0 +1,38 @@
+-- sql1
+CREATE TABLE t1(v1 int);
+
+/*
+
+*/
+
+-- sql2
+CREATE TABLE t2(v2 int);
+
+/*
+
+*/
+
+-- Test whether join works correctly.
+SELECT * FROM t1, t2 WHERE t1.v1 = t2.v2;
+
+/*
+=== logical
+I'm a naive db, so I don't now how to process
+CREATE TABLE t1(v1 int);
+
+CREATE TABLE t2(v2 int);
+
+CREATE TABLE t3(v3 int);
+SELECT * FROM t1, t2 WHERE t1.v1 = t2.v2;
+
+
+=== physical
+I'm a naive db, so I don't now how to process
+CREATE TABLE t1(v1 int);
+
+CREATE TABLE t2(v2 int);
+
+CREATE TABLE t3(v3 int);
+SELECT * FROM t1, t2 WHERE t1.v1 = t2.v2;
+*/
+
diff --git a/naivedb/tests/ok.yml b/naivedb/tests/ok.yml
@@ -8,6 +8,6 @@
     SELECT * FROM t1, t2 WHERE t1.v1 = t2.v2;
   desc: Test whether join works correctly.
   before: ["*sql1", "*sql2", "CREATE TABLE t3(v3 int);"]
-  test:
+  tasks:
   - logical
   - physical
diff --git a/src/apply.rs b/src/apply.rs
@@ -1,55 +1,113 @@
 use std::future::Future;
 use std::path::Path;
 
-use anyhow::{anyhow, Context, Result};
+use anyhow::{anyhow, Context, Error, Result};
+use console::style;
+use futures_util::{stream, StreamExt, TryFutureExt};
 
-use crate::{discover_tests, parse_test_cases, ParsedTestCase, PlannerTestRunner, TestCase};
+use crate::{
+    discover_tests, parse_test_cases, ParsedTestCase, PlannerTestRunner, TestCase, RESULT_SUFFIX,
+};
 
 pub async fn planner_test_apply<F, Ft, R>(path: impl AsRef<Path>, runner_fn: F) -> Result<()>
 where
     F: Fn() -> Ft + Send + Sync + 'static,
     Ft: Future<Output = Result<R>> + Send,
-    R: PlannerTestRunner,
+    R: PlannerTestRunner + 'static,
 {
-    let tests = discover_tests(path)?;
+    let tests = discover_tests(path)?
+        .into_iter()
+        .map(|path| {
+            let path = path?;
+            let filename = path
+                .file_name()
+                .context("unable to extract filename")?
+                .to_os_string();
+            let testname = filename
+                .to_str()
+                .context("unable to convert to string")?
+                .to_string();
+            Ok::<_, Error>((path, testname))
+        })
+        .collect::<Result<Vec<_>, _>>()?;
+    let test_stream = stream::iter(tests).map(|(path, testname)| {
+        let runner_fn = &runner_fn;
+        let testname_x = testname.clone();
+        async {
+            let mut runner = runner_fn().await?;
+            tokio::spawn(async move {
+                let testcases = tokio::fs::read(&path).await?;
+                let testcases: Vec<TestCase> = serde_yaml::from_slice(&testcases)?;
+                let testcases = parse_test_cases(testcases)?;
+                let mut generated_result = String::new();
+                for testcase in testcases {
+                    let runner_result = runner.run(&testcase).await;
+                    generate_result(&testcase, &runner_result, &mut generated_result)?;
+                }
+                let path = {
+                    let mut path = path;
+                    path.set_extension(RESULT_SUFFIX);
+                    path
+                };
+                tokio::fs::write(&path, generated_result).await?;
+
+                Ok::<_, Error>(())
+            })
+            .await??;
+            Ok::<_, Error>(testname)
+        }
+        .map_err(|e| (e, testname_x))
+    });
+
+    let mut test_stream =
+        test_stream.buffer_unordered(std::thread::available_parallelism()?.into());
+
     let mut test_discovered = false;
-    for entry in tests {
-        test_discovered = true;
-        let path = entry.context("failed to read glob entry")?;
-        let filename = path.file_name().context("unable to extract filename")?;
-        let testname = filename.to_str().context("unable to convert to string")?;
-        let mut runner = runner_fn().await?;
-        let testcases = tokio::fs::read(&path).await?;
-        let testcases: Vec<TestCase> = serde_yaml::from_slice(&testcases)?;
-        let testcases = parse_test_cases(testcases)?;
-        let mut generated_result = String::new();
-        println!("{}", testname);
-        for testcase in testcases {
-            let runner_result = runner.run(&testcase).await?;
-            generate_result(&testcase, &runner_result, &mut generated_result)?;
+    let mut failed_cases = vec![];
+    while let Some(item) = test_stream.next().await {
+        match item {
+            Ok(name) => println!("{} {}", style("[DONE]").green().bold(), name),
+            Err((e, name)) => {
+                println!("{} {}: {:#}", style("[FAIL]").red().bold(), name, e);
+                failed_cases.push(name);
+            }
         }
-        let path = {
-            let mut path = path;
-            path.set_extension("sql");
-            path
-        };
-        tokio::fs::write(&path, generated_result).await?;
+        test_discovered = true;
     }
+
     if !test_discovered {
         return Err(anyhow!("no test discovered"));
     }
+
+    if !failed_cases.is_empty() {
+        println!("Failed cases: {:#?}", failed_cases);
+        return Err(anyhow!("Cannot apply planner test"));
+    }
+
     Ok(())
 }
 
 /// Generate a text-based result based on test case and runner result
 pub fn generate_result(
     testcase: &ParsedTestCase,
-    runner_result: &str,
+    runner_result: &Result<String>,
     mut r: impl std::fmt::Write,
 ) -> Result<()> {
-    writeln!(r, "-----")?;
+    match (&testcase.id, &testcase.desc) {
+        (Some(id), Some(desc)) => writeln!(r, "-- {}: {}", id, desc)?,
+        (Some(id), None) => writeln!(r, "-- {}", id)?,
+        (None, Some(desc)) => writeln!(r, "-- {}", desc)?,
+        (None, None) => writeln!(r, "-- (no id or description)")?,
+    }
     writeln!(r, "{}", testcase.sql)?;
-    writeln!(r, "{}", runner_result)?;
+    match runner_result {
+        Ok(runner_result) => {
+            writeln!(r, "/*\n{}\n*/", runner_result.trim_end())?;
+        }
+        Err(err) => {
+            writeln!(r, "/*\nError\n{}\n*/", err)?;
+        }
+    }
     writeln!(r)?;
     Ok(())
 }
diff --git a/src/lib.rs b/src/lib.rs
@@ -19,7 +19,7 @@ pub struct TestCase {
     pub desc: Option<String>,
     pub sql: String,
     pub before: Option<Vec<String>>,
-    pub test: Option<Vec<String>>,
+    pub tasks: Option<Vec<String>>,
 }
 
 /// A parsed test case.
@@ -29,24 +29,25 @@ pub struct ParsedTestCase {
     pub desc: Option<String>,
     pub sql: String,
     pub before_sql: Vec<String>,
-    pub test: Vec<String>,
+    pub tasks: Vec<String>,
 }
 
 /// A planner test runner.
 #[async_trait]
-pub trait PlannerTestRunner {
+pub trait PlannerTestRunner: Send {
     /// Run a test case and return the result
     async fn run(&mut self, test_case: &ParsedTestCase) -> Result<String>;
 }
 
-pub fn parse_test_cases(test: Vec<TestCase>) -> Result<Vec<ParsedTestCase>> {
-    resolve_testcase_id(test)
+pub fn parse_test_cases(tests: Vec<TestCase>) -> Result<Vec<ParsedTestCase>> {
+    resolve_testcase_id(tests)
 }
 
-const SUFFIX: &str = ".yml";
+const TEST_SUFFIX: &str = ".yml";
+const RESULT_SUFFIX: &str = "planner.sql";
 
 pub fn discover_tests(path: impl AsRef<Path>) -> Result<Paths> {
-    let pattern = format!("**/[!_]*{}", SUFFIX);
+    let pattern = format!("**/[!_]*{}", TEST_SUFFIX);
     let path = path.as_ref().join(&pattern);
     let path = path.to_str().context("non utf-8 path")?;
     let paths = glob::glob(path).context("failed to discover test")?;
@@ -70,7 +71,7 @@ mod tests {
     SELECT * FROM t1, t2 WHERE t1.v1 = t2.v2;
   desc: "Test whether join works correctly."
   before: ["*sql1", "*sql2", "CREATE TABLE t3(v3 int);"]
-  test:
+  tasks:
     - logical
     - physical
         "#
diff --git a/src/resolve_id.rs b/src/resolve_id.rs
@@ -54,7 +54,7 @@ pub fn resolve_testcase_id(testcases: Vec<TestCase>) -> Result<Vec<ParsedTestCas
                 id: testcase.id,
                 desc: testcase.desc,
                 sql: testcase.sql,
-                test: testcase.test.unwrap_or_default(),
+                tasks: testcase.tasks.unwrap_or_default(),
             })
         })
         .collect::<Result<Vec<_>>>()
diff --git a/src/test_runner.rs b/src/test_runner.rs

-Original file line number
+Diff line change
@@ @@ -0,0 +1,6 @@ @@
 +- sql: |
 +    PANIC
 +  desc: Test whether panic works
 +  tasks:
 +  - logical
 +  - physical