{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DataKinds #-}
module Main (main) where
import Lmdb.Types
import Lmdb.Connection
import Lmdb.Internal (runEncoding)
import Data.Monoid (All(..))
import Control.Monad
import qualified Lmdb.Connection as Lmdb
import qualified Lmdb.Codec as Codec
import qualified Lmdb.Map as Map
import qualified Lmdb.Multimap as Multimap
import qualified Data.List as List
import qualified Data.ByteString.Char8 as BC8
import qualified Data.Set as Set
import qualified Data.Map as CMap
import Text.Read (readMaybe)
import Data.ByteString (ByteString)
import Data.Foldable
import Data.Word
import System.Random
import Foreign.Marshal.Alloc
import System.Directory
import Test.HUnit (Assertion,(@?=),assertBool,assertFailure)
import Test.Framework (testGroup, Test, defaultMain)
import Test.QuickCheck (Gen, Arbitrary(..), choose, Property, arbitraryBoundedEnum)
import Test.QuickCheck.Monadic (monadicIO, assert, run, pre, PropertyM)
import Test.Framework.Providers.QuickCheck2 (testProperty)
import Test.Framework.Providers.HUnit (testCase)
import Debug.Trace
import Pipes
import Control.Concurrent
import Control.Concurrent.MVar
import Data.Text (Text)
import qualified Data.Text as Text
import qualified Pipes.Prelude as Pipes
main :: IO ()
main = do
let dir = "data/test/"
createDirectoryIfMissing True dir
defaultMain tests
-- For some reason, this does not end up deleting the
-- directory. Hmm...
-- removeDirectoryRecursive dir
tests :: [Test]
tests =
[ testProperty "Put Get Law (single transaction)" putGetLaw
, testProperty "Put Get Law (separate transactions)" putGetLawSeparate
, testProperty "Text Codec" (propCodecIso Codec.text)
, testProperty "ByteString Codec" (propCodecIso Codec.byteString)
, testProperty "Put Get Law (Text)" textCodecTest
, testProperty "Native Word Ordering" wordOrdering
, testProperty "Custom Color Ordering" customColorOrdering
, testProperty "Multimap Streaming Values at Key" multimapStreamingValues
]
testDbiSettings :: DatabaseSettings Word Word
testDbiSettings = makeSettings
(SortNative NativeSortInteger) Codec.word Codec.word
textualDbiSettings :: DatabaseSettings Text Text
textualDbiSettings = makeSettings
(SortNative NativeSortLexographic) Codec.text Codec.text
wordTextDbiSettings :: DatabaseSettings Word Text
wordTextDbiSettings = makeSettings
(SortNative NativeSortInteger) Codec.word Codec.text
-- multiTestSettings :: MultiDatabaseSettings Word Word
-- multiTestSettings = makeMultiSettings
-- (SortNative NativeSortInteger)
-- (SortNative NativeSortInteger)
-- Codec.word
-- Codec.word
multiColorSettings :: MultiDatabaseSettings Word Color
multiColorSettings = makeMultiSettings
(SortNative NativeSortInteger)
(SortCustom $ CustomSortSafe compare)
Codec.word
(Codec.throughByteString colorToByteString colorFromByteString)
customSortColorSettings :: DatabaseSettings Color Word
customSortColorSettings = makeSettings
(SortCustom $ CustomSortSafe compare)
(Codec.throughByteString colorToByteString colorFromByteString)
Codec.word
data Color = Red | Green | Blue | Purple | Orange | Yellow
deriving (Show,Read,Eq,Ord,Enum,Bounded)
colorToByteString :: Color -> ByteString
colorToByteString = BC8.pack . show
colorFromByteString :: ByteString -> Maybe Color
colorFromByteString = readMaybe . BC8.unpack
putGetLaw :: Word -> Word -> Property
putGetLaw k v = monadicIO $ do
(assert =<<) $ run $ do
withOneDb testDbiSettings $ \env db -> do
withTransaction env $ \txn -> do
Map.insert' txn db k v
m <- Map.lookup' (readonly txn) db k
return (m == Just v)
putGetLawSeparate :: Word -> Word -> Property
putGetLawSeparate k v = monadicIO $ do
(assert =<<) $ run $ do
withOneDb testDbiSettings $ \env db -> do
withTransaction env $ \txn -> do
Map.insert' txn db k v
m <- withTransaction env $ \txn -> do
Map.lookup' (readonly txn) db k
return (m == Just v)
wordOrdering :: [Word] -> Property
wordOrdering ks = monadicIO $ do
(assert =<<) $ run $ do
withOneDb wordTextDbiSettings $ \env db -> do
withTransaction env $ \txn -> do
forM_ ks $ \k -> Map.insertSuccess' txn db k Text.empty
let ksNoDups = Set.toList $ Set.fromList ks
withCursor txn db $ \cur -> do
xs <- Pipes.toListM (Map.firstForward cur)
return (map keyValueKey xs == ksNoDups)
customColorOrdering :: [Color] -> Property
customColorOrdering ks = monadicIO $ do
(assert =<<) $ run $ do
withOneDb customSortColorSettings $ \env db -> do
withTransaction env $ \txn -> do
forM_ ks $ \k -> Map.insertSuccess' txn db k 55
let ksNoDups = Set.toList $ Set.fromList ks
withCursor txn db $ \cur -> do
xs <- Pipes.toListM (Map.firstForward cur)
return (map keyValueKey xs == ksNoDups)
multimapStreamingValues :: [(Word,Color)] -> Property
multimapStreamingValues xs = monadicIO $ do
(assert =<<) $ run $ do
withOneMultiDb multiColorSettings $ \env db -> do
withTransaction env $ \txn -> do
let xs2 = map (\(a,b) -> (a,Set.singleton b)) xs
expected = CMap.toList $ fmap Set.toList $ CMap.fromListWith mappend xs2
withMultiCursor txn db $ \cur -> do
forM_ xs $ \(k,v) -> Multimap.insert cur k v
All correct <- fmap mconcat $ forM expected $ \(k,expVals) -> do
vals <- Pipes.toListM $ Multimap.lookupValues cur k
return $ All $ vals == expVals
return correct
textCodecTest :: Text -> Text -> Property
textCodecTest k v = monadicIO $ do
pre (not $ Text.null k)
pre (Text.length k < 512)
(assertWith =<<) $ run $ do
withOneDb textualDbiSettings $ \env db -> do
withTransaction env $ \txn -> do
Map.insert' txn db k v
withTransaction env $ \txn -> do
m <- Map.lookup' (readonly txn) db k
return $ case m of
Just found -> if found == v
then Nothing
else Just $ concat
[ "Looking up key ["
, Text.unpack k
, "]: expected ["
, Text.unpack v
, "] but got ["
, Text.unpack found
, "]"
]
Nothing -> Just $ concat
[ "Tried to look up key ["
, Text.unpack k
, "] after insertion but no value was found"
]
propCodecIso :: (Eq a,Show a) => Codec s a -> a -> Property
propCodecIso (Codec encoding (Decoding decode)) original = do
let SizedPoke sz f = runEncoding encoding original
monadicIO $ (assertWith =<<) $ run $ allocaBytes (fromIntegral sz) $ \ptr -> do
f ptr
copied <- decode sz ptr
return $ if copied == original
then Nothing
else Just $ concat
[ "Original value: ["
, show original
, "] Decoded value: ["
, show copied
, "]"
]
assertWith :: Monad m => Maybe String -> PropertyM m ()
assertWith Nothing = return ()
assertWith (Just e) = fail e
withOneDb ::
DatabaseSettings k v
-> (Environment 'ReadWrite -> Database k v -> IO a)
-> IO a
withOneDb s1 f = do
dirNum <- randomIO :: IO Word64
let dir = "data/test/" ++ show dirNum
createDirectory dir
env <- initializeReadWriteEnvironment 1000000 5 25 dir
db1 <- withTransaction env $ \txn -> do
openDatabase txn (Just "db1") s1
a <- f env db1
closeDatabase env db1
closeEnvironment env
removeDirectoryRecursive dir
return a
withOneMultiDb ::
MultiDatabaseSettings k v
-> (Environment 'ReadWrite -> MultiDatabase k v -> IO a)
-> IO a
withOneMultiDb s1 f = do
dirNum <- randomIO :: IO Word64
let dir = "data/test/" ++ show dirNum
createDirectory dir
env <- initializeReadWriteEnvironment 1000000 5 25 dir
db1 <- withTransaction env $ \txn -> do
openMultiDatabase txn (Just "db1") s1
a <- f env db1
closeMultiDatabase env db1
closeEnvironment env
removeDirectoryRecursive dir
return a
instance Arbitrary Color where
arbitrary = arbitraryBoundedEnum
instance Arbitrary Text where
arbitrary = fmap Text.pack arbitrary
shrink = map Text.pack . shrink . Text.unpack
instance Arbitrary ByteString where
arbitrary = fmap BC8.pack arbitrary
shrink = map BC8.pack . shrink . BC8.unpack