Speed up tests

glean/test/lib/TestDB.hs

@@ -7,14 +7,19 @@
 -}
 
 module TestDB (
+  WithDB,
   withTestDB, withWritableTestDB, withStackedTestDB,
-  dbTestCase, dbTestCaseWritable, dbTestCaseSettings, createTestDB
+  dbTestCase, dbTestCaseWritable, dbTestCaseSettings, createTestDB,
+  withDbTests,
 ) where
 
 import Data.Default
 import Data.Either
+import Foreign.Marshal.Utils
 import Test.HUnit
 
+import Util.IO
+
 import Glean.Database.Storage (DBVersion(..), currentVersion, writableVersions)
 import Glean.Database.Test
 import Glean.Database.Types
@@ -29,25 +34,28 @@ import qualified Glean.Schema.Sys as Sys
 
 import TestData
 
-afterComplete :: (Env -> Thrift.Repo -> IO a) -> Env -> Thrift.Repo -> IO a
+-- | An action that runs on a test DB
+type WithDB a = Env -> Thrift.Repo -> IO a
+
+afterComplete :: WithDB a -> WithDB a
 afterComplete action env repo = do
   completeTestDB env repo
   action env repo
 
-withTestDB :: [Setting] -> (Env -> Thrift.Repo -> IO a) -> IO a
+withTestDB :: [Setting] -> WithDB a -> IO a
 withTestDB settings = withWritableTestDB settings . afterComplete
 
-createTestDB :: Env -> Thrift.Repo -> IO ()
+createTestDB :: WithDB ()
 createTestDB env repo = do
   kickOffTestDB env repo id
   writeTestDB env repo testFacts
 
-withWritableTestDB :: [Setting] -> (Env -> Thrift.Repo -> IO a) -> IO a
+withWritableTestDB :: [Setting] -> WithDB a -> IO a
 withWritableTestDB settings action = withEmptyTestDB settings $ \env repo -> do
   writeTestDB env repo testFacts
   action env repo
 
-withStackedTestDB :: [Setting] -> (Env -> Thrift.Repo -> IO a) -> IO a
+withStackedTestDB :: [Setting] -> WithDB a -> IO a
 withStackedTestDB settings action = withTestEnv settings $ \env -> do
   kickOffTestDB env repo1 id
   writeTestDB env repo1 testFacts1
@@ -62,30 +70,53 @@ withStackedTestDB settings action = withTestEnv settings $ \env -> do
     repo1 = Thrift.Repo "dbtest-repo" "1"
     repo2 = Thrift.Repo "dbtest-repo" "2"
 
