bloodhound-1.0.0.0: tests/Test/KnnModelSpec.hs
{-# LANGUAGE OverloadedStrings #-}
module Test.KnnModelSpec (spec) where
import Data.ByteString.Char8 qualified as BS
import Data.ByteString.Lazy.Char8 qualified as LBS
import Database.Bloodhound.OpenSearch3.Requests qualified as OS3Requests
import Database.Bloodhound.OpenSearch3.Types
import TestsUtils.Import
import Prelude
-- | A representative response body for @GET /_plugins/_knn/models/{model_id}@,
-- drawn from the OpenSearch vector-search API documentation. @model_blob@ is
-- truncated; the real value is a large base64 serialization of the trained
-- graph.
sampleFullResponse :: LBS.ByteString
sampleFullResponse =
"{\
\ \"model_id\" : \"test-model\",\
\ \"model_blob\" : \"SXdGbIAAAAAAAAAAAA...\",\
\ \"state\" : \"created\",\
\ \"timestamp\" : \"2021-11-15T18:45:07.505369036Z\",\
\ \"description\" : \"Default\",\
\ \"error\" : \"\",\
\ \"space_type\" : \"l2\",\
\ \"dimension\" : 128,\
\ \"engine\" : \"faiss\"\
\}"
-- | A metadata-only response: a caller that passes @?filter_path=model_id,state@
-- (or @_source_excludes=model_blob@) receives a subset. The parser must
-- tolerate any field being absent.
samplePartialResponse :: LBS.ByteString
samplePartialResponse =
"{\
\ \"model_id\" : \"test-model\",\
\ \"state\" : \"created\"\
\}"
-- | A model still undergoing training has no @model_blob@ yet and reports a
-- @training@ state.
sampleTrainingResponse :: LBS.ByteString
sampleTrainingResponse =
"{\
\ \"model_id\" : \"train-model\",\
\ \"state\" : \"training\"\
\}"
-- | A representative @POST /_plugins/_knn/models/_search@ response: the
-- standard OpenSearch search envelope wrapping hits from
-- @.opensearch-knn-models@. The first hit carries a full @_source@; the
-- second is metadata-only, mirroring the documented
-- @_source_excludes=model_blob@ read. Both must decode into 'KnnModel'
-- under 'SearchResult'.
sampleSearchModelsResponse :: LBS.ByteString
sampleSearchModelsResponse =
"{\
\ \"took\" : 3,\
\ \"timed_out\" : false,\
\ \"_shards\" : {\"total\" : 1, \"successful\" : 1, \"skipped\" : 0, \"failed\" : 0},\
\ \"hits\" : {\
\ \"total\" : {\"value\" : 2, \"relation\" : \"eq\"},\
\ \"max_score\" : 1.0,\
\ \"hits\" : [\
\ {\
\ \"_index\" : \".opensearch-knn-models\",\
\ \"_id\" : \"test-model\",\
\ \"_score\" : 1.0,\
\ \"_source\" : {\
\ \"model_id\" : \"test-model\",\
\ \"state\" : \"created\",\
\ \"space_type\" : \"l2\",\
\ \"dimension\" : 128,\
\ \"engine\" : \"faiss\"\
\ }\
\ },\
\ {\
\ \"_index\" : \".opensearch-knn-models\",\
\ \"_id\" : \"meta-only\",\
\ \"_score\" : 1.0,\
\ \"_source\" : {\
\ \"model_id\" : \"meta-only\",\
\ \"state\" : \"training\"\
\ }\
\ }\
\ ]\
\ }\
\}"
spec :: Spec
spec = describe "k-NN Get Model API" $ do
describe "KnnModelState JSON" $ do
it "encodes the three documented states as bare strings" $ do
encode KnnModelStateCreated `shouldBe` "\"created\""
encode KnnModelStateFailed `shouldBe` "\"failed\""
encode KnnModelStateTraining `shouldBe` "\"training\""
it "decodes the three documented states" $ do
decode "\"created\"" `shouldBe` Just KnnModelStateCreated
decode "\"failed\"" `shouldBe` Just KnnModelStateFailed
decode "\"training\"" `shouldBe` Just KnnModelStateTraining
it "rejects an unknown state" $ do
let decoded = decode "\"loaded\"" :: Maybe KnnModelState
decoded `shouldBe` Nothing
describe "KnnEngine JSON (OpenSearchKnnEngine reuse)" $ do
it "encodes the documented engines as bare strings" $ do
encode KnnEngineFaiss `shouldBe` "\"faiss\""
encode KnnEngineNmslib `shouldBe` "\"nmslib\""
encode KnnEngineLucene `shouldBe` "\"lucene\""
it "decodes the documented engines" $ do
decode "\"faiss\"" `shouldBe` Just KnnEngineFaiss
decode "\"nmslib\"" `shouldBe` Just KnnEngineNmslib
decode "\"lucene\"" `shouldBe` Just KnnEngineLucene
it "rejects an unknown engine" $ do
let decoded = decode "\"annoy\"" :: Maybe OpenSearchKnnEngine
decoded `shouldBe` Nothing
describe "KnnModel JSON" $ do
it "decodes a fully-populated response" $ do
let Just decoded = decode sampleFullResponse :: Maybe KnnModel
knnModelModelId decoded `shouldBe` "test-model"
knnModelState decoded `shouldBe` Just KnnModelStateCreated
knnModelDimension decoded `shouldBe` Just 128
knnModelEngine decoded `shouldBe` Just KnnEngineFaiss
knnModelModelBlob decoded `shouldBe` Just "SXdGbIAAAAAAAAAAAA..."
it "decodes a metadata-only (filter_path) response" $ do
let Just decoded = decode samplePartialResponse :: Maybe KnnModel
knnModelModelId decoded `shouldBe` "test-model"
knnModelState decoded `shouldBe` Just KnnModelStateCreated
knnModelModelBlob decoded `shouldBe` Nothing
knnModelDimension decoded `shouldBe` Nothing
knnModelEngine decoded `shouldBe` Nothing
it "decodes a training-state model without a blob" $ do
let Just decoded = decode sampleTrainingResponse :: Maybe KnnModel
knnModelModelId decoded `shouldBe` "train-model"
knnModelState decoded `shouldBe` Just KnnModelStateTraining
knnModelModelBlob decoded `shouldBe` Nothing
it "decodes a failed-state model carrying a non-empty error" $ do
let Just decoded =
decode
"{\
\ \"model_id\":\"bad-model\",\
\ \"state\":\"failed\",\
\ \"error\":\"training failed: cluster timeout\"\
\}" ::
Maybe KnnModel
knnModelModelId decoded `shouldBe` "bad-model"
knnModelState decoded `shouldBe` Just KnnModelStateFailed
knnModelError decoded `shouldBe` Just "training failed: cluster timeout"
it "tolerates unknown fields (forward compatibility)" $ do
let Just decoded = decode "{\"model_id\":\"x\",\"future_field\":42}" :: Maybe KnnModel
knnModelModelId decoded `shouldBe` "x"
it "requires the model_id field" $ do
let decoded = decode "{\"state\":\"created\"}" :: Maybe KnnModel
decoded `shouldBe` Nothing
it "decodes an empty error string as Just \"\"" $ do
let Just decoded = decode sampleFullResponse :: Maybe KnnModel
knnModelError decoded `shouldBe` Just ""
it "rejects a top-level array" $ do
let decoded = decode "[]" :: Maybe KnnModel
decoded `shouldBe` Nothing
it "rejects malformed JSON" $ do
let decoded = decode "{ not json" :: Maybe KnnModel
decoded `shouldBe` Nothing
it "round-trips a full body through ToJSON/FromJSON" $ do
let Just decoded = decode sampleFullResponse :: Maybe KnnModel
(decode . encode) decoded `shouldBe` Just decoded
it "round-trips a partial body through ToJSON/FromJSON" $ do
let Just decoded = decode samplePartialResponse :: Maybe KnnModel
(decode . encode) decoded `shouldBe` Just decoded
it "ToJSON omits Nothing fields (no null leakage)" $ do
let Just decoded = decode samplePartialResponse :: Maybe KnnModel
encoded = LBS.toStrict (encode decoded)
encoded `shouldSatisfy` not . BS.isInfixOf "\"model_blob\""
encoded `shouldSatisfy` not . BS.isInfixOf "\"dimension\""
it "ToJSON emits a field when it is present" $ do
let Just decoded = decode sampleFullResponse :: Maybe KnnModel
encoded = LBS.toStrict (encode decoded)
encoded `shouldSatisfy` BS.isInfixOf "\"engine\""
encoded `shouldSatisfy` BS.isInfixOf "\"dimension\""
describe "getKnnModel endpoint shape" $ do
-- Pure checks against the BHRequest shape; no live backend needed.
it "GETs /_plugins/_knn/models/{model_id} for the given ModelId" $ do
let req = OS3Requests.getKnnModel (ModelId "test-model")
getRawEndpoint (bhRequestEndpoint req)
`shouldBe` ["_plugins", "_knn", "models", "test-model"]
getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []
it "uses the model id as a single opaque path segment" $ do
let req = OS3Requests.getKnnModel (ModelId "faiss-ivf-pq_42")
getRawEndpoint (bhRequestEndpoint req)
`shouldBe` ["_plugins", "_knn", "models", "faiss-ivf-pq_42"]
it "uses a distinct id in the path when given a different ModelId" $ do
let req = OS3Requests.getKnnModel (ModelId "other-model")
getRawEndpoint (bhRequestEndpoint req)
`shouldBe` ["_plugins", "_knn", "models", "other-model"]
it "uses the GET method" $ do
let req = OS3Requests.getKnnModel (ModelId "test-model")
bhRequestMethod req `shouldBe` "GET"
it "does not attach a body (GET semantics)" $ do
let req = OS3Requests.getKnnModel (ModelId "test-model")
bhRequestBody req `shouldBe` Nothing
describe "KnnDeleteModelResponse JSON" $ do
it "decodes the documented {model_id, result} response" $ do
let Just decoded = decode "{\"model_id\":\"m\",\"result\":\"deleted\"}" :: Maybe KnnDeleteModelResponse
knnDeleteModelResponseModelId decoded `shouldBe` "m"
knnDeleteModelResponseResult decoded `shouldBe` "deleted"
it "echoes the model_id passed in the path" $ do
let Just decoded = decode "{\"model_id\":\"faiss-ivf_42\",\"result\":\"deleted\"}" :: Maybe KnnDeleteModelResponse
knnDeleteModelResponseModelId decoded `shouldBe` "faiss-ivf_42"
it "round-trips through ToJSON/FromJSON" $ do
let original = KnnDeleteModelResponse "test-model" "deleted"
(decode . encode) original `shouldBe` Just original
it "tolerates unknown fields (forward compatibility)" $ do
let Just decoded = decode "{\"model_id\":\"x\",\"result\":\"deleted\",\"future_field\":42}" :: Maybe KnnDeleteModelResponse
knnDeleteModelResponseModelId decoded `shouldBe` "x"
it "requires the model_id field" $ do
let decoded = decode "{\"result\":\"deleted\"}" :: Maybe KnnDeleteModelResponse
decoded `shouldBe` Nothing
it "requires the result field" $ do
let decoded = decode "{\"model_id\":\"x\"}" :: Maybe KnnDeleteModelResponse
decoded `shouldBe` Nothing
it "rejects a top-level array" $ do
let decoded = decode "[]" :: Maybe KnnDeleteModelResponse
decoded `shouldBe` Nothing
it "rejects malformed JSON" $ do
let decoded = decode "{ not json" :: Maybe KnnDeleteModelResponse
decoded `shouldBe` Nothing
describe "deleteKnnModel endpoint shape" $ do
-- Pure checks against the BHRequest shape; no live backend needed.
it "DELETEs /_plugins/_knn/models/{model_id} for the given ModelId" $ do
let req = OS3Requests.deleteKnnModel (ModelId "test-model")
getRawEndpoint (bhRequestEndpoint req)
`shouldBe` ["_plugins", "_knn", "models", "test-model"]
getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []
it "uses the model id as a single opaque path segment" $ do
let req = OS3Requests.deleteKnnModel (ModelId "faiss-ivf-pq_42")
getRawEndpoint (bhRequestEndpoint req)
`shouldBe` ["_plugins", "_knn", "models", "faiss-ivf-pq_42"]
it "uses a distinct id in the path when given a different ModelId" $ do
let req = OS3Requests.deleteKnnModel (ModelId "other-model")
getRawEndpoint (bhRequestEndpoint req)
`shouldBe` ["_plugins", "_knn", "models", "other-model"]
it "uses the DELETE method" $ do
let req = OS3Requests.deleteKnnModel (ModelId "test-model")
bhRequestMethod req `shouldBe` "DELETE"
it "does not attach a body (DELETE semantics)" $ do
let req = OS3Requests.deleteKnnModel (ModelId "test-model")
bhRequestBody req `shouldBe` Nothing
-- The DELETE and GET endpoints share an identical path and differ only in
-- HTTP method, so a copy-paste regression that pointed deleteKnnModel at
-- the GET builder would be invisible without this guard. Mirrors the
-- clearKnnCache/warmupKnnIndex distinctness precedent in Test.KnnClearCacheSpec.
it "is distinct from getKnnModel (same path, different method)" $ do
let modelId = ModelId "test-model"
delReq = OS3Requests.deleteKnnModel modelId
getReq = OS3Requests.getKnnModel modelId
getRawEndpoint (bhRequestEndpoint delReq)
`shouldBe` getRawEndpoint (bhRequestEndpoint getReq)
bhRequestMethod delReq `shouldBe` "DELETE"
bhRequestMethod getReq `shouldBe` "GET"
bhRequestMethod delReq `shouldNotBe` bhRequestMethod getReq
describe "searchKnnModels endpoint shape" $ do
-- Pure checks against the BHRequest shape; no live backend needed.
-- Mirrors the submitOSAsyncSearch body-forwarding tests and the
-- searchModels (ML Commons) path test.
it "POSTs to /_plugins/_knn/models/_search with no query string" $ do
let req = OS3Requests.searchKnnModels (mkSearch (Just (MatchAllQuery Nothing)) Nothing)
bhRequestMethod req `shouldBe` "POST"
getRawEndpoint (bhRequestEndpoint req)
`shouldBe` ["_plugins", "_knn", "models", "_search"]
getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []
it "attaches the encoded Search as the JSON body" $ do
let search = mkSearch (Just (MatchAllQuery Nothing)) Nothing
req = OS3Requests.searchKnnModels search
bhRequestBody req `shouldBe` Just (encode search)
it "tracks the input Search: a different query yields a different body" $ do
let reqA = OS3Requests.searchKnnModels (mkSearch (Just (MatchAllQuery Nothing)) Nothing)
reqB = OS3Requests.searchKnnModels (mkSearch (Just (TermQuery (Term "user" "bob") Nothing)) Nothing)
bhRequestBody reqA `shouldNotBe` bhRequestBody reqB
describe "searchKnnModels SearchResult KnnModel JSON" $ do
it "decodes the standard search envelope into SearchResult KnnModel" $ do
let Just decoded = decode sampleSearchModelsResponse :: Maybe (SearchResult KnnModel)
took decoded `shouldBe` 3
timedOut decoded `shouldBe` False
let modelHits = hits (searchHits decoded)
length modelHits `shouldBe` 2
it "decodes a fully-populated _source as KnnModel" $ do
let Just decoded = decode sampleSearchModelsResponse :: Maybe (SearchResult KnnModel)
firstHit = head (hits (searchHits decoded))
Just model = hitSource firstHit
knnModelModelId model `shouldBe` "test-model"
knnModelState model `shouldBe` Just KnnModelStateCreated
knnModelDimension model `shouldBe` Just 128
knnModelEngine model `shouldBe` Just KnnEngineFaiss
it "decodes a metadata-only _source (model_blob excluded) as KnnModel" $ do
let Just decoded = decode sampleSearchModelsResponse :: Maybe (SearchResult KnnModel)
secondHit = hits (searchHits decoded) !! 1
Just model = hitSource secondHit
knnModelModelId model `shouldBe` "meta-only"
knnModelState model `shouldBe` Just KnnModelStateTraining
knnModelModelBlob model `shouldBe` Nothing
knnModelDimension model `shouldBe` Nothing
it "rejects a response missing the hits envelope" $ do
let decoded = decode "{\"took\":1,\"timed_out\":false,\"_shards\":{}}" :: Maybe (SearchResult KnnModel)
decoded `shouldBe` Nothing