ppad-lmdb-0.1.0: test/Main.hs

{-# LANGUAGE OverloadedStrings #-}

module Main where

import qualified Control.Exception as E
import Control.Monad (forM_)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as BSC
import qualified Data.Map.Strict as Map
import qualified Database.LMDB as L
import System.Directory
  ( createDirectory
  , doesDirectoryExist
  , getTemporaryDirectory
  , removeDirectoryRecursive
  )
import System.FilePath ((</>))
import Test.Tasty
import Test.Tasty.HUnit
import Test.Tasty.QuickCheck as QC

main :: IO ()
main = defaultMain $ testGroup "ppad-lmdb"
  [ unitTests
  , propTests
  ]

-- helpers --------------------------------------------------------------------

-- | Run an action against a fresh temporary db directory. The
--   directory is created (clean) before the action and removed after.
withTmpEnv :: L.EnvFlags -> (L.Env -> IO a) -> IO a
withTmpEnv envFlags k = do
  tmpRoot <- getTemporaryDirectory
  let dir = tmpRoot </> "ppad-lmdb-test"
  exists <- doesDirectoryExist dir
  if exists then removeDirectoryRecursive dir else pure ()
  createDirectory dir
  result <- L.withEnv dir envFlags k
  removeDirectoryRecursive dir
  pure result

-- shared bigger map for tests
flags :: L.EnvFlags
flags = L.defaultEnvFlags
  { L.envMapSize = 64 * 1024 * 1024
  , L.envMaxDbs  = 4
  }

writeKV
  :: L.Env -> [(BS.ByteString, BS.ByteString)] -> IO ()
writeKV env kvs = L.withWriteTxn env $ \txn -> do
  dbi <- L.openDbi txn Nothing True
  forM_ kvs $ \(k, v) -> L.put txn dbi k v

-- unit tests -----------------------------------------------------------------

unitTests :: TestTree
unitTests = testGroup "unit"
  [ testCase "put then get roundtrip" $
      withTmpEnv flags $ \env -> do
        writeKV env [("hello", "world")]
        v <- L.withReadTxn env $ \txn -> do
          dbi <- L.openDbi txn Nothing False
          L.get txn dbi "hello"
        v @?= Just "world"

  , testCase "get on missing key returns Nothing" $
      withTmpEnv flags $ \env -> do
        writeKV env [("a", "1")]
        v <- L.withReadTxn env $ \txn -> do
          dbi <- L.openDbi txn Nothing False
          L.get txn dbi "absent"
        v @?= Nothing

  , testCase "del returns True only when key existed" $
      withTmpEnv flags $ \env -> do
        writeKV env [("a", "1"), ("b", "2")]
        (existed, repeated, missing) <- L.withWriteTxn env $ \txn -> do
          dbi <- L.openDbi txn Nothing False
          e   <- L.del txn dbi "a"
          e2  <- L.del txn dbi "a"
          m   <- L.del txn dbi "never-existed"
          pure (e, e2, m)
        existed  @?= True
        repeated @?= False
        missing  @?= False

  , testCase "write txn aborted on exception" $
      withTmpEnv flags $ \env -> do
        -- attempt that throws after a write
        let attempt = L.withWriteTxn env $ \txn -> do
              dbi <- L.openDbi txn Nothing True
              L.put txn dbi "k" "v"
              E.throwIO (userError "boom")
        r <- E.try attempt :: IO (Either IOError ())
        case r of
          Left _  -> pure ()
          Right _ -> assertFailure "expected the user error"
        -- key must not be visible to a fresh read txn
        v <- L.withReadTxn env $ \txn -> do
          dbi <- L.openDbi txn Nothing True
          L.get txn dbi "k"
        v @?= Nothing

  , testCase "cursor walks keys in ascending order" $
      withTmpEnv flags $ \env -> do
        let kvs = [ (BSC.pack [c], BSC.pack [c])
                  | c <- "abcdefghij" ]
        writeKV env kvs
        seen <- L.withReadTxn env $ \txn -> do
          dbi <- L.openDbi txn Nothing False
          L.withCursor txn dbi (cursorWalk [])
        seen @?= kvs

  , testCase "cursorSeek lands on first >= key" $
      withTmpEnv flags $ \env -> do
        writeKV env [ ("a","1"), ("c","3"), ("e","5") ]
        r <- L.withReadTxn env $ \txn -> do
          dbi <- L.openDbi txn Nothing False
          L.withCursor txn dbi $ \cur ->
            L.cursorSeek cur "b"
        r @?= Just ("c", "3")
  ]

-- walk a cursor, accumulating pairs from current position to end
cursorWalk
  :: [(BS.ByteString, BS.ByteString)]
  -> L.Cursor s
  -> IO [(BS.ByteString, BS.ByteString)]
cursorWalk acc cur = do
  step <- L.cursorFirst cur
  go (reverse acc) step
  where
    go ys Nothing      = pure (reverse ys)
    go ys (Just kv)    = do
      next <- L.cursorNext cur
      go (kv : ys) next

-- properties -----------------------------------------------------------------

newtype Key = Key BS.ByteString deriving Show
newtype Val = Val BS.ByteString deriving Show

-- LMDB rejects empty keys; restrict to non-empty bounded strings.
instance Arbitrary Key where
  arbitrary = do
    n <- QC.chooseInt (1, 32)
    Key . BS.pack <$> QC.vectorOf n QC.arbitrary

instance Arbitrary Val where
  arbitrary = do
    n <- QC.chooseInt (0, 64)
    Val . BS.pack <$> QC.vectorOf n QC.arbitrary

propTests :: TestTree
propTests = testGroup "props"
  [ QC.testProperty "put / get roundtrip" $ \(Key k) (Val v) ->
      QC.ioProperty $ withTmpEnv flags $ \env -> do
        writeKV env [(k, v)]
        got <- L.withReadTxn env $ \txn -> do
          dbi <- L.openDbi txn Nothing False
          L.get txn dbi k
        pure (got == Just v)

  , QC.testProperty "cursor enumerates Map.toAscList" $ \kvs0 ->
      let kvs = uniqueKeys [(k, v) | (Key k, Val v) <- kvs0]
      in  QC.ioProperty $ withTmpEnv flags $ \env -> do
            writeKV env kvs
            seen <- L.withReadTxn env $ \txn -> do
              dbi <- L.openDbi txn Nothing False
              L.withCursor txn dbi (cursorWalk [])
            pure (seen == Map.toAscList (Map.fromList kvs))
  ]

-- keep only the last write for each key, matching LMDB's overwrite
-- semantics under MDB_NOOVERWRITE-not-set
uniqueKeys
  :: [(BS.ByteString, BS.ByteString)]
  -> [(BS.ByteString, BS.ByteString)]
uniqueKeys = Map.toList . Map.fromList