-testCases :: [Setting] -> (Env -> Thrift.Repo -> IO ()) -> Test
-testCases testCaseSettings action = TestList
-  [ TestLabel (label1 ++ label2) $
-    TestCase $ with (settings <> testCaseSettings) action
-    | (label1, with) <-
-        [ ("", withWritableTestDB)
-        , ("stacked/", withStackedTestDB) ]
-    , (label2, settings) <-
+newtype RunTest = RunTest (forall a . WithDB a -> IO a)
+
+toTestCases :: [(String, RunTest)] -> WithDB () -> Test
+toTestCases testCases action = TestList
+  [ TestLabel label $ TestCase $ with action
+    | (label, RunTest with) <- testCases
+  ]
+
+dbFlavours :: [Setting] -> [(String, RunTest)]
+dbFlavours testCaseSettings =
+  [ (label1 ++ label2, run)
+    | (label2, settings) <-
         [ ("memory", [setMemoryStorage])
         , ("rocksdb", [])
         ]
         ++
         [ ("rocksdb-" ++ show (unDBVersion v), [setDBVersion v])
           | v <- writableVersions, v /= currentVersion ]
+    , let allSettings = settings <> testCaseSettings
+    , (label1, run) <-
+        [ ("", RunTest (withWritableTestDB allSettings))
+        , ("stacked/", RunTest (withStackedTestDB allSettings)) ]
   ]
 
-dbTestCase :: (Env -> Thrift.Repo -> IO ()) -> Test
-dbTestCase = testCases [] . afterComplete
+-- | Run a test on several flavour of test DB. For a test suite
+-- with multiple tests, use 'withDbTests' instead.
+dbTestCase :: WithDB () -> Test
+dbTestCase = toTestCases (dbFlavours []) . afterComplete
+
+-- | Like dbTestCase, but the test DBs are shared between multiple
+-- tests, rather than being created afresh for each test. Useful
+-- for speeding up test suites that have many individual tests.
+withDbTests :: ((WithDB () -> Test) -> IO a) -> IO a
+withDbTests fn =
+  withMany lazify (dbFlavours []) $ fn . toTestCases
+  where
+  lazify :: (String, RunTest) -> ((String, RunTest) -> IO a) -> IO a
+  lazify (label, RunTest run) fn =
+    withLazy (run . afterComplete . curry) $ \get ->
+      fn (label, RunTest (\with -> do (env, repo) <- get; with env repo))
 
-dbTestCaseWritable :: (Env -> Thrift.Repo -> IO ()) -> Test
-dbTestCaseWritable = testCases []
+dbTestCaseWritable :: WithDB () -> Test
+dbTestCaseWritable = toTestCases (dbFlavours [])
 
-dbTestCaseSettings :: [Setting] -> (Env -> Thrift.Repo -> IO ()) -> Test
-dbTestCaseSettings settings action = testCases settings (afterComplete action)
+dbTestCaseSettings :: [Setting] -> WithDB () -> Test
+dbTestCaseSettings settings action =
+  toTestCases (dbFlavours settings) (afterComplete action)
 
 writeTestDB :: Env -> Thrift.Repo -> (forall m. NewFact m => m ()) -> IO ()
 writeTestDB env repo facts = do

glean/test/tests/Angle/AngleTest.hs

@@ -34,24 +34,24 @@ import TestData
 import TestDB
 
 main :: IO ()
-main = withUnitTest $ testRunner $ TestList
-  [ TestLabel "angle" $ angleTest id
-  , TestLabel "angle/page" $ angleTest (limit 1)
-  , TestLabel "angleDot" angleDotTest
-  , TestLabel "angleNegation" $ angleNegationTest id
-  , TestLabel "angleNegation/page" $ angleNegationTest (limit 1)
-  , TestLabel "angleIfThenElse" $ angleIfThenElse id
-  , TestLabel "angleIfThenElse/page" $ angleIfThenElse (limit 1)
-  , TestLabel "angleTypeTest" angleTypeTest
+main = withUnitTest $ withDbTests $ \dbTestCase -> testRunner $ TestList
+  [ TestLabel "angle" $ angleTest dbTestCase id
+  , TestLabel "angle/page" $ angleTest dbTestCase (limit 1)
+  , TestLabel "angleDot" $ angleDotTest dbTestCase
+  , TestLabel "angleNegation" $ angleNegationTest dbTestCase id
+  , TestLabel "angleNegation/page" $ angleNegationTest dbTestCase (limit 1)
+  , TestLabel "angleIfThenElse" $ angleIfThenElse dbTestCase id
+  , TestLabel "angleIfThenElse/page" $ angleIfThenElse dbTestCase (limit 1)
+  , TestLabel "angleTypeTest" $ angleTypeTest dbTestCase
   ]
 
 ignorePredK :: Glean.Test.KitchenSink -> Glean.Test.KitchenSink
 ignorePredK k = k {
   Glean.Test.kitchenSink_pred = def,
   Glean.Test.kitchenSink_sum_ = def }
 
-angleTest :: (forall a . Query a -> Query a) -> Test
-angleTest modify = dbTestCase $ \env repo -> do
+angleTest :: (WithDB () -> Test) -> (forall a . Query a -> Query a) -> Test
+angleTest dbTestCase modify = dbTestCase $ \env repo -> do
   -- match zero results
   results <- runQuery_ env repo $ modify $ angle @Sys.Blob
     [s|
@@ -469,6 +469,17 @@ angleTest modify = dbTestCase $ \env repo -> do
        (Nat 6, Nat 2),
        (Nat 8, Nat 3), (Nat 12, Nat 3), (Nat 15, Nat 3)] ] r
 
+  -- Test prim.reverse
+  results <- runQuery_ env repo $ modify $ angle @Glean.Test.Predicate
+    [s|
+        glean.test.Predicate {
+            string_ = X
+        } where
+        X = prim.reverse (prim.reverse X)
+    |]
+  print results
+  assertEqual "angle - string reverse" 4 (length results)
+
   -- Test numeric comparison primitives
   r <- runQuery_ env repo $ angleData @() "prim.gtNat 2 1"
   print r
@@ -597,8 +608,8 @@ angleTest modify = dbTestCase $ \env repo -> do
     [[Nat 1, Nat 2, Nat 1, Nat 2]] r
 
 
-angleDotTest :: Test
-angleDotTest = dbTestCase $ \env repo -> do
+angleDotTest :: (WithDB () -> Test) -> Test
+angleDotTest dbTestCase = dbTestCase $ \env repo -> do
 
   -- record selection
   r <- runQuery_ env repo $ angleData @Text
@@ -697,8 +708,11 @@ angleDotTest = dbTestCase $ \env repo -> do
 
 
 -- if statements
-angleIfThenElse :: (forall a . Query a -> Query a) -> Test
-angleIfThenElse modify = dbTestCase $ \env repo -> do
+angleIfThenElse
+  :: (WithDB () -> Test)
+  -> (forall a . Query a -> Query a)
+  -> Test
+angleIfThenElse dbTestCase modify = dbTestCase $ \env repo -> do
 
   r <- runQuery_ env repo $ modify $ angleData @Nat
     "if never : {} then 1 else 2"
@@ -812,8 +826,11 @@ angleIfThenElse modify = dbTestCase $ \env repo -> do
   assertEqual "reordering disjunctions" 5 (length r)
 
 
-angleNegationTest :: (forall a . Query a -> Query a) -> Test
-angleNegationTest modify = dbTestCase $ \env repo -> do
+angleNegationTest
+  :: (WithDB () -> Test)
+  -> (forall a . Query a -> Query a)
+  -> Test
+angleNegationTest dbTestCase modify = dbTestCase $ \env repo -> do
   -- Negation
 
   -- negating a term fails
@@ -918,8 +935,8 @@ angleNegationTest modify = dbTestCase $ \env repo -> do
   assertEqual "negation -  3" 1 (length r)
 
 -- type checking
-angleTypeTest :: Test
-angleTypeTest = dbTestCase $ \env repo -> do
+angleTypeTest :: (WithDB () -> Test) -> Test
+angleTypeTest dbTestCase = dbTestCase $ \env repo -> do
     -- Test for correct handling of maybe, bool, and enums in the type checker
   r <- runQuery_ env repo $ recursive $ angleData @Nat
     [s|

glean/test/tests/Angle/ArrayTest.hs

@@ -26,19 +26,21 @@ import Glean.Types
 import TestDB
 
 main :: IO ()
-main = withUnitTest $ testRunner $ TestList
-  [ TestLabel "array" $ angleArray id
-  , TestLabel "array/page" $ angleArray (limit 1)
+main = withUnitTest $ withDbTests $ \dbTestCase -> testRunner $ TestList
+  [ TestLabel "array" $ angleArray dbTestCase id
+  , TestLabel "array/page" $ angleArray dbTestCase (limit 1)
   ]
 
-angleArray :: (forall a . Query a -> Query a) -> Test
-angleArray modify = TestList
-  [ TestLabel "generators" $ angleArrayGenerator modify
-  , TestLabel "prefix" $ angleArrayPrefix modify
+angleArray :: (WithDB () -> Test) -> (forall a . Query a -> Query a) -> Test
+angleArray dbTestCase modify = TestList
+  [ TestLabel "generators" $ angleArrayGenerator dbTestCase modify
+  , TestLabel "prefix" $ angleArrayPrefix dbTestCase modify
   ]
 
-angleArrayGenerator :: (forall a . Query a -> Query a) -> Test
-angleArrayGenerator modify = TestList
+angleArrayGenerator
+  :: (WithDB () -> Test)
+  -> (forall a . Query a -> Query a) -> Test
+angleArrayGenerator dbTestCase modify = TestList
   [ TestLabel "array of pred" $
     dbTestCase $ \env repo -> do
     -- fetch all elements of an array
@@ -95,8 +97,10 @@ angleArrayGenerator modify = TestList
       (sort [Byte 1, Byte (fromIntegral (255 :: Word8))]) (sort results)
   ]
 
-angleArrayPrefix :: (forall a . Query a -> Query a) -> Test
-angleArrayPrefix modify = TestList
+angleArrayPrefix
+  :: (WithDB () -> Test)
+  -> (forall a . Query a -> Query a) -> Test
+angleArrayPrefix dbTestCase modify = TestList
   [ TestLabel "nat" $ TestList
     [ TestLabel "nested" $ dbTestCase $ \env repo -> do
         results <- runQuery_ env repo $ modify $ angle @Glean.Test.Predicate

glean/test/tests/Angle/DataTest.hs

@@ -24,13 +24,13 @@ import Glean.Types
 import TestDB
 
 main :: IO ()
-main = withUnitTest $ testRunner $ TestList
-  [ TestLabel "angleData" $ angleDataTest id
-  , TestLabel "angleData/page" $ angleDataTest (limit 1)
+main = withUnitTest $ withDbTests $ \dbTestCase -> testRunner $ TestList
+  [ TestLabel "angleData" $ angleDataTest dbTestCase id
+  , TestLabel "angleData/page" $ angleDataTest dbTestCase (limit 1)
   ]
 
-angleDataTest :: (forall a . Query a -> Query a) -> Test
-angleDataTest modify = dbTestCase $ \env repo -> do
+angleDataTest :: (WithDB () -> Test) -> (forall a . Query a -> Query a) -> Test
+angleDataTest dbTestCase modify = dbTestCase $ \env repo -> do
   results <- runQuery_ env repo $ modify $ angleData
     [s| { 3, false } : { x : nat, y : bool } |]
   assertEqual "angleData 1" [(Nat 3, False)] results

glean/test/tests/Angle/MiscTest.hs

@@ -37,17 +37,17 @@ import Glean.Types
 import TestDB
 
 main :: IO ()
-main = withUnitTest $ testRunner $ TestList
-  [ TestLabel "justKeys" $ justKeys id
-  , TestLabel "justKeys/page" $ justKeys (limit 1)
-  , TestLabel "reorder" reorderTest
-  , TestLabel "scoping" scopingTest
-  , TestLabel "queryOptions" angleQueryOptions
-  , TestLabel "limitBytes" limitTest
+main = withUnitTest $ withDbTests $ \dbTestCase -> testRunner $ TestList
+  [ TestLabel "justKeys" $ justKeys dbTestCase id
+  , TestLabel "justKeys/page" $ justKeys dbTestCase (limit 1)
+  , TestLabel "reorder" $ reorderTest dbTestCase
+  , TestLabel "scoping" $ scopingTest dbTestCase
+  , TestLabel "queryOptions" $ angleQueryOptions dbTestCase
+  , TestLabel "limitBytes" $ limitTest dbTestCase
   , TestLabel "fullScans" fullScansTest
   , TestLabel "newold" $ newOldTest id
-  , TestLabel "justCheck" justCheckTest
-  , TestLabel "warnDiag" warnDiagTest
+  , TestLabel "justCheck" $ justCheckTest dbTestCase
+  , TestLabel "warnDiag" $ warnDiagTest dbTestCase
   ]
 
 newOldTest :: (forall a . Query a -> Query a) -> Test
@@ -64,8 +64,8 @@ newOldTest modify = TestCase $ withStackedTestDB [] $ \env repo -> do
     [ "abba", "anonymous", "azimuth", "blubber", "book", "foo" ]
     results
 
-warnDiagTest :: Test
-warnDiagTest = dbTestCase $ \env repo -> do
+warnDiagTest :: (WithDB () -> Test) -> Test
+warnDiagTest dbTestCase = dbTestCase $ \env repo -> do
   let (Query q1) = dbgPredHasFacts $ Angle.query $
         predicate @Glean.Test.EmptyPred wild
       (Query q2) = Angle.query $ predicate @Cxx.Name wild
@@ -78,8 +78,8 @@ warnDiagTest = dbTestCase $ \env repo -> do
   where getWarns r = filter (\d -> "Warning" `isPrefixOf` Text.unpack d)
              $ userQueryResults_diagnostics r
 
-scopingTest :: Test
-scopingTest = dbTestCase $ \env repo -> do
+scopingTest :: (WithDB () -> Test) -> Test
+scopingTest dbTestCase = dbTestCase $ \env repo -> do
   r <- try $ runQuery_ env repo $ angle @Glean.Test.Predicate
     [s|
       cxx1.Name X
@@ -142,8 +142,8 @@ scopingTest = dbTestCase $ \env repo -> do
   print (r :: [Glean.Test.NothingTest])
   assertEqual "angle - nothingTest" 1 (length r)
 
-justKeys :: (forall a . Query a -> Query a) -> Test
-justKeys modify = dbTestCase $ \env repo -> do
+justKeys :: (WithDB () -> Test) -> (forall a . Query a -> Query a) -> Test
+justKeys dbTestCase modify = dbTestCase $ \env repo -> do
   results <- runQuery_ env repo $ modify $ keys $ allFacts @Cxx.Name
   assertEqual "angle - justKeys" 11 (length results)
   assertBool "angle - justKeys" $ "abba" `elem` results
@@ -189,8 +189,8 @@ factsSearched ref lookupPid maybeStats = do
       \ /
       "d"
 -}
-reorderTest :: Test
-reorderTest = dbTestCase $ \env repo -> do
+reorderTest :: (WithDB () -> Test) -> Test
+reorderTest dbTestCase = dbTestCase $ \env repo -> do
   si <- getSchemaInfo env (Just repo) def { getSchemaInfo_omit_source = True }
   let lookupPid = Map.fromList
         [ (ref,pid) | (pid,ref) <- Map.toList (schemaInfo_predicateIds si) ]
@@ -391,8 +391,8 @@ reorderTest = dbTestCase $ \env repo -> do
     |]
   assertEqual "reorder unbound choice" 3 (length results)
 
-angleQueryOptions :: Test
-angleQueryOptions = dbTestCase $ \env repo -> do
+angleQueryOptions :: (WithDB () -> Test) -> Test
+angleQueryOptions dbTestCase = dbTestCase $ \env repo -> do
   let omitResults :: Bool -> Query q -> Query q
       omitResults omit (Query query) = Query q
         where
@@ -421,8 +421,8 @@ angleQueryOptions = dbTestCase $ \env repo -> do
   assertEqual "queryOptions - not omitting results" (counts r) (2, 2)
 
 
-limitTest :: Test
-limitTest = dbTestCase $ \env repo -> do
+limitTest :: (WithDB () -> Test) -> Test
+limitTest dbTestCase = dbTestCase $ \env repo -> do
   (results, truncated) <- runQuery env repo $ limitBytes 30 $ recursive $
     Angle.query $ predicate @Glean.Test.Edge wild
   -- each Edge will be
@@ -437,8 +437,8 @@ limitTest = dbTestCase $ \env repo -> do
     Angle.query $ predicate @Glean.Test.Edge wild
   assertBool "limitBytes" (length results == 2 && truncated)
 
-justCheckTest :: Test
-justCheckTest = dbTestCase $ \env repo -> do
+justCheckTest :: (WithDB () -> Test) -> Test
+justCheckTest dbTestCase = dbTestCase $ \env repo -> do
   results <- runQuery_ env repo $ justCheck $ Angle.query $
     predicate @Cxx.Name wild
   assertEqual "just check" 0 (length results)