KnnOs2Spec.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}

module Test.KnnOs2Spec (spec) where

import Data.ByteString.Lazy.Char8 qualified as LBS
import Database.Bloodhound.OpenSearch2.Requests qualified as OS2Requests
import Database.Bloodhound.OpenSearch2.Types
import Network.HTTP.Types.Method qualified as NHTM
import TestsUtils.Import
import Prelude

-- | Sample identifiers for endpoint-shape assertions. The k-NN plugin uses
-- opaque 22-character node IDs; any 'Text' round-trips, so these are short
-- stand-ins purely for URL-segment checks.
testNodeId :: KnnNodeId
testNodeId = KnnNodeId "nodeIdOne"

testStatName :: KnnStatName
testStatName = KnnStatName "hit_count"

testModelId :: ModelId
testModelId = ModelId "test-model"

-- | A minimal 'KnnTrainingRequest' (FAISS IVF) — only the shape matters for
-- endpoint tests, never the field values.
sampleTrainingRequest :: KnnTrainingRequest
sampleTrainingRequest =
  KnnTrainingRequest
    { knnTrainingRequestIndex = [qqIndexName|train-idx|],
      knnTrainingRequestField = FieldName "vector",
      knnTrainingRequestDimension = 4,
      knnTrainingRequestMethod =
        KnnTrainingMethod
          { knnTrainingMethodName = "ivf",
            knnTrainingMethodEngine = KnnEngineFaiss,
            knnTrainingMethodSpaceType = Just "l2",
            knnTrainingMethodParameters = Nothing
          },
      knnTrainingRequestSpaceType = Nothing,
      knnTrainingRequestMaxTrainingVectorCount = Nothing,
      knnTrainingRequestSearchSize = Nothing,
      knnTrainingRequestDescription = Nothing
    }

spec :: Spec
spec = describe "k-NN plugin endpoint shapes - OpenSearch 2" $ do
  describe "getKnnStats" $ do
    it "GETs /_plugins/_knn/stats with no node/stat filters" $ do
      let req = OS2Requests.getKnnStats Nothing Nothing
      bhRequestMethod req `shouldBe` NHTM.methodGet
      getRawEndpoint (bhRequestEndpoint req) `shouldBe` ["_plugins", "_knn", "stats"]
      bhRequestBody req `shouldSatisfy` isNothing

    it "inserts the node id and stat name into the path when given" $ do
      let req = OS2Requests.getKnnStats (Just testNodeId) (Just testStatName)
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", unKnnNodeId testNodeId, "stats", unKnnStatName testStatName]

    it "emits node id only when no stat is requested" $ do
      let req = OS2Requests.getKnnStats (Just testNodeId) Nothing
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", unKnnNodeId testNodeId, "stats"]

  describe "getKnnModel" $ do
    it "GETs /_plugins/_knn/models/{model_id} with no body" $ do
      let req = OS2Requests.getKnnModel testModelId
      bhRequestMethod req `shouldBe` NHTM.methodGet
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "models", unModelId testModelId]
      bhRequestBody req `shouldSatisfy` isNothing

  describe "trainKnnModel" $ do
    it "POSTs the training body to /_plugins/_knn/models/{model_id}/_train" $ do
      let req = OS2Requests.trainKnnModel testModelId sampleTrainingRequest
      bhRequestMethod req `shouldBe` NHTM.methodPost
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "models", unModelId testModelId, "_train"]
      bhRequestBody req `shouldSatisfy` isJust
      getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []

  describe "trainKnnModelWith" $ do
    it "adds ?preference={node_id} when a node is supplied" $ do
      let req = OS2Requests.trainKnnModelWith (Just testNodeId) testModelId sampleTrainingRequest
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "models", unModelId testModelId, "_train"]
      getRawEndpointQueries (bhRequestEndpoint req)
        `shouldBe` [("preference", Just (unKnnNodeId testNodeId))]

    it "omits ?preference when no node is supplied (regression)" $ do
      let req = OS2Requests.trainKnnModelWith Nothing testModelId sampleTrainingRequest
      getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []

  describe "deleteKnnModel" $ do
    it "DELETEs /_plugins/_knn/models/{model_id} (distinct method from getKnnModel)" $ do
      let req = OS2Requests.deleteKnnModel testModelId
      bhRequestMethod req `shouldBe` NHTM.methodDelete
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "models", unModelId testModelId]
      bhRequestBody req `shouldSatisfy` isNothing

  describe "searchKnnModels" $ do
    it "POSTs the search body to /_plugins/_knn/models/_search" $ do
      let req = OS2Requests.searchKnnModels (mkSearch Nothing Nothing)
      bhRequestMethod req `shouldBe` NHTM.methodPost
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "models", "_search"]
      bhRequestBody req `shouldSatisfy` isJust

    it "serializes the Search DSL as the request body" $ do
      let req = OS2Requests.searchKnnModels (mkSearch Nothing Nothing)
      let body = bhRequestBody req
      body `shouldSatisfy` isJust
      -- An empty search encodes as an object, never an array or null.
      case body of
        Just bs -> LBS.head bs `shouldBe` '{'
        Nothing -> expectationFailure "expected a request body"