EnrichSpec.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.EnrichSpec (spec) where

import Data.Aeson
import Data.Aeson.KeyMap qualified as KM
import Data.ByteString.Lazy.Char8 qualified as LBS
import Data.List.NonEmpty (NonEmpty (..))
import Data.List.NonEmpty qualified as NE
import TestsUtils.Import
import Prelude

-- | A canonical @match@ policy body matching the ES 7.17 docs example.
sampleMatchPolicyBytes :: LBS.ByteString
sampleMatchPolicyBytes =
  "{\
  \  \"match\": {\
  \    \"indices\": [\"users\"],\
  \    \"match_field\": \"email\",\
  \    \"enrich_fields\": [\"first_name\", \"last_name\"]\
  \  }\
  \}"

-- | A @geo_match@ policy that exercises a non-default 'EnrichPolicyType'
-- and a bare-string @indices@ (ES accepts a single string there).
sampleGeoPolicyBytes :: LBS.ByteString
sampleGeoPolicyBytes =
  "{\
  \  \"geo_match\": {\
  \    \"indices\": \"users-geo-*\",\
  \    \"match_field\": \"geo_location\",\
  \    \"enrich_fields\": \"city\"\
  \  }\
  \}"

-- | A policy carrying an optional @query@ and an unknown sibling field
-- (@description@) which the @epcExtras@ catch-all must preserve.
samplePolicyWithExtrasBytes :: LBS.ByteString
samplePolicyWithExtrasBytes =
  "{\
  \  \"range\": {\
  \    \"indices\": [\"orders\"],\
  \    \"match_field\": \"customer_id\",\
  \    \"enrich_fields\": [\"tier\"],\
  \    \"query\": {\"term\": {\"active\": true}},\
  \    \"description\": \"customer-tier lookup\"\
  \  }\
  \}"

-- | Full @GET /_enrich/policy@ response with two entries.
sampleListResponseBytes :: LBS.ByteString
sampleListResponseBytes =
  "{\
  \  \"policies\": [\
  \    {\
  \      \"name\": \"users-policy\",\
  \      \"config\": {\
  \        \"match\": {\
  \          \"indices\": [\"users\"],\
  \          \"match_field\": \"email\",\
  \          \"enrich_fields\": [\"first_name\"]\
  \        }\
  \      }\
  \    },\
  \    {\
  \      \"name\": \"geo-policy\",\
  \      \"config\": {\
  \        \"geo_match\": {\
  \          \"indices\": [\"geodata\"],\
  \          \"match_field\": \"location\",\
  \          \"enrich_fields\": [\"city\", \"state\"]\
  \        }\
  \      }\
  \    }\
  \  ]\
  \}"

-- | @GET /_enrich/_execute_stats@ response, ES 7.x shape:
-- @last_execution_time@ as epoch-millis number.
sampleExecuteStatsEpochMillisBytes :: LBS.ByteString
sampleExecuteStatsEpochMillisBytes =
  "{\
  \  \"executor_pool\": \"generic\",\
  \  \"execute_stats\": [\
  \    {\
  \      \"policy\": \"users-policy\",\
  \      \"last_execution_time\": 1719394582880,\
  \      \"running_policy_index_search\": false,\
  \      \"phase\": \"ENDED\"\
  \    }\
  \  ]\
  \}"

-- | @GET /_enrich/_execute_stats@ response, ES 8.x\/9.x shape:
-- @last_execution_time@ as an ISO-8601 string, plus an unknown sibling
-- field (@remote_indices@) that the @eseeExtras@ catch-all must keep.
sampleExecuteStatsIsoBytes :: LBS.ByteString
sampleExecuteStatsIsoBytes =
  "{\
  \  \"executor_pool\": \"generic\",\
  \  \"execute_stats\": [\
  \    {\
  \      \"policy\": \"geo-policy\",\
  \      \"last_execution_time\": \"2024-06-26T10:16:22.880Z\",\
  \      \"phase\": \"RUNNING\",\
  \      \"remote_indices\": [\"remote:geodata\"]\
  \    }\
  \  ]\
  \}"

