Main.hs (5906B)
1 {-# LANGUAGE OverloadedStrings #-} 2 3 module Main where 4 5 import qualified Control.Exception as E 6 import Control.Monad (forM_) 7 import qualified Data.ByteString as BS 8 import qualified Data.ByteString.Char8 as BSC 9 import qualified Data.Map.Strict as Map 10 import qualified Database.LMDB as L 11 import System.Directory 12 ( createDirectory 13 , doesDirectoryExist 14 , getTemporaryDirectory 15 , removeDirectoryRecursive 16 ) 17 import System.FilePath ((</>)) 18 import Test.Tasty 19 import Test.Tasty.HUnit 20 import Test.Tasty.QuickCheck as QC 21 22 main :: IO () 23 main = defaultMain $ testGroup "ppad-lmdb" 24 [ unitTests 25 , propTests 26 ] 27 28 -- helpers -------------------------------------------------------------------- 29 30 -- | Run an action against a fresh temporary db directory. The 31 -- directory is created (clean) before the action and removed after. 32 withTmpEnv :: L.EnvFlags -> (L.Env -> IO a) -> IO a 33 withTmpEnv envFlags k = do 34 tmpRoot <- getTemporaryDirectory 35 let dir = tmpRoot </> "ppad-lmdb-test" 36 exists <- doesDirectoryExist dir 37 if exists then removeDirectoryRecursive dir else pure () 38 createDirectory dir 39 result <- L.withEnv dir envFlags k 40 removeDirectoryRecursive dir 41 pure result 42 43 -- shared bigger map for tests 44 flags :: L.EnvFlags 45 flags = L.defaultEnvFlags 46 { L.envMapSize = 64 * 1024 * 1024 47 , L.envMaxDbs = 4 48 } 49 50 writeKV 51 :: L.Env -> [(BS.ByteString, BS.ByteString)] -> IO () 52 writeKV env kvs = L.withWriteTxn env $ \txn -> do 53 dbi <- L.openDbi txn Nothing True 54 forM_ kvs $ \(k, v) -> L.put txn dbi k v 55 56 -- unit tests ----------------------------------------------------------------- 57 58 unitTests :: TestTree 59 unitTests = testGroup "unit" 60 [ testCase "put then get roundtrip" $ 61 withTmpEnv flags $ \env -> do 62 writeKV env [("hello", "world")] 63 v <- L.withReadTxn env $ \txn -> do 64 dbi <- L.openDbi txn Nothing False 65 L.get txn dbi "hello" 66 v @?= Just "world" 67 68 , testCase "get on missing key returns Nothing" $ 69 withTmpEnv flags $ \env -> do 70 writeKV env [("a", "1")] 71 v <- L.withReadTxn env $ \txn -> do 72 dbi <- L.openDbi txn Nothing False 73 L.get txn dbi "absent" 74 v @?= Nothing 75 76 , testCase "del returns True only when key existed" $ 77 withTmpEnv flags $ \env -> do 78 writeKV env [("a", "1"), ("b", "2")] 79 (existed, repeated, missing) <- L.withWriteTxn env $ \txn -> do 80 dbi <- L.openDbi txn Nothing False 81 e <- L.del txn dbi "a" 82 e2 <- L.del txn dbi "a" 83 m <- L.del txn dbi "never-existed" 84 pure (e, e2, m) 85 existed @?= True 86 repeated @?= False 87 missing @?= False 88 89 , testCase "write txn aborted on exception" $ 90 withTmpEnv flags $ \env -> do 91 -- attempt that throws after a write 92 let attempt = L.withWriteTxn env $ \txn -> do 93 dbi <- L.openDbi txn Nothing True 94 L.put txn dbi "k" "v" 95 E.throwIO (userError "boom") 96 r <- E.try attempt :: IO (Either IOError ()) 97 case r of 98 Left _ -> pure () 99 Right _ -> assertFailure "expected the user error" 100 -- key must not be visible to a fresh read txn 101 v <- L.withReadTxn env $ \txn -> do 102 dbi <- L.openDbi txn Nothing True 103 L.get txn dbi "k" 104 v @?= Nothing 105 106 , testCase "cursor walks keys in ascending order" $ 107 withTmpEnv flags $ \env -> do 108 let kvs = [ (BSC.pack [c], BSC.pack [c]) 109 | c <- "abcdefghij" ] 110 writeKV env kvs 111 seen <- L.withReadTxn env $ \txn -> do 112 dbi <- L.openDbi txn Nothing False 113 L.withCursor txn dbi (cursorWalk []) 114 seen @?= kvs 115 116 , testCase "cursorSeek lands on first >= key" $ 117 withTmpEnv flags $ \env -> do 118 writeKV env [ ("a","1"), ("c","3"), ("e","5") ] 119 r <- L.withReadTxn env $ \txn -> do 120 dbi <- L.openDbi txn Nothing False 121 L.withCursor txn dbi $ \cur -> 122 L.cursorSeek cur "b" 123 r @?= Just ("c", "3") 124 ] 125 126 -- walk a cursor, accumulating pairs from current position to end 127 cursorWalk 128 :: [(BS.ByteString, BS.ByteString)] 129 -> L.Cursor s 130 -> IO [(BS.ByteString, BS.ByteString)] 131 cursorWalk acc cur = do 132 step <- L.cursorFirst cur 133 go (reverse acc) step 134 where 135 go ys Nothing = pure (reverse ys) 136 go ys (Just kv) = do 137 next <- L.cursorNext cur 138 go (kv : ys) next 139 140 -- properties ----------------------------------------------------------------- 141 142 newtype Key = Key BS.ByteString deriving Show 143 newtype Val = Val BS.ByteString deriving Show 144 145 -- LMDB rejects empty keys; restrict to non-empty bounded strings. 146 instance Arbitrary Key where 147 arbitrary = do 148 n <- QC.chooseInt (1, 32) 149 Key . BS.pack <$> QC.vectorOf n QC.arbitrary 150 151 instance Arbitrary Val where 152 arbitrary = do 153 n <- QC.chooseInt (0, 64) 154 Val . BS.pack <$> QC.vectorOf n QC.arbitrary 155 156 propTests :: TestTree 157 propTests = testGroup "props" 158 [ QC.testProperty "put / get roundtrip" $ \(Key k) (Val v) -> 159 QC.ioProperty $ withTmpEnv flags $ \env -> do 160 writeKV env [(k, v)] 161 got <- L.withReadTxn env $ \txn -> do 162 dbi <- L.openDbi txn Nothing False 163 L.get txn dbi k 164 pure (got == Just v) 165 166 , QC.testProperty "cursor enumerates Map.toAscList" $ \kvs0 -> 167 let kvs = uniqueKeys [(k, v) | (Key k, Val v) <- kvs0] 168 in QC.ioProperty $ withTmpEnv flags $ \env -> do 169 writeKV env kvs 170 seen <- L.withReadTxn env $ \txn -> do 171 dbi <- L.openDbi txn Nothing False 172 L.withCursor txn dbi (cursorWalk []) 173 pure (seen == Map.toAscList (Map.fromList kvs)) 174 ] 175 176 -- keep only the last write for each key, matching LMDB's overwrite 177 -- semantics under MDB_NOOVERWRITE-not-set 178 uniqueKeys 179 :: [(BS.ByteString, BS.ByteString)] 180 -> [(BS.ByteString, BS.ByteString)] 181 uniqueKeys = Map.toList . Map.fromList