SearchShardsSpec.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}
{-# OPTIONS_GHC -Wno-incomplete-uni-patterns -Wno-x-partial #-}

module Test.SearchShardsSpec (spec) where

import Data.Aeson (Value (..), decode)
import Data.Aeson.KeyMap qualified as KM
import Data.ByteString.Lazy.Char8 qualified as BL8
import Data.Map.Strict qualified as M
import Database.Bloodhound.Common.Requests qualified as Requests
  ( getSearchShards,
    getSearchShardsWith,
  )
import TestsUtils.Common
import TestsUtils.Import

spec :: Spec
spec = do
  describe "SearchShardsResponse JSON" $ do
    it "decodes the canonical ES doc example with one STARTED primary" $
      let bytes =
            BL8.pack
              "{\"nodes\":{\"CBH3C4RiQcK3CxFvPJyFtQ\":{\"name\":\"elasticsearch1\",\"ephemeral_id\":\"DGt8SykXTlWz3GEQtdNfUA\",\"transport_address\":\"172.19.0.6:9300\",\"attributes\":{\"ml.machine_memory\":\"2147483648\"},\"roles\":[\"data\",\"master\"]}},\"indices\":{\"bloodhound-tests-twitter-gettask-opts\":{}},\"shards\":[[{\"state\":\"STARTED\",\"primary\":true,\"node\":\"CBH3C4RiQcK3CxFvPJyFtQ\",\"relocating_node\":null,\"shard\":0,\"index\":\"bloodhound-tests-twitter-gettask-opts\",\"allocation_id\":{\"id\":\"A4BybQlTS_mWcfqMiOR2zg\"}}]]}"
          Just resp = decode bytes
          nodes = searchShardsResponseNodes resp
          Just (firstShardCopies : _) = Just (searchShardsResponseShards resp)
          firstShard = head firstShardCopies
          Just nodeId = searchShardsShardNode firstShard
          Just nodeInfo = M.lookup nodeId nodes
       in do
            M.size nodes `shouldBe` 1
            searchShardsNodeName nodeInfo `shouldBe` Just "elasticsearch1"
            searchShardsNodeTransportAddress nodeInfo `shouldBe` Just "172.19.0.6:9300"
            searchShardsNodeRoles nodeInfo `shouldBe` Just ["data", "master"]
            M.size (searchShardsResponseIndices resp) `shouldBe` 1
            searchShardsShardState firstShard `shouldBe` SearchShardsStateStarted
            searchShardsShardPrimary firstShard `shouldBe` True
            searchShardsShardRelocatingNode firstShard `shouldBe` Nothing
            searchShardsShardAllocationId firstShard
              `shouldBe` Just (M.singleton "id" "A4BybQlTS_mWcfqMiOR2zg")
            unIndexName (searchShardsShardIndex firstShard)
              `shouldBe` "bloodhound-tests-twitter-gettask-opts"

    it "decodes a response with empty nodes/indices/shards (no matching index)" $
      let bytes = BL8.pack "{\"nodes\":{},\"indices\":{},\"shards\":[]}"
          Just resp = decode bytes :: Maybe SearchShardsResponse
       in do
            M.null (searchShardsResponseNodes resp) `shouldBe` True
            M.null (searchShardsResponseIndices resp) `shouldBe` True
            searchShardsResponseShards resp `shouldBe` []

    it "treats missing nodes/indices/shards fields as empty" $
      let bytes = BL8.pack "{}"
          Just resp = decode bytes :: Maybe SearchShardsResponse
       in do
            M.null (searchShardsResponseNodes resp) `shouldBe` True
            M.null (searchShardsResponseIndices resp) `shouldBe` True
            searchShardsResponseShards resp `shouldBe` []

    it "decodes an unassigned shard (state=UNASSIGNED, node=null)" $
      let bytes =
            BL8.pack
              "{\"nodes\":{},\"indices\":{\"foo\":{}},\"shards\":[[{\"state\":\"UNASSIGNED\",\"primary\":false,\"node\":null,\"relocating_node\":null,\"shard\":1,\"index\":\"foo\"}]]}"
          Just resp = decode bytes :: Maybe SearchShardsResponse
          Just (copy : _) = Just (searchShardsResponseShards resp)
          shard = head copy
       in do
            searchShardsShardState shard `shouldBe` SearchShardsStateUnassigned
            searchShardsShardNode shard `shouldBe` Nothing
            searchShardsShardPrimary shard `shouldBe` False

    it "preserves unknown sibling fields on a shard copy via searchShardsShardOther" $
      let bytes =
            BL8.pack
              "{\"nodes\":{},\"indices\":{\"foo\":{}},\"shards\":[[{\"state\":\"STARTED\",\"primary\":true,\"node\":\"n1\",\"relocating_node\":null,\"shard\":0,\"index\":\"foo\",\"allocation_id\":{\"id\":\"x\"},\"future_field\":42}]]}"
          Just resp = decode bytes :: Maybe SearchShardsResponse
          Just (copy : _) = Just (searchShardsResponseShards resp)
          shard = head copy
          Just (Object other) = searchShardsShardOther shard
       in KM.size other `shouldSatisfy` (>= 1)

    it "preserves unknown sibling fields on a node via searchShardsNodeOther" $
      let bytes =
            BL8.pack
              "{\"nodes\":{\"n1\":{\"name\":\"n\",\"future_attr\":\"x\"}},\"indices\":{},\"shards\":[]}"
          Just resp = decode bytes :: Maybe SearchShardsResponse
          Just nodeInfo = M.lookup "n1" (searchShardsResponseNodes resp)
          Just (Object other) = searchShardsNodeOther nodeInfo
       in do
            searchShardsNodeName nodeInfo `shouldBe` Just "n"
            KM.size other `shouldSatisfy` (>= 1)

    it "round-trips an unknown shard state through SearchShardsStateOther" $
      let bytes =
            BL8.pack
              "{\"nodes\":{},\"indices\":{\"foo\":{}},\"shards\":[[{\"state\":\"WEIRD\",\"primary\":true,\"node\":\"n1\",\"relocating_node\":null,\"shard\":0,\"index\":\"foo\"}]]}"
          Just resp = decode bytes :: Maybe SearchShardsResponse
          Just (copy : _) = Just (searchShardsResponseShards resp)
          shard = head copy
       in searchShardsShardState shard `shouldBe` SearchShardsStateOther "WEIRD"

  describe "SearchShardsOptions params" $ do
    it "renders nothing for defaultSearchShardsOptions" $
      searchShardsOptionsParams defaultSearchShardsOptions `shouldBe` []

    it "renders local as a boolean string" $
      let opts = defaultSearchShardsOptions {sshoLocal = Just True}
       in lookup "local" (searchShardsOptionsParams opts) `shouldBe` Just (Just "true")

    it "renders allow_no_indices as a boolean string" $
      let opts = defaultSearchShardsOptions {sshoAllowNoIndices = Just False}
       in lookup "allow_no_indices" (searchShardsOptionsParams opts)
            `shouldBe` Just (Just "false")

    it "renders expand_wildcards comma-joined" $
      let opts =
            defaultSearchShardsOptions
              { sshoExpandWildcards = Just [ExpandWildcardsOpen, ExpandWildcardsClosed]
              }
       in lookup "expand_wildcards" (searchShardsOptionsParams opts)
            `shouldBe` Just (Just "open,closed")

    it "renders preference verbatim" $
      let opts = defaultSearchShardsOptions {sshoPreference = Just "_primary"}
       in lookup "preference" (searchShardsOptionsParams opts) `shouldBe` Just (Just "_primary")

    it "renders routing comma-joined" $
      let opts = defaultSearchShardsOptions {sshoRouting = Just ["a", "b"]}
       in lookup "routing" (searchShardsOptionsParams opts) `shouldBe` Just (Just "a,b")

    it "renders master_timeout via timeUnitsSuffix" $
      let opts = defaultSearchShardsOptions {sshoMasterTimeout = Just (TimeUnitSeconds, 30)}
       in lookup "master_timeout" (searchShardsOptionsParams opts) `shouldBe` Just (Just "30s")

    it "omits master_timeout by default" $
      lookup "master_timeout" (searchShardsOptionsParams defaultSearchShardsOptions)
        `shouldBe` Nothing

  describe "Search Shards endpoint shape" $ do
    -- Pure checks against the BHRequest shape; no live backend needed.
    it "POSTs /_search_shards with no body" $ do
      let req = Requests.getSearchShards Nothing
      bhRequestMethod req `shouldBe` "POST"
      getRawEndpoint (bhRequestEndpoint req) `shouldBe` ["_search_shards"]
      bhRequestBody req `shouldBe` Nothing
      getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []

    it "treats Just [] the same as Nothing (single _search_shards segment)" $ do
      let reqNo = Requests.getSearchShards Nothing
          reqEmpty = Requests.getSearchShards (Just [])
      getRawEndpoint (bhRequestEndpoint reqNo)
        `shouldBe` getRawEndpoint (bhRequestEndpoint reqEmpty)

    it "comma-joins a non-empty index list into the leading segment" $ do
      let req = Requests.getSearchShards (Just [[qqIndexName|twitter-1|], [qqIndexName|twitter-2|]])
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["twitter-1,twitter-2", "_search_shards"]

    it "honours SearchShardsOptions in the query string" $ do
      let req =
            Requests.getSearchShardsWith
              (Just [[qqIndexName|twitter-1|]])
              defaultSearchShardsOptions {sshoPreference = Just "_primary"}
      getRawEndpointQueries (bhRequestEndpoint req)
        `shouldBe` [("preference", Just "_primary")]

  describe "Search Shards endpoint" $ do
    it "returns at least one assigned primary shard for the test index" $
      withTestEnv $ do
        _ <- insertData
        resp <- performBHRequest $ getSearchShards (Just [testIndex])
        liftIO $ do
          let shards = searchShardsResponseShards resp
          shards `shouldSatisfy` (not . null)
          let allShards = concat shards
          allShards `shouldSatisfy` any searchShardsShardPrimary
          let testIndexShards =
                filter
                  ((==) (unIndexName testIndex) . unIndexName . searchShardsShardIndex)
                  allShards
          testIndexShards `shouldSatisfy` (not . null)

    it "exposes a node entry for every assigned shard's node id" $
      withTestEnv $ do
        _ <- insertData
        resp <- performBHRequest $ getSearchShards (Just [testIndex])
        liftIO $ do
          let nodes = searchShardsResponseNodes resp
              assignedNodeIds =
                [ nid
                | copy <- concat (searchShardsResponseShards resp),
                  Just nid <- [searchShardsShardNode copy]
                ]
              uniqueNodeIds = uniq assignedNodeIds
          uniqueNodeIds `shouldSatisfy` (not . null)
          mapM_ (\nid -> M.member nid nodes `shouldBe` True) uniqueNodeIds

    it "passes preference and routing via SearchShardsOptions without error" $
      withTestEnv $ do
        _ <- insertData
        resp <-
          performBHRequest $
            getSearchShardsWith
              (Just [testIndex])
              defaultSearchShardsOptions
                { sshoPreference = Just "_local",
                  sshoRouting = Just ["user_1"]
                }
        liftIO $
          searchShardsResponseShards resp `shouldSatisfy` (not . null)

-- | Local helper because we cannot use the lens from the module under
-- test in two different roles easily without an explicit accessor.
uniq :: (Eq a) => [a] -> [a]
uniq = go []
  where
    go _ [] = []
    go seen (x : xs)
      | x `elem` seen = go seen xs
      | otherwise = x : go (x : seen) xs