spec :: Spec
spec = describe "Enrich policy API" $ do
  describe "EnrichPolicyId JSON" $ do
    it "round-trips as a bare JSON string" $ do
      let Just decoded = decode "\"my-policy\"" :: Maybe EnrichPolicyId
      unEnrichPolicyId decoded `shouldBe` "my-policy"
      encode (EnrichPolicyId "my-policy") `shouldBe` "\"my-policy\""

    it "rejects a non-string value" $ do
      let decoded = decode "42" :: Maybe EnrichPolicyId
      decoded `shouldBe` Nothing

  describe "EnrichPolicyType JSON" $ do
    it "'enrichPolicyTypeText' spells the documented keys" $ do
      enrichPolicyTypeText EnrichPolicyTypeMatch `shouldBe` "match"
      enrichPolicyTypeText EnrichPolicyTypeGeoMatch `shouldBe` "geo_match"
      enrichPolicyTypeText EnrichPolicyTypeRange `shouldBe` "range"

    it "absorbs an unknown policy type into the Custom branch" $ do
      let Just (EnrichPolicyTypeCustom t) =
            decode "\"mismatch\"" :: Maybe EnrichPolicyType
      t `shouldBe` "mismatch"

  describe "EnrichPolicy JSON (PUT body)" $ do
    it "decodes a canonical match policy" $ do
      let Just decoded = decode sampleMatchPolicyBytes :: Maybe EnrichPolicy
      epTypeTag decoded `shouldBe` EnrichPolicyTypeMatch
      let cfg = epConfig decoded
      NE.toList (epcIndices cfg) `shouldBe` [IndexPattern "users"]
      epcMatchField cfg `shouldBe` FieldName "email"
      NE.toList (epcEnrichFields cfg)
        `shouldBe` [FieldName "first_name", FieldName "last_name"]
      epcQuery cfg `shouldBe` Nothing
      epcExtras cfg `shouldBe` KM.empty

    it "decodes a geo_match policy with bare-string indices/enrich_fields" $ do
      let Just decoded = decode sampleGeoPolicyBytes :: Maybe EnrichPolicy
      epTypeTag decoded `shouldBe` EnrichPolicyTypeGeoMatch
      let cfg = epConfig decoded
      -- A bare string is normalised to a singleton list.
      NE.toList (epcIndices cfg) `shouldBe` [IndexPattern "users-geo-*"]
      NE.toList (epcEnrichFields cfg) `shouldBe` [FieldName "city"]

    it "preserves an optional query and unknown sibling fields in extras" $ do
      let Just decoded = decode samplePolicyWithExtrasBytes :: Maybe EnrichPolicy
          cfg = epConfig decoded
      epTypeTag decoded `shouldBe` EnrichPolicyTypeRange
      -- The @query@ body is carried as an opaque Value.
      epcQuery cfg `shouldSatisfy` isJust
      -- The unknown @description@ key survives in the extras catch-all.
      KM.lookup "description" (epcExtras cfg) `shouldSatisfy` isJust

    it "encodes the type-tagged shape with array indices/enrich_fields" $ do
      let Just decoded = decode sampleMatchPolicyBytes :: Maybe EnrichPolicy
          encoded = encode decoded
          Just roundTripped = decode encoded :: Maybe EnrichPolicy
      -- Round-trip is stable.
      epTypeTag roundTripped `shouldBe` EnrichPolicyTypeMatch
      NE.toList (epcEnrichFields (epConfig roundTripped))
        `shouldBe` [FieldName "first_name", FieldName "last_name"]

    it "round-trips a policy with extras through encode . decode" $ do
      let Just decoded = decode samplePolicyWithExtrasBytes :: Maybe EnrichPolicy
          Just roundTripped = decode (encode decoded) :: Maybe EnrichPolicy
          cfg = epConfig roundTripped
      KM.lookup "description" (epcExtras cfg) `shouldSatisfy` isJust
      -- The known @query@ key is NOT duplicated into extras.
      KM.lookup "query" (epcExtras cfg) `shouldBe` Nothing

    it "rejects an empty object body" $ do
      let decoded = decode "{}" :: Maybe EnrichPolicy
      decoded `shouldBe` Nothing

    it "rejects a multi-key body" $ do
      let decoded =
            decode "{\"match\":{},\"geo_match\":{}}" :: Maybe EnrichPolicy
      decoded `shouldBe` Nothing

    it "rejects a policy missing required 'indices'" $ do
      let decoded =
            decode
              "{\"match\":{\"match_field\":\"email\",\"enrich_fields\":[\"x\"]}}" ::
              Maybe EnrichPolicy
      decoded `shouldBe` Nothing

    it "rejects a policy with an empty 'enrich_fields' array" $ do
      let decoded =
            decode
              "{\"match\":{\"indices\":[\"u\"],\"match_field\":\"e\",\"enrich_fields\":[]}}" ::
              Maybe EnrichPolicy
      decoded `shouldBe` Nothing

    it "rejects a policy with an empty 'indices' array" $ do
      let decoded =
            decode
              "{\"match\":{\"indices\":[],\"match_field\":\"e\",\"enrich_fields\":[\"x\"]}}" ::
              Maybe EnrichPolicy
      decoded `shouldBe` Nothing

  describe "EnrichPoliciesResponse JSON (GET response)" $ do
    it "decodes the {policies: [...]} envelope" $ do
      let Just decoded =
            decode sampleListResponseBytes :: Maybe EnrichPoliciesResponse
      length (unEnrichPoliciesResponse decoded) `shouldBe` 2
      let firstInfo = head (unEnrichPoliciesResponse decoded)
      epiName firstInfo `shouldBe` EnrichPolicyId "users-policy"
      epTypeTag (epiPolicy firstInfo)
        `shouldBe` EnrichPolicyTypeMatch
      let secondInfo = unEnrichPoliciesResponse decoded !! 1
      epTypeTag (epiPolicy secondInfo)
        `shouldBe` EnrichPolicyTypeGeoMatch

    it "re-encodes to the same envelope shape" $ do
      let Just decoded =
            decode sampleListResponseBytes :: Maybe EnrichPoliciesResponse
          reEncoded = encode decoded
          Just reparsed =
            decode reEncoded :: Maybe EnrichPoliciesResponse
      length (unEnrichPoliciesResponse reparsed) `shouldBe` 2

  describe "EnrichPolicyPhase JSON" $ do
    it "decodes documented values" $ do
      decode "\"ENDED\"" `shouldBe` Just EnrichPolicyPhaseEnded
      decode "\"RUNNING\"" `shouldBe` Just EnrichPolicyPhaseRunning

    it "encodes documented values back to the wire spelling" $ do
      encode EnrichPolicyPhaseEnded `shouldBe` "\"ENDED\""
      encode EnrichPolicyPhaseRunning `shouldBe` "\"RUNNING\""

    it "absorbs an unknown phase into the Custom branch" $ do
      let Just (EnrichPolicyPhaseCustom t) =
            decode "\"WAITING_FOR_LEADER\"" :: Maybe EnrichPolicyPhase
      t `shouldBe` "WAITING_FOR_LEADER"

  describe "EnrichExecuteStatsResponse JSON" $ do
    it "decodes the ES 7.x epoch-millis shape" $ do
      let Just decoded =
            decode sampleExecuteStatsEpochMillisBytes ::
              Maybe EnrichExecuteStatsResponse
      eesrExecutorPool decoded `shouldBe` Just "generic"
      length (eesrStats decoded) `shouldBe` 1
      let entry = head (eesrStats decoded)
      eseePolicy entry `shouldBe` EnrichPolicyId "users-policy"
      eseeRunningPolicyIndexSearch entry `shouldBe` Just False
      eseePhase entry `shouldBe` Just EnrichPolicyPhaseEnded
      -- @last_execution_time@ is carried as an opaque Value (epoch millis).
      eseeLastExecutionTime entry `shouldSatisfy` isJust

    it "decodes the ES 8.x/9.x ISO-8601 shape and preserves unknown fields" $ do
      let Just decoded =
            decode sampleExecuteStatsIsoBytes ::
              Maybe EnrichExecuteStatsResponse
      let entry = head (eesrStats decoded)
      eseePolicy entry `shouldBe` EnrichPolicyId "geo-policy"
      eseePhase entry `shouldBe` Just EnrichPolicyPhaseRunning
      -- The unknown @remote_indices@ sibling is preserved in extras.
      KM.lookup "remote_indices" (eseeExtras entry) `shouldSatisfy` isJust
      -- @running_policy_index_search@ is absent on the wire, so it is Nothing.
      eseeRunningPolicyIndexSearch entry `shouldBe` Nothing

    it "round-trips the execute-stats response" $ do
      let Just decoded =
            decode sampleExecuteStatsEpochMillisBytes ::
              Maybe EnrichExecuteStatsResponse
          Just roundTripped =
            decode (encode decoded) :: Maybe EnrichExecuteStatsResponse
      length (eesrStats roundTripped) `shouldBe` 1

    it "preserves unknown execute-stats fields through encode . decode" $ do
      let Just decoded =
            decode sampleExecuteStatsIsoBytes ::
              Maybe EnrichExecuteStatsResponse
          Just roundTripped =
            decode (encode decoded) :: Maybe EnrichExecuteStatsResponse
          entry = head (eesrStats roundTripped)
      -- The @remote_indices@ extra survives the round-trip.
      KM.lookup "remote_indices" (eseeExtras entry) `shouldSatisfy` isJust

  describe "EnrichExecuteOptions params" $ do
    it "default options render no query string" $ do
      enrichExecuteOptionsParams defaultEnrichExecuteOptions
        `shouldBe` []

    it "wait_for_completion is rendered as a bool string" $ do
      let opts =
            defaultEnrichExecuteOptions
              { eeoWaitForCompletion = Just False
              }
      enrichExecuteOptionsParams opts
        `shouldBe` [("wait_for_completion", Just "false")]

    it "wait_for_completion=True renders 'true'" $ do
      let opts =
            defaultEnrichExecuteOptions
              { eeoWaitForCompletion = Just True
              }
      enrichExecuteOptionsParams opts
        `shouldBe` [("wait_for_completion", Just "true")]

  describe "defaultEnrichPolicyConfig" $ do
    it "builds a minimal config with single-element lists" $ do
      let cfg =
            defaultEnrichPolicyConfig
              (IndexPattern "users")
              (FieldName "email")
              (FieldName "first_name" :| [])
      NE.toList (epcIndices cfg) `shouldBe` [IndexPattern "users"]
      epcMatchField cfg `shouldBe` FieldName "email"
      NE.toList (epcEnrichFields cfg) `shouldBe` [FieldName "first_name"]
      epcQuery cfg `shouldBe` Nothing
      epcExtras cfg `shouldBe` KM.empty