NeuralWarmupSpec.hs

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

module Test.NeuralWarmupSpec (spec) where

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 Neural Search plugin warmup API
-- (@POST /_plugins/_neural/warmup\/{index}@) for OpenSearch 3. These
-- mirror 'Test.OSAsyncSearchSpec's shape tests but assert the
-- @\/_plugins\/_neural\/warmup@ path, the POST method, and the
-- empty-object body. JSON decoding of the shared 'ShardsResult' response
-- is re-verified here on the warmup response shape
-- (@{\"_shards\":{...}}@). The wire shape for the neural warmup API is
-- live-verified against OS 3.7.0 in bead bloodhound-at5. The neural-sparse
-- warmup route is OS 3.x only: OS 1.3.x does not ship the neural-search
-- plugin (bloodhound-ie4j) and OS 2.19.5 does not register the route or
-- the sparse_vector field type (bloodhound-sh7o), so neither OS1 nor OS2
-- ships a builder and no cross-backend parity block is meaningful.
spec :: Spec
spec = describe "Neural Warmup API" $ do
  describe "warmupNeuralIndex endpoint shape (OpenSearch 3)" $ do
    it "POSTs to /_plugins/_neural/warmup/{index} when given an IndexName" $ do
      let req = OS3Requests.warmupNeuralIndex [[qqIndexName|my-index|]]
      bhRequestMethod req `shouldBe` "POST"
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_neural", "warmup", "my-index"]
      getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []

    it "attaches an empty JSON object as the body" $ do
      let req = OS3Requests.warmupNeuralIndex [[qqIndexName|my-index|]]
      bhRequestBody req `shouldBe` Just (encode (object []))

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

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

  -- ---------------------------------------------------------------- --
  -- Multi-index (comma-separated) path rendering: bloodhound-6jy.    --
  -- The Neural 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 "warmupNeuralIndex multi-index path rendering (OpenSearch 3)" $ do
    it "renders [i1,i2,i3] as a single comma-joined path segment" $ do
      let req =
            OS3Requests.warmupNeuralIndex
              [[qqIndexName|index1|], [qqIndexName|index2|], [qqIndexName|index3|]]
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_neural", "warmup", "index1,index2,index3"]

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

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

  describe "ShardsResult response decoding" $ do
    -- Pure unit tests (no OS round-trip): lock the 'FromJSON ShardsResult'
    -- decoder against the wire shape returned by
    -- @POST /_plugins/_neural/warmup\/{index}@. Live OS 3.7.0 confirmed
    -- the response is the @_shards@ envelope
    -- (@{"_shards":{"total":2,"successful":1,"failed":0}}@), matching
    -- @flushIndex@ \/ @refreshIndex@ \/ @clearIndexCache@ (bead
    -- bloodhound-at5).
    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

  -- ---------------------------------------------------------------- --
  -- Live round-trip: gated per-OS with os{n}OnlyIT (the              --
  -- 'Test.KnnWarmupSpec' pattern).                                     --
  --                                                                    --
  -- Wire shape live-verified against OS 3.7.0 (port 9205) in bead     --
  -- bloodhound-at5: the response is                                    --
  -- @{"_shards":{"total":2,"successful":1,"failed":0}}@ — the         --
  -- @_shards@ envelope, decoded into 'ShardsResult'. The endpoint     --
  -- requires an index with @index.sparse: true@ and a @sparse_vector@  --
  -- field mapping.                                                     --
  --                                                                    --
  -- The neural-sparse warmup route is OS 3.x only: OS 1.3.x does not  --
  -- ship the neural-search plugin (bloodhound-ie4j), and OS 2.19.5    --
  -- does not register the warmup route or the @sparse_vector@ field   --
  -- type (bloodhound-sh7o), so the live round-trip can only be        --
  -- exercised on OS 3.                                                 --
  -- ---------------------------------------------------------------- --
  describe "warmupNeuralIndex live round-trip" $ do
    os3It <- runIO os3OnlyIT

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

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