TestDB.hs

{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

module TestDB (
  WithDB,
  withTestDB, withWritableTestDB, withStackedTestDB,
  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
import Glean.Backend.Types as Backend
import Glean.Typed
import Glean.Schema.Resolve
import qualified Glean.Types as Thrift

import qualified Glean.Schema.Cxx1 as Cxx
import qualified Glean.Schema.GleanTest as Glean.Test
import qualified Glean.Schema.Sys as Sys

import TestData

-- | 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] -> WithDB a -> IO a
withTestDB settings = withWritableTestDB settings . afterComplete

createTestDB :: WithDB ()
createTestDB env repo = do
  kickOffTestDB env repo id
  writeTestDB env repo testFacts

withWritableTestDB :: [Setting] -> WithDB a -> IO a
withWritableTestDB settings action = withEmptyTestDB settings $ \env repo -> do
  writeTestDB env repo testFacts
  action env repo

withStackedTestDB :: [Setting] -> WithDB a -> IO a
withStackedTestDB settings action = withTestEnv settings $ \env -> do
  kickOffTestDB env repo1 id
  writeTestDB env repo1 testFacts1
  completeTestDB env repo1
  kickOffTestDB env repo2 $ \x -> x
    { Thrift.kickOff_dependencies = Just $ stacked repo1 }
  writeTestDB env repo2 testFacts2
  action env repo2
  where
    stacked (Thrift.Repo name hash) =
      Thrift.Dependencies_stacked $ Thrift.Stacked name hash Nothing
    repo1 = Thrift.Repo "dbtest-repo" "1"
    repo2 = Thrift.Repo "dbtest-repo" "2"

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)) ]
  ]

-- | 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 :: WithDB () -> Test
dbTestCaseWritable = toTestCases (dbFlavours [])

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
  backend_schema <- parseAndResolveSchema . Thrift.schemaInfo_schema <$>
    Backend.getSchemaInfo env (Just repo) def
  assertBool "schema1" (isRight backend_schema)

  let allPredicates =
        [ Cxx.allPredicates
        , Glean.Test.allPredicates
        , Sys.allPredicates
        ]
  writeFactsIntoDB env repo allPredicates facts