MultiSearchSpec.hs

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

module Test.MultiSearchSpec where

import Data.Aeson.Key qualified as K
import Data.Aeson.KeyMap qualified as KM
import Data.ByteString.Lazy.Char8 qualified as LC
import Data.Either (rights)
import Data.List (sort)
import Database.Bloodhound.Common.Types
import TestsUtils.Common
import TestsUtils.Import

-- | Smart constructor that mirrors the pre-bloodhound-04f.7.3 2-field form
-- of 'MultiSearchItem'. Keeps the integration-test fixtures readable while
-- the record itself gains per-header fields.
mkItem :: Maybe IndexName -> Search -> MultiSearchItem
mkItem idx s =
  MultiSearchItem
    { multiSearchItemIndex = idx,
      multiSearchItemRouting = Nothing,
      multiSearchItemSearchType = Nothing,
      multiSearchItemPreference = Nothing,
      multiSearchItemAllowPartialSearchResults = Nothing,
      multiSearchItemSearch = s
    }

-- | Decode a ToJSON-encoded value back to a sorted list of key names so the
-- tests can assert on the header line as a set rather than relying on
-- aeson's key ordering.
headerKeys :: LC.ByteString -> [Text]
headerKeys bs = case decode bs :: Maybe Object of
  Just o -> sort (map K.toText (KM.keys o))
  Nothing -> error "headerKeys: decode failed"

spec :: Spec
spec = do
  describe "MultiSearchItem ToJSON header rendering" $ do
    let bareSearch = mkSearch (Just (MatchAllQuery Nothing)) Nothing
        logs = [qqIndexName|logs|]

    it "renders only index when no per-header fields are set" $ do
      let item = mkItem (Just logs) bareSearch
      encode item `shouldBe` "{\"index\":\"logs\"}"

    it "renders an empty object when no fields (including index) are set" $ do
      let item = mkItem Nothing bareSearch
      encode item `shouldBe` "{}"

    it "renders routing alongside index" $ do
      let item = (mkItem (Just logs) bareSearch) {multiSearchItemRouting = Just "r1"}
      headerKeys (encode item) `shouldBe` ["index", "routing"]

    it "renders search_type using the SearchType Text encoding" $ do
      let item = (mkItem (Just logs) bareSearch) {multiSearchItemSearchType = Just SearchTypeDfsQueryThenFetch}
      encode item `shouldBe` "{\"index\":\"logs\",\"search_type\":\"dfs_query_then_fetch\"}"

    it "renders preference and allow_partial_search_results" $ do
      let item =
            (mkItem (Just logs) bareSearch)
              { multiSearchItemPreference = Just "_local",
                multiSearchItemAllowPartialSearchResults = Just False
              }
      headerKeys (encode item)
        `shouldBe` ["allow_partial_search_results", "index", "preference"]

    it "renders every per-header field simultaneously" $ do
      let item =
            (mkItem (Just logs) bareSearch)
              { multiSearchItemRouting = Just "r1",
                multiSearchItemSearchType = Just SearchTypeDfsQueryThenFetch,
                multiSearchItemPreference = Just "_local",
                multiSearchItemAllowPartialSearchResults = Just True
              }
      -- aeson 2.x renders object keys in alphabetical order; the
      -- assertion pins the full wire shape so regressions in either
      -- the key set or the value encoding surface here.
      encode item
        `shouldBe` "{\"allow_partial_search_results\":true,\"index\":\"logs\",\"preference\":\"_local\",\"routing\":\"r1\",\"search_type\":\"dfs_query_then_fetch\"}"

  describe "SearchOptions max_concurrent_searches (msearch URI param)" $ do
    it "renders max_concurrent_searches as a decimal string" $ do
      let opts = defaultSearchOptions {soMaxConcurrentSearches = Just 8}
      searchOptionsParams opts SearchTypeQueryThenFetch
        `shouldBe` [("max_concurrent_searches", Just "8")]

  describe "MultiSearchResponse per-item error decoding" $ do
    -- No server round-trip: the decoder is exercised directly so the
    -- assertions are stable across ES/OS versions and don't depend on
    -- a live cluster. The wire shapes mirror what ES returns for a
    -- per-item 404 (missing index) and a normal success sub-search.
    it "decodes a per-item error as Left EsError" $ do
      let perItemError =
            LC.pack
              "{\"responses\":[{\"error\":\"index_not_found_exception\",\
              \\"status\":404}]}"
      let Just (resp :: MultiSearchResponse Tweet) = decode perItemError
      multiSearchResponses resp
        `shouldBe` [Left (EsError (Just 404) "index_not_found_exception")]

    it "decodes the modern ES error object form (error.reason) as Left EsError" $ do
      -- ES 7+ returns per-item errors as objects carrying
      -- {"type":..., "reason":...}. EsError's FromJSON p2 alternative
      -- extracts .reason as the message. See BHRequest.hs EsError parser.
      let perItemErrorObject =
            LC.pack
              "{\"responses\":[{\"error\":{\"type\":\"index_not_found_exception\",\
              \\"reason\":\"no such index [foo]\",\
              \\"index\":\"foo\"},\"status\":404}]}"
      let Just (resp :: MultiSearchResponse Tweet) = decode perItemErrorObject
      multiSearchResponses resp
        `shouldBe` [Left (EsError (Just 404) "no such index [foo]")]

    it "preserves order when mixing successes and errors" $ do
      let mixed =
            LC.pack
              "{\"responses\":[\
              \{\"took\":1,\"timed_out\":false,\"_shards\":{},\"hits\":{\"total\":{\"value\":0,\"relation\":\"eq\"},\"max_score\":null,\"hits\":[]}},\
              \{\"error\":\"index_not_found_exception\",\"status\":404}\
              \]}"
      let Just (resp :: MultiSearchResponse Tweet) = decode mixed
      length (multiSearchResponses resp) `shouldBe` 2
      -- Order preserved: first item is a successful SearchResult, the
      -- second is a per-item EsError carrying the 404 status.
      length (rights (multiSearchResponses resp)) `shouldBe` 1
      length [() | Left _ <- multiSearchResponses resp] `shouldBe` 1
      case multiSearchResponses resp of
        [Right _, Left _] -> pure ()
        _ -> expectationFailure "expected [Right _, Left _] in order"

    it "still decodes an all-success payload as all Right" $ do
      let allOk =
            LC.pack
              "{\"responses\":[\
              \{\"took\":1,\"timed_out\":false,\"_shards\":{},\"hits\":{\"total\":{\"value\":0,\"relation\":\"eq\"},\"max_score\":null,\"hits\":[]}},\
              \{\"took\":2,\"timed_out\":false,\"_shards\":{},\"hits\":{\"total\":{\"value\":0,\"relation\":\"eq\"},\"max_score\":null,\"hits\":[]}}\
              \]}"
      let Just (resp :: MultiSearchResponse Tweet) = decode allOk
      multiSearchResponses resp `shouldSatisfy` all (either (const False) (const True))
      length (multiSearchResponses resp) `shouldBe` 2

  describe "multi-search API" $ do
    it "executes multiple searches in a single request" $
      withTestEnv $ do
        _ <- insertData
        _ <- insertOther
        let query1 = TermQuery (Term "user" "bitemyapp") Nothing
        let search1 = mkSearch (Just query1) Nothing
        let query2 = TermQuery (Term "user" "notmyapp") Nothing
        let search2 = mkSearch (Just query2) Nothing
        let items = mkItem Nothing search1 :| [mkItem Nothing search2]
        result <- tryPerformBHRequest $ multiSearch @Tweet items
        case result of
          Left _ -> liftIO $ expectationFailure "Multi-search failed"
          Right resp -> do
            let responses = multiSearchResponses resp
            liftIO $ responses `shouldSatisfy` all (either (const False) (const True))
            liftIO $ length responses `shouldBe` 2

    it "executes multiple searches on a specific index" $
      withTestEnv $ do
        _ <- insertData
        let query1 = TermQuery (Term "user" "bitemyapp") Nothing
        let search1 = mkSearch (Just query1) Nothing
        let query2 = MatchAllQuery Nothing
        let search2 = mkSearch (Just query2) Nothing
        let searches = search1 :| [search2]
        result <- tryPerformBHRequest $ multiSearchByIndex @Tweet testIndex searches
        case result of
          Left _ -> liftIO $ expectationFailure "Multi-search by index failed"
          Right resp -> do
            let responses = multiSearchResponses resp
            liftIO $ responses `shouldSatisfy` all (either (const False) (const True))
            liftIO $ length responses `shouldBe` 2
            let firstResult = head (rights responses)
            let hitsCount = maybe 0 value (hitsTotal (searchHits firstResult))
            liftIO $ hitsCount `shouldSatisfy` (> 0)

    it "handles searches on different indices" $
      withTestEnv $ do
        _ <- insertData
        _ <- insertOther
        let query1 = TermQuery (Term "user" "bitemyapp") Nothing
        let search1 = mkSearch (Just query1) Nothing
        let query2 = TermQuery (Term "user" "notmyapp") Nothing
        let search2 = mkSearch (Just query2) Nothing
        let items =
              mkItem (Just testIndex) search1
                :| [ mkItem (Just testIndex) search2
                   ]
        result <- tryPerformBHRequest $ multiSearch @Tweet items
        case result of
          Left _ -> liftIO $ expectationFailure "Multi-search failed"
          Right resp -> do
            let responses = multiSearchResponses resp
            liftIO $ responses `shouldSatisfy` all (either (const False) (const True))
            liftIO $ length responses `shouldBe` 2

    it "multiSearchWith forwards URI params (max_concurrent_searches, typed_keys)" $
      withTestEnv $ do
        _ <- insertData
        _ <- insertOther
        let query1 = TermQuery (Term "user" "bitemyapp") Nothing
        let search1 = mkSearch (Just query1) Nothing
        let query2 = TermQuery (Term "user" "notmyapp") Nothing
        let search2 = mkSearch (Just query2) Nothing
        let items = mkItem Nothing search1 :| [mkItem Nothing search2]
            opts =
              defaultSearchOptions
                { soMaxConcurrentSearches = Just 1,
                  soTypedKeys = Just True
                }
        result <- tryPerformBHRequest $ multiSearchWith @Tweet opts items
        case result of
          Left _ -> liftIO $ expectationFailure "Multi-search with options failed"
          Right resp -> do
            let responses = multiSearchResponses resp
            liftIO $ responses `shouldSatisfy` all (either (const False) (const True))
            liftIO $ length responses `shouldBe` 2

    it "per-header fields land on the header line and are accepted by ES" $
      withTestEnv $ do
        _ <- insertData
        let query = TermQuery (Term "user" "bitemyapp") Nothing
        let search = mkSearch (Just query) Nothing
        let item =
              (mkItem (Just testIndex) search)
                { multiSearchItemPreference = Just "_local",
                  multiSearchItemSearchType = Just SearchTypeQueryThenFetch,
                  multiSearchItemAllowPartialSearchResults = Just True
                }
        result <- tryPerformBHRequest $ multiSearch @Tweet (item :| [])
        case result of
          Left _ -> liftIO $ expectationFailure "Multi-search with header fields failed"
          Right resp -> do
            let responses = multiSearchResponses resp
            liftIO $ responses `shouldSatisfy` all (either (const False) (const True))
            liftIO $ length responses `shouldBe` 1