packages feed

bloodhound-1.0.0.0: tests/Test/KnnWarmupSpec.hs

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

module Test.KnnWarmupSpec (spec) where

import Database.Bloodhound.OpenSearch1.Client qualified as OS1.Client
import Database.Bloodhound.OpenSearch1.Requests qualified as OS1Requests
import Database.Bloodhound.OpenSearch2.Client qualified as OS2.Client
import Database.Bloodhound.OpenSearch2.Requests qualified as OS2Requests
import Database.Bloodhound.OpenSearch3.Client qualified as OS3.Client
import Database.Bloodhound.OpenSearch3.Requests qualified as OS3Requests
import TestsUtils.Common
import TestsUtils.Import
import Prelude

-- | Pure endpoint-shape tests for the k-NN plugin Warmup API
-- (@GET /_plugins/_knn/warmup\/{index}@) across OpenSearch 1, 2, and 3.
-- These mirror 'Test.NeuralWarmupSpec's shape tests but assert the
-- @\/_plugins\/_knn\/warmup@ path, the GET method, the absent body,
-- and cross-backend parity. JSON decoding of the shared 'ShardsResult'
-- response is re-verified here on the warmup response shape
-- (@{\"_shards\":{\"total\":...,\"successful\":...,\"failed\":...}}@,
-- live-verified against OS 2.19.5 \/ 3.7.0) so a regression in the
-- pre-existing instance is caught in this spec too.
spec :: Spec
spec = describe "k-NN Warmup API" $ do
  describe "warmupKnnIndex endpoint shape (OpenSearch 3)" $ do
    it "GETs /_plugins/_knn/warmup/{index} when given an IndexName" $ do
      let req = OS3Requests.warmupKnnIndex [[qqIndexName|my-index|]]
      bhRequestMethod req `shouldBe` "GET"
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "warmup", "my-index"]
      getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []

    it "carries no request body" $ do
      let req = OS3Requests.warmupKnnIndex [[qqIndexName|my-index|]]
      bhRequestBody req `shouldBe` Nothing

    it "uses a distinct index in the path when given a different IndexName" $ do
      let req = OS3Requests.warmupKnnIndex [[qqIndexName|other-index|]]
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "warmup", "other-index"]

    it "keeps the index name as a single opaque path segment" $ do
      let req = OS3Requests.warmupKnnIndex [[qqIndexName|my-index_42|]]
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "warmup", "my-index_42"]

  describe "warmupKnnIndex endpoint shape (OpenSearch 2)" $ do
    it "GETs /_plugins/_knn/warmup/{index} when given an IndexName" $ do
      let req = OS2Requests.warmupKnnIndex [[qqIndexName|my-index|]]
      bhRequestMethod req `shouldBe` "GET"
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "warmup", "my-index"]
      getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []

    it "carries no request body" $ do
      let req = OS2Requests.warmupKnnIndex [[qqIndexName|my-index|]]
      bhRequestBody req `shouldBe` Nothing

    it "uses a distinct index in the path when given a different IndexName" $ do
      let req = OS2Requests.warmupKnnIndex [[qqIndexName|other-index|]]
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "warmup", "other-index"]

  describe "warmupKnnIndex endpoint shape (OpenSearch 1)" $ do
    it "GETs /_plugins/_knn/warmup/{index} when given an IndexName" $ do
      let req = OS1Requests.warmupKnnIndex [[qqIndexName|my-index|]]
      bhRequestMethod req `shouldBe` "GET"
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "warmup", "my-index"]
      getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []

    it "carries no request body" $ do
      let req = OS1Requests.warmupKnnIndex [[qqIndexName|my-index|]]
      bhRequestBody req `shouldBe` Nothing

    it "uses a distinct index in the path when given a different IndexName" $ do
      let req = OS1Requests.warmupKnnIndex [[qqIndexName|other-index|]]
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "warmup", "other-index"]

  -- ---------------------------------------------------------------- --
  -- Multi-index (comma-separated) path rendering: bloodhound-6jy.    --
  -- The k-NN Warmup API documents comma-separated lists and index    --
  -- patterns in the {index} path segment; the [IndexName] argument   --
  -- is rendered as a single comma-joined segment, matching the       --
  -- OpenSearch multi-target syntax.                                  --
  -- ---------------------------------------------------------------- --
  describe "warmupKnnIndex multi-index path rendering (OpenSearch 3)" $ do
    it "renders [i1,i2,i3] as a single comma-joined path segment" $ do
      let req =
            OS3Requests.warmupKnnIndex
              [[qqIndexName|index1|], [qqIndexName|index2|], [qqIndexName|index3|]]
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "warmup", "index1,index2,index3"]

    it "preserves index name order in the comma-joined segment" $ do
      let req =
            OS3Requests.warmupKnnIndex
              [[qqIndexName|zebra|], [qqIndexName|alpha|], [qqIndexName|mike|]]
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "warmup", "zebra,alpha,mike"]

    it "renders an empty segment for an empty index list (server will error)" $ do
      let req = OS3Requests.warmupKnnIndex []
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "warmup", ""]

  describe "warmupKnnIndex multi-index path rendering (OpenSearch 2)" $ do
    it "renders [i1,i2] as a single comma-joined path segment" $ do
      let req =
            OS2Requests.warmupKnnIndex [[qqIndexName|index1|], [qqIndexName|index2|]]
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "warmup", "index1,index2"]

  describe "warmupKnnIndex multi-index path rendering (OpenSearch 1)" $ do
    it "renders [i1,i2] as a single comma-joined path segment" $ do
      let req =
            OS1Requests.warmupKnnIndex [[qqIndexName|index1|], [qqIndexName|index2|]]
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "warmup", "index1,index2"]

  describe "warmupKnnIndex cross-backend parity" $ do
    -- All three backends must emit the same path, GET method, absent
    -- body, and empty query string for the same 'IndexName'.
    let idx = [qqIndexName|shared-index|]

    it "OS1, OS2, OS3 produce the same path segments" $ do
      let path1 = getRawEndpoint (bhRequestEndpoint (OS1Requests.warmupKnnIndex [idx]))
          path2 = getRawEndpoint (bhRequestEndpoint (OS2Requests.warmupKnnIndex [idx]))
          path3 = getRawEndpoint (bhRequestEndpoint (OS3Requests.warmupKnnIndex [idx]))
      path1 `shouldBe` path2
      path2 `shouldBe` path3

    it "OS1, OS2, OS3 all use GET" $ do
      let m1 = bhRequestMethod (OS1Requests.warmupKnnIndex [idx])
          m2 = bhRequestMethod (OS2Requests.warmupKnnIndex [idx])
          m3 = bhRequestMethod (OS3Requests.warmupKnnIndex [idx])
      m1 `shouldBe` "GET"
      m2 `shouldBe` "GET"
      m3 `shouldBe` "GET"

    it "OS1, OS2, OS3 all carry no query string" $ do
      let q1 = getRawEndpointQueries (bhRequestEndpoint (OS1Requests.warmupKnnIndex [idx]))
          q2 = getRawEndpointQueries (bhRequestEndpoint (OS2Requests.warmupKnnIndex [idx]))
          q3 = getRawEndpointQueries (bhRequestEndpoint (OS3Requests.warmupKnnIndex [idx]))
      q1 `shouldBe` []
      q2 `shouldBe` []
      q3 `shouldBe` []

    it "OS1, OS2, OS3 all carry no body" $ do
      let b1 = bhRequestBody (OS1Requests.warmupKnnIndex [idx])
          b2 = bhRequestBody (OS2Requests.warmupKnnIndex [idx])
          b3 = bhRequestBody (OS3Requests.warmupKnnIndex [idx])
      b1 `shouldBe` Nothing
      b2 `shouldBe` Nothing
      b3 `shouldBe` Nothing

  describe "ShardsResult response decoding" $ do
    -- Pure unit tests (no OS round-trip): lock the 'FromJSON ShardsResult'
    -- decoder against the actual wire shape returned by
    -- @GET /_plugins/_knn/warmup\/{index}@. Live OS 2.19.5 \/ 3.7.0
    -- confirmed the response is the @_shards@ envelope (matching
    -- @flushIndex@ \/ @refreshIndex@ \/ @clearIndexCache@), not the
    -- @{\"acknowledged\":true}@ envelope the previous 'Acknowledged'
    -- return type assumed.
    it "decodes {\"_shards\":{\"total\":2,\"successful\":1,\"skipped\":0,\"failed\":0}} into ShardsResult" $ do
      let payload = "{\"_shards\":{\"total\":2,\"successful\":1,\"skipped\":0,\"failed\":0}}"
          Just decoded = decode payload :: Maybe ShardsResult
      decoded `shouldBe` ShardsResult (ShardResult 2 1 0 0)

    it "rejects a payload missing the _shards envelope" $ do
      let payload = "{\"acknowledged\":true}"
          decoded = decode payload :: Maybe ShardsResult
      decoded `shouldBe` Nothing

    it "round-trips a captured live warmup response through the decoder" $ do
      let original = "{\"_shards\":{\"total\":2,\"successful\":1,\"failed\":0}}"
          Just decoded = decode original :: Maybe ShardsResult
      shardTotal (srShards decoded) `shouldBe` 2

  -- ---------------------------------------------------------------- --
  -- Live round-trip: gated per-OS with os{n}OnlyIT (the               --
  -- 'Test.MLModelSpec' pattern). The k-NN Warmup API exists on        --
  -- OS 1, 2, and 3, so all three backends are exercised.              --
  --                                                                    --
  -- Wire shape live-verified against OS 2.19.5 (port 9204) and OS     --
  -- 3.7.0 (port 9205) in bead bloodhound-vvj: the response is         --
  -- @{"_shards":{"total":2,"successful":1,"failed":0}}@ — the         --
  -- @_shards@ envelope, decoded into 'ShardsResult'. The HTTP         --
  -- method (POST vs GET) was previously wrong and fixed in            --
  -- bloodhound-04f.6.2.9.                                             --
  -- ---------------------------------------------------------------- --
  describe "warmupKnnIndex live round-trip" $ do
    os1It <- runIO os1OnlyIT
    os2It <- runIO os2OnlyIT
    os3It <- runIO os3OnlyIT

    let assertShards :: ShardsResult -> IO ()
        assertShards (ShardsResult sr) = do
          shardTotal sr `shouldSatisfy` (>= 1)
          shardsFailed sr `shouldBe` 0

    os1It "OS1: returns a ShardsResult with at least one shard and no failures" $
      withTestEnv $ do
        _ <- tryEsError resetOsKnnIndex
        result <- OS1.Client.warmupKnnIndex osKnnTestIndex
        liftIO $ assertShards result

    os2It "OS2: returns a ShardsResult with at least one shard and no failures" $
      withTestEnv $ do
        _ <- tryEsError resetOsKnnIndex
        result <- OS2.Client.warmupKnnIndex osKnnTestIndex
        liftIO $ assertShards result

    os3It "OS3: returns a ShardsResult with at least one shard and no failures" $
      withTestEnv $ do
        _ <- tryEsError resetOsKnnIndex
        result <- OS3.Client.warmupKnnIndex osKnnTestIndex
        liftIO $ assertShards result