import Control.Exception (bracket)
import Data.Either
import qualified Llama.Backend as Backend
import qualified Llama.ChatTemplate as ChatTemplate
import Llama.Context
import Llama.Model
import Test.Tasty (TestTree, defaultMain, testGroup)
import Test.Tasty.HUnit
main :: IO ()
main = defaultMain tests
tests :: TestTree
tests = do
testGroup
"tests"
[ testBackend
, testChatApplyTemplate
, checkContextFunctions
, checkIntFunctions
, checkBoolFunctions
, testModels
-- , synchronizeContextTests ctx
-- , warmupModeTests ctx
-- , batchInitTests
-- , batchGetOneTests
-- , freeBatchTests
-- , encodeDecodeTests ctx
-- , threadCountTests ctx
-- , embeddingsTests ctx
-- , causalAttentionTests ctx
]
testBackend :: TestTree
testBackend =
testGroup
"Llama.Backend tests"
[ testCase "llamaBackendInit and llamaBackendFree should not throw exceptions" $ do
Backend.llamaBackendInit
Backend.llamaBackendFree
return ()
, testCase "llamaBackendInit and llamaBackendFree should be idempotent" $ do
Backend.llamaBackendInit
Backend.llamaBackendFree
Backend.llamaBackendFree
return ()
, testCase "llamaBackendInit can be called multiple times" $ do
Backend.llamaBackendInit
Backend.llamaBackendInit
Backend.llamaBackendFree
Backend.llamaBackendFree
return ()
, testCase "llamaBackendInit and llamaBackendFree should not throw exceptions" $ do
bracket Backend.llamaBackendInit (const Backend.llamaBackendFree) $ \_ -> return ()
]
testChatApplyTemplate :: TestTree
testChatApplyTemplate =
testGroup
"chatApplyTemplate tests"
[ testCase "chatApplyTemplate with no template" $ do
messages <-
ChatTemplate.chatApplyTemplate
Nothing
[ChatTemplate.ChatMessage "role" "content"]
False
4096
assertBool "result is not Left" (isRight messages)
, -- testCase "chatApplyTemplate with custom template" $ do
-- messages <- ChatTemplate.chatApplyTemplate
-- (Just "template") [ChatTemplate.ChatMessage "role" "content"] False 4096
-- assertBool "result is not Left" (isRight messages), -- throwing error as of now
testCase "chatApplyTemplate with empty messages" $ do
messages <- ChatTemplate.chatApplyTemplate Nothing [] False 4096
assertBool "result is not Left" (isRight messages)
]
checkBoolFunctions :: TestTree
checkBoolFunctions =
testGroup
"Boolean Functions"
[ testCase "supportsRpc" $ do
result <- supportsRpc
assertBool "supportsRpc returns a boolean value" (not result) -- My env doesn't support
, testCase "supportsGpuOffload" $ do
result <- supportsGpuOffload
assertBool "supportsGpuOffload returns a boolean value" (not result)
, testCase "supportsMLock" $ do
result <- supportsMLock
assertBool "supportsMLock returns a boolean value" result
, testCase "supportsMMap" $ do
result <- supportsMMap
assertBool "supportsMMap returns a boolean value" result
]
checkIntFunctions :: TestTree
checkIntFunctions =
testGroup
"Integer Functions"
[ testCase "getMaxDevices" $ do
result <- getMaxDevices
assertBool "getMaxDevices returns a non-negative integer" (result >= 0)
, testCase "getTimeUs" $ do
result <- getTimeUs
assertBool "getTimeUs returns a non-negative integer" (result >= 0)
]
checkContextFunctions :: TestTree
checkContextFunctions =
testGroup
"Context Functions"
[ testCase "defaultContextParams" $ do
params <- defaultContextParams
assertBool
"defaultContextParams returns a valid LlamaContextParams"
(not $ null $ show params)
]
testModels :: TestTree
testModels =
testGroup
"Model functions"
[ testCase "loadModelFromFile" testLoadModelFromFile
, testCase "initContextFromModel" testInitContextFromModel
, testCase "getModelVocab" testGetModelVocab
]
testLoadModelFromFile :: Assertion
testLoadModelFromFile = do
params <- defaultModelParams
result <- loadModelFromFile "Qwen3-0.6B-Q4_K_M.gguf" params
case result of
Right _ -> assertBool "model loaded successfully" True
Left err -> assertFailure err
testInitContextFromModel :: Assertion
testInitContextFromModel = do
params <- defaultModelParams
result <- loadModelFromFile "Qwen3-0.6B-Q4_K_M.gguf" params
case result of
Right model -> do
contextParams <- defaultContextParams
result1 <- initContextFromModel model contextParams
case result1 of
Right _ -> assertBool "context initialized successfully" True
Left err -> assertFailure err
Left err -> assertFailure err
testGetModelVocab :: Assertion
testGetModelVocab = do
params <- defaultModelParams
result <- loadModelFromFile "Qwen3-0.6B-Q4_K_M.gguf" params
case result of
Right model -> do
result1 <- getModelVocab model
case result1 of
Right _ -> assertBool "vocab retrieved successfully" True
Left err -> assertFailure err
Left err -> assertFailure err
{-
testGetContextModel :: Assertion
testGetContextModel = do
params <- defaultModelParams
result <- loadModelFromFile "Qwen3-0.6B-Q4_K_M.gguf" params
case result of
Right model -> do
contextParams <- defaultContextParams
result <- initContextFromModel model contextParams
case result of
Right context -> do
model' <- getContextModel context
assertEqual "models are equal" model model'
Left err -> assertFailure err
Left err -> assertFailure err
modelParams <- liftIO defaultModelParams
eModel <- liftIO $ loadModelFromFile "Qwen3-0.6B-Q4_K_M.gguf" modelParams
case eModel of
Left err -> putStrLn $ "Something went wrong: " ++ err
Right model -> do
ctxParams <- defaultContextParams
eCtx <- initContextFromModel model ctxParams
case eCtx of
Left e -> putStrLn $ "Something went wrong" ++ e
Right ctx -> do
batchInitTests :: TestTree
batchInitTests = testGroup "batchInit tests"
[ testCase "batchInit succeeds" $ do
batch <- batchInit 10 256 1
freeBatch (getBatchPtr batch)
]
batchGetOneTests :: TestTree
batchGetOneTests = testGroup "batchGetOne tests"
[ testCase "batchGetOne succeeds" $ do
batch <- batchGetOne [1, 2, 3]
freeBatch (getBatchPtr batch)
]
freeBatchTests :: TestTree
freeBatchTests = testGroup "freeBatch tests"
[ testCase "freeBatch succeeds" $ do
batch <- batchInit 10 256 1
freeBatch (getBatchPtr batch)
-- No assertion, just check if it doesn't crash
]
getBatchPtr :: Batch -> Ptr LlamaBatch
getBatchPtr (Batch ptr) = ptr
encodeDecodeTests :: Context -> TestTree
encodeDecodeTests ctx = testGroup "encodeDecode tests"
[ testCase "encodeBatch succeeds" $ do
batch <- batchInit 10 256 1
result <- encodeBatch ctx batch
assertBool "Encoding failed" (isRight result)
freeBatch (getBatchPtr batch)
-- decodeBatch test
result' <- decodeBatch ctx batch
assertBool "Decoding failed" (isRight result')
]
threadCountTests :: Context -> TestTree
threadCountTests ctx = testGroup "threadCount tests"
[ testCase "setThreadCount succeeds" $ do
setThreadCount ctx 4
count <- getThreadCount ctx
assertBool "Thread count not set" (count == 4)
, testCase "getBatchThreadCount succeeds" $ do
setThreadCounts ctx 4 2
count <- getBatchThreadCount ctx
assertBool "Batch thread count not set" (count == 2)
]
embeddingsTests :: Context -> TestTree
embeddingsTests ctx = testGroup "embeddings tests"
[ testCase "setEmbeddingsEnabled succeeds" $ do
setEmbeddingsEnabled ctx True
enabled <- areEmbeddingsEnabled ctx
assertBool "Embeddings not enabled" enabled
]
causalAttentionTests :: Context -> TestTree
causalAttentionTests ctx = testGroup "causalAttention tests"
[ testCase "setCausalAttention succeeds" $ do
setCausalAttention ctx True
-- No assertion, just check if it doesn't crash
]
warmupModeTests :: Context -> TestTree
warmupModeTests ctx = testGroup "warmupMode tests"
[ testCase "setWarmupMode succeeds" $ do
setWarmupMode ctx True
-- No assertion, just check if it doesn't crash
]
synchronizeContextTests :: Context -> TestTree
synchronizeContextTests ctx = testGroup "synchronizeContext tests"
[ testCase "synchronizeContext succeeds" $ do
synchronizeContext ctx
-- No assertion, just check if it doesn't crash
]
-}