From 4cbfc4529294fcf3395f3d96b06c045827cd44c1 Mon Sep 17 00:00:00 2001
From: Leni Aniva <v@leni.sh>
Date: Sat, 6 Apr 2024 14:07:13 -0700
Subject: [PATCH] test: Parallel testing infrastructure

---
 Test/Common.lean      | 29 +++++++++++++++++++++++++++--
 Test/Environment.lean | 11 ++++++-----
 Test/Integration.lean | 16 ++++++++--------
 Test/Library.lean     | 14 +++++---------
 Test/Main.lean        | 25 +++++++++++++++----------
 Test/Metavar.lean     | 13 ++-----------
 Test/Proofs.lean      | 14 ++------------
 Test/Serial.lean      | 20 ++++++++------------
 8 files changed, 73 insertions(+), 69 deletions(-)

diff --git a/Test/Common.lean b/Test/Common.lean
index 378dce8..2711ff2 100644
--- a/Test/Common.lean
+++ b/Test/Common.lean
@@ -8,6 +8,7 @@ open Lean
 
 namespace Pantograph
 
+-- Auxiliary functions
 namespace Protocol
 /-- Set internal names to "" -/
 def Goal.devolatilize (goal: Goal): Goal :=
@@ -37,9 +38,32 @@ def TacticResult.toString : TacticResult → String
   | .parseError error => s!".parseError {error}"
   | .indexError index => s!".indexError {index}"
 
-def assertUnreachable (message: String): LSpec.TestSeq := LSpec.check message false
+namespace Test
 
-open Lean
+def expectationFailure (desc: String) (error: String): LSpec.TestSeq := LSpec.test desc (LSpec.ExpectationFailure "ok _" error)
+
+-- Test running infrastructure
+
+def addPrefix (pref: String) (tests: List (String × α)): List (String  × α) :=
+  tests.map (λ (name, x) => (pref ++ "/" ++ name, x))
+
+/-- Runs test in parallel. Filters test name if given -/
+def runTestGroup (filter: Option String) (tests: List (String × IO LSpec.TestSeq)): IO LSpec.TestSeq := do
+  let tests: List (String × IO LSpec.TestSeq) := match filter with
+    | .some filter => tests.filter (λ (name, _) => filter.isPrefixOf name)
+    | .none => tests
+  let tasks: List (String × Task _) ← tests.mapM (λ (name, task) => do
+    return (name, ← EIO.asTask task))
+  let all := tasks.foldl (λ acc (name, task) =>
+    let v: Except IO.Error LSpec.TestSeq := Task.get task
+    match v with
+    | .ok case => acc ++ (LSpec.group name case)
+    | .error e => acc ++ (expectationFailure name e.toString)
+    ) LSpec.TestSeq.done
+  return all
+
+
+def assertUnreachable (message: String): LSpec.TestSeq := LSpec.check message false
 
 def runCoreMSeq (env: Environment) (coreM: CoreM LSpec.TestSeq) (options: Array String := #[]): IO LSpec.TestSeq := do
   let coreContext: Core.Context ← createCoreContext options
@@ -59,5 +83,6 @@ def defaultTermElabMContext: Lean.Elab.Term.Context := {
     declName? := some "_pantograph".toName,
     errToSorry := false
   }
+end Test
 
 end Pantograph
diff --git a/Test/Environment.lean b/Test/Environment.lean
index 7a398da..927793d 100644
--- a/Test/Environment.lean
+++ b/Test/Environment.lean
@@ -94,10 +94,11 @@ def test_inspect: IO LSpec.TestSeq := do
     ) LSpec.TestSeq.done
   runCoreMSeq env inner
 
-def suite: IO LSpec.TestSeq := do
-  return LSpec.group "Environment" $
-    (LSpec.group "Catalog" (← test_catalog)) ++
-    (LSpec.group "Symbol visibility" (← test_symbol_visibility)) ++
-    (LSpec.group "Inspect" (← test_inspect))
+def suite: List (String × IO LSpec.TestSeq) :=
+  [
+    ("Catalog", test_catalog),
+    ("Symbol Visibility", test_symbol_visibility),
+    ("Inspect", test_inspect),
+  ]
 
 end Pantograph.Test.Environment
diff --git a/Test/Integration.lean b/Test/Integration.lean
index 5f2523b..4f3bcba 100644
--- a/Test/Integration.lean
+++ b/Test/Integration.lean
@@ -112,7 +112,7 @@ def test_tactic : IO LSpec.TestSeq :=
       Protocol.GoalTacticResult))
   ]
 
-def test_env : IO LSpec.TestSeq :=
+def test_env_add_inspect : IO LSpec.TestSeq :=
   let name1 := "Pantograph.mystery"
   let name2 := "Pantograph.mystery2"
   subroutine_runner [
@@ -148,13 +148,13 @@ def test_env : IO LSpec.TestSeq :=
       Protocol.EnvInspectResult))
   ]
 
-def suite: IO LSpec.TestSeq := do
-
-  return LSpec.group "Integration" $
-    (LSpec.group "Option modify" (← test_option_modify)) ++
-    (LSpec.group "Malformed command" (← test_malformed_command)) ++
-    (LSpec.group "Tactic" (← test_tactic)) ++
-    (LSpec.group "Env" (← test_env))
+def suite: List (String × IO LSpec.TestSeq) :=
+  [
+    ("Option modify", test_option_modify),
+    ("Malformed command", test_malformed_command),
+    ("Tactic", test_tactic),
+    ("env.add env.inspect", test_env_add_inspect),
+  ]
 
 
 end Pantograph.Test.Integration
diff --git a/Test/Library.lean b/Test/Library.lean
index 8c935ee..d995374 100644
--- a/Test/Library.lean
+++ b/Test/Library.lean
@@ -8,11 +8,7 @@ open Pantograph
 
 namespace Pantograph.Test.Library
 
-def test_expr_echo: IO LSpec.TestSeq := do
-  let env: Environment ← importModules
-    (imports := #[`Init])
-    (opts := {})
-    (trustLevel := 1)
+def test_expr_echo (env: Environment): IO LSpec.TestSeq := do
   let inner: CoreM LSpec.TestSeq := do
     let prop_and_proof := "⟨∀ (x: Prop), x → x, λ (x: Prop) (h: x) => h⟩"
     let tests := LSpec.TestSeq.done
@@ -34,9 +30,9 @@ def test_expr_echo: IO LSpec.TestSeq := do
     return tests
   runCoreMSeq env (options := #["pp.proofs.threshold=100"]) inner
 
-def suite: IO LSpec.TestSeq := do
-
-  return LSpec.group "Library" $
-    (LSpec.group "ExprEcho" (← test_expr_echo))
+def suite (env: Environment): List (String × IO LSpec.TestSeq) :=
+  [
+    ("expr_echo", test_expr_echo env),
+  ]
 
 end Pantograph.Test.Library
diff --git a/Test/Main.lean b/Test/Main.lean
index 6baffad..42b9d5e 100644
--- a/Test/Main.lean
+++ b/Test/Main.lean
@@ -8,16 +8,21 @@ import Test.Serial
 
 open Pantograph.Test
 
-def main := do
+def main (args: List String) := do
+  let name_filter := args.head?
   Lean.initSearchPath (← Lean.findSysroot)
+  let env_default: Lean.Environment ← Lean.importModules
+    (imports := #[`Init])
+    (opts := {})
+    (trustLevel := 1)
 
-  let suites := [
-    Environment.suite,
-    Integration.suite,
-    Library.suite,
-    Metavar.suite,
-    Proofs.suite,
-    Serial.suite,
+  let suites: List (String × List (String × IO LSpec.TestSeq)) := [
+    ("Environment", Environment.suite),
+    ("Integration", Integration.suite),
+    ("Library", Library.suite env_default),
+    ("Metavar", Metavar.suite env_default),
+    ("Proofs", Proofs.suite env_default),
+    ("Serial", Serial.suite env_default),
   ]
-  let all ← suites.foldlM (λ acc m => do pure $ acc ++ (← m)) LSpec.TestSeq.done
-  LSpec.lspecIO $ all
+  let tests: List (String × IO LSpec.TestSeq) := suites.foldl (λ acc (name, suite) => acc ++ (addPrefix name suite)) []
+  LSpec.lspecIO (← runTestGroup name_filter tests)
diff --git a/Test/Metavar.lean b/Test/Metavar.lean
index 9595b35..433326d 100644
--- a/Test/Metavar.lean
+++ b/Test/Metavar.lean
@@ -268,11 +268,7 @@ def test_partial_continuation: TestM Unit := do
   return ()
 
 
-def suite: IO LSpec.TestSeq := do
-  let env: Lean.Environment ← Lean.importModules
-    (imports := #["Init"].map (λ str => { module := str.toName, runtimeOnly := false }))
-    (opts := {})
-    (trustLevel := 1)
+def suite (env: Environment): List (String × IO LSpec.TestSeq) :=
   let tests := [
     ("Instantiate", test_instantiate_mvar),
     ("2 < 5", test_m_couple),
@@ -280,11 +276,6 @@ def suite: IO LSpec.TestSeq := do
     ("Proposition Generation", test_proposition_generation),
     ("Partial Continuation", test_partial_continuation)
   ]
-  let tests ← tests.foldlM (fun acc tests => do
-    let (name, tests) := tests
-    let tests ← proofRunner env tests
-    return acc ++ (LSpec.group name tests)) LSpec.TestSeq.done
-
-  return LSpec.group "Metavar" tests
+  tests.map (fun (name, test) => (name, proofRunner env test))
 
 end Pantograph.Test.Metavar
diff --git a/Test/Proofs.lean b/Test/Proofs.lean
index 133eee0..a4a1927 100644
--- a/Test/Proofs.lean
+++ b/Test/Proofs.lean
@@ -293,12 +293,7 @@ def proof_or_comm: TestM Unit := do
       ]
     }
 
-
-def suite: IO LSpec.TestSeq := do
-  let env: Lean.Environment ← Lean.importModules
-    (imports := #[{ module := "Init".toName, runtimeOnly := false}])
-    (opts := {})
-    (trustLevel := 1)
+def suite (env: Environment): List (String × IO LSpec.TestSeq) :=
   let tests := [
     ("Nat.add_comm", proof_nat_add_comm false),
     ("Nat.add_comm manual", proof_nat_add_comm true),
@@ -306,11 +301,6 @@ def suite: IO LSpec.TestSeq := do
     ("arithmetic", proof_arith),
     ("Or.comm", proof_or_comm)
   ]
-  let tests ← tests.foldlM (fun acc tests => do
-    let (name, tests) := tests
-    let tests ← proofRunner env tests
-    return acc ++ (LSpec.group name tests)) LSpec.TestSeq.done
-
-  return LSpec.group "Proofs" tests
+  tests.map (fun (name, test) => (name, proofRunner env test))
 
 end Pantograph.Test.Proofs
diff --git a/Test/Serial.lean b/Test/Serial.lean
index f186bbb..0a46acc 100644
--- a/Test/Serial.lean
+++ b/Test/Serial.lean
@@ -72,17 +72,13 @@ def test_instance (env: Environment): IO LSpec.TestSeq := do
     return LSpec.TestSeq.done
   runMetaMSeq env metaM
 
-def suite: IO LSpec.TestSeq := do
-  let env: Environment ← importModules
-    (imports := #["Init"].map (λ str => { module := str.toName, runtimeOnly := false }))
-    (opts := {})
-    (trustLevel := 1)
-
-  return LSpec.group "Serialization" $
-    (LSpec.group "name_to_ast" test_name_to_ast) ++
-    (LSpec.group "Expression binder" (← test_expr_to_binder env)) ++
-    (LSpec.group "Sexp from symbol" (← test_sexp_of_symbol env)) ++
-    (LSpec.group "Sexp from expr" (← test_sexp_of_expr env)) ++
-    (LSpec.group "Instance" (← test_instance env))
+def suite (env: Environment): List (String × IO LSpec.TestSeq) :=
+  [
+    ("name_to_ast", do pure test_name_to_ast),
+    ("Expression binder", test_expr_to_binder env),
+    ("Sexp from symbol", test_sexp_of_symbol env),
+    ("Sexp from expr", test_sexp_of_expr env),
+    ("Instance", test_instance env),
+  ]
 
 end Pantograph.Test.Serial