RerouteSpec.hs

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

module Test.RerouteSpec where

import Data.Aeson qualified as Aeson
import Data.List qualified as L
import Data.List.NonEmpty qualified as NE
import Data.Text qualified as T
import Database.Bloodhound.Common.Client qualified as Client
import Database.Bloodhound.Common.Requests qualified as Common
import TestsUtils.Common
import TestsUtils.Import

-- | Encode a 'ToJSON' value and decode it straight back as an untyped
-- 'Value'. Object key order is normalised away by the round-trip, so the
-- resulting 'Value' compares equal regardless of the encoder's key order.
encodedAs :: (ToJSON a) => a -> Maybe Value
encodedAs = Aeson.decode . Aeson.encode

spec :: Spec
spec = do
  describe "RerouteCommand ToJSON" $ do
    let idx = [qqIndexName|twitter|]
        shard = ShardId 0
        n1 = NodeName "node-1"
        n2 = NodeName "node-2"

    it "renders move as {\"move\":{index,shard,from_node,to_node}}" $ do
      let cmd = RerouteMove idx shard n1 n2
      encodedAs cmd
        `shouldBe` Just
          ( object
              [ "move"
                  .= object
                    [ "index" .= ("twitter" :: Text),
                      "shard" .= (0 :: Int),
                      "from_node" .= ("node-1" :: Text),
                      "to_node" .= ("node-2" :: Text)
                    ]
              ]
          )

    it "renders cancel without allow_primary when it is Nothing" $ do
      let cmd = RerouteCancel idx shard n1 Nothing
      encodedAs cmd
        `shouldBe` Just
          ( object
              [ "cancel"
                  .= object
                    [ "index" .= ("twitter" :: Text),
                      "shard" .= (0 :: Int),
                      "node" .= ("node-1" :: Text)
                    ]
              ]
          )

    it "renders cancel with allow_primary when it is Just" $ do
      let cmd = RerouteCancel idx shard n1 (Just True)
      encodedAs cmd
        `shouldBe` Just
          ( object
              [ "cancel"
                  .= object
                    [ "index" .= ("twitter" :: Text),
                      "shard" .= (0 :: Int),
                      "node" .= ("node-1" :: Text),
                      "allow_primary" .= True
                    ]
              ]
          )

    it "renders allocate_replica as {\"allocate_replica\":{index,shard,node}}" $ do
      let cmd = RerouteAllocateReplica idx shard n1
      encodedAs cmd
        `shouldBe` Just
          ( object
              [ "allocate_replica"
                  .= object
                    [ "index" .= ("twitter" :: Text),
                      "shard" .= (0 :: Int),
                      "node" .= ("node-1" :: Text)
                    ]
              ]
          )

    it "renders allocate_stale_primary with mandatory accept_data_loss" $ do
      let cmd = RerouteAllocateStalePrimary idx shard n1 True
      encodedAs cmd
        `shouldBe` Just
          ( object
              [ "allocate_stale_primary"
                  .= object
                    [ "index" .= ("twitter" :: Text),
                      "shard" .= (0 :: Int),
                      "node" .= ("node-1" :: Text),
                      "accept_data_loss" .= True
                    ]
              ]
          )

    it "renders allocate_empty_primary with mandatory accept_data_loss" $ do
      let cmd = RerouteAllocateEmptyPrimary idx shard n2 False
      encodedAs cmd
        `shouldBe` Just
          ( object
              [ "allocate_empty_primary"
                  .= object
                    [ "index" .= ("twitter" :: Text),
                      "shard" .= (0 :: Int),
                      "node" .= ("node-2" :: Text),
                      "accept_data_loss" .= False
                    ]
              ]
          )

    it "passes RerouteCmdOther through verbatim" $ do
      let v = object ["allocate" .= object ["foo" .= ("bar" :: Text)]] :: Value
      -- The opaque command object is sent unchanged.
      encodedAs (RerouteCmdOther v) `shouldBe` Just v

  describe "rerouteCluster request body" $ do
    it "wraps the commands in a {\"commands\":[...]} object" $ do
      let idx = [qqIndexName|twitter|]
          cmds =
            [ RerouteMove idx (ShardId 0) (NodeName "a") (NodeName "b"),
              RerouteCancel idx (ShardId 1) (NodeName "a") Nothing
            ]
          req = Common.rerouteCluster cmds
          Just body = bhRequestBody req
      Aeson.decode body `shouldBe` Just (object ["commands" .= cmds])

  describe "rerouteOptionsParams URI rendering" $ do
    let normalize :: [(T.Text, Maybe T.Text)] -> [(T.Text, Maybe T.Text)]
        normalize = L.sort

    it "defaultRerouteOptions emits no params" $
      rerouteOptionsParams defaultRerouteOptions `shouldBe` []

    it "renders dry_run as lowercase true/false" $ do
      let opts = defaultRerouteOptions {roDryRun = Just True}
      normalize (rerouteOptionsParams opts) `shouldBe` [("dry_run", Just "true")]
      let opts' = defaultRerouteOptions {roDryRun = Just False}
      normalize (rerouteOptionsParams opts') `shouldBe` [("dry_run", Just "false")]

    it "renders explain as lowercase true/false" $ do
      let opts = defaultRerouteOptions {roExplain = Just True}
      normalize (rerouteOptionsParams opts) `shouldBe` [("explain", Just "true")]

    it "renders metric as a comma-separated list" $ do
      let opts = defaultRerouteOptions {roMetric = Just ("blocks" NE.:| ["nodes"])}
      normalize (rerouteOptionsParams opts) `shouldBe` [("metric", Just "blocks,nodes")]

    it "renders master_timeout and timeout as <n><suffix>" $ do
      let opts =
            defaultRerouteOptions
              { roMasterTimeout = Just (TimeUnitSeconds, 30),
                roTimeout = Just (TimeUnitMinutes, 1)
              }
      normalize (rerouteOptionsParams opts)
        `shouldBe` [("master_timeout", Just "30s"), ("timeout", Just "1m")]

  describe "RerouteResponse FromJSON" $ do
    it "decodes a bare acknowledgement" $ do
      let payload = "{\"acknowledged\":true}"
          Just decoded = Aeson.decode payload :: Maybe RerouteResponse
      rrAcknowledged decoded `shouldBe` True
      rrState decoded `shouldBe` Nothing
      rrExplanations decoded `shouldBe` Nothing

    it "decodes an acknowledgement carrying a cluster state (dry_run)" $ do
      let payload =
            "{\"acknowledged\":true,\"state\":{\"cluster_name\":\"es\",\
            \\"cluster_uuid\":\"uuid\",\"version\":7,\"state_uuid\":\"s\",\
            \\"master_node\":\"m\",\"nodes\":{},\"metadata\":{},\
            \\"routing_table\":{}}}"
          Just decoded = Aeson.decode payload :: Maybe RerouteResponse
      rrAcknowledged decoded `shouldBe` True
      rrState decoded `shouldSatisfy` isJust

    it "decodes an acknowledgement carrying explanations (explain)" $ do
      let payload =
            "{\"acknowledged\":true,\"state\":{\"cluster_name\":\"es\",\
            \\"cluster_uuid\":\"u\",\"version\":1,\"state_uuid\":\"s\",\
            \\"nodes\":{}},\"explanations\":[{\"node\":\"n\",\"decisions\":[]}]}"
          Just decoded = Aeson.decode payload :: Maybe RerouteResponse
      rrAcknowledged decoded `shouldBe` True
      rrExplanations decoded `shouldSatisfy` (maybe False (not . null))

    it "rejects a payload missing the acknowledged key" $
      (Aeson.decode "{\"state\":{}}" :: Maybe RerouteResponse) `shouldBe` Nothing

  -- ------------------------------------------------------------------ --
  -- rerouteCluster endpoint shape: pure checks against the BHRequest     --
  -- value; no live backend needed. Mirrors the updateClusterSettings     --
  -- shape tests in Test.ClusterSettingsSpec.                            --
  -- ------------------------------------------------------------------ --
  describe "rerouteCluster endpoint shape" $ do
    let cmds = [RerouteAllocateReplica [qqIndexName|twitter|] (ShardId 0) (NodeName "n")]

    it "POSTs /_cluster/reroute with no queries by default" $ do
      let req = Common.rerouteCluster cmds
      getRawEndpoint (bhRequestEndpoint req) `shouldBe` ["_cluster", "reroute"]
      getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []

    it "attaches the encoded commands as the request body" $ do
      let req = Common.rerouteCluster cmds
      bhRequestBody req `shouldBe` Just (encode (object ["commands" .= cmds]))

    it "forwards dry_run via rerouteClusterWith" $ do
      let opts = defaultRerouteOptions {roDryRun = Just True}
          req = Common.rerouteClusterWith opts cmds
      getRawEndpoint (bhRequestEndpoint req) `shouldBe` ["_cluster", "reroute"]
      getRawEndpointQueries (bhRequestEndpoint req)
        `shouldBe` [("dry_run", Just "true")]

  -- ------------------------------------------------------------------ --
  -- Live round-trip: a dry_run reroute applies nothing to the cluster   --
  -- but exercises the full RerouteResponse decoder (acknowledged +      --
  -- ClusterState) against a real backend.                              --
  -- ------------------------------------------------------------------ --
  describe "rerouteClusterWith live dry-run round-trip" $
    it "dry_run=true with empty commands acknowledges and returns the state" $
      withTestEnv $
        do
          let opts = defaultRerouteOptions {roDryRun = Just True}
          resp <- Client.rerouteClusterWith opts []
          liftIO $ rrAcknowledged resp `shouldBe` True
          -- ES 9 dropped the (very large) cluster @state@ from the
          -- dry_run reroute response — only @acknowledged@ comes back.
          -- ES 7/8 still include it, so the state assertion is gated
          -- to majors below 9.
          major <- liftIO fetchMajorVersion
          liftIO $ when (major < 9) $ rrState resp `shouldSatisfy` isJust