MlAnomalySpec.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.MlAnomalySpec (spec) where

import Data.Aeson
import Data.Aeson.KeyMap qualified as KeyMap
import Data.ByteString.Lazy.Char8 qualified as LBS
import Data.Scientific (Scientific)
import Database.Bloodhound.ElasticSearch9.Requests qualified as RequestsES9
import Database.Bloodhound.ElasticSearch9.Types qualified as Types
import TestsUtils.Import

spec :: Spec
spec =
  describe "ML anomaly APIs (/_ml/anomaly_detectors, /_ml/datafeeds, /_ml/filters, /_ml/calendars)" $ do
    describe "identities and enums" $ do
      it "round-trips MlJobId / MlModelSnapshotId / MlDatafeedId / MlFilterId / MlCalendarId" $ do
        decode (encode (Types.MlJobId "j")) `shouldBe` Just (Types.MlJobId "j")
        decode (encode (Types.MlModelSnapshotId "s"))
          `shouldBe` Just (Types.MlModelSnapshotId "s")
        decode (encode (Types.MlDatafeedId "f"))
          `shouldBe` Just (Types.MlDatafeedId "f")
        decode (encode (Types.MlFilterId "fi"))
          `shouldBe` Just (Types.MlFilterId "fi")
        decode (encode (Types.MlCalendarId "c"))
          `shouldBe` Just (Types.MlCalendarId "c")

      it "round-trips documented job states and absorbs unknowns" $ do
        mapM_
          (\s -> decode (encode s) `shouldBe` Just s)
          [ Types.MlJobStateOpening,
            Types.MlJobStateOpened,
            Types.MlJobStateClosing,
            Types.MlJobStateClosed,
            Types.MlJobStateFailed
          ]
        decode "\"paused\"" `shouldBe` Just (Types.MlJobStateCustom "paused")

      it "maps result types to documented path segments" $ do
        Types.mlAnomalyResultTypeSegment Types.MlAnomalyResultBuckets
          `shouldBe` ("buckets" :: Text)
        Types.mlAnomalyResultTypeSegment Types.MlAnomalyResultRecords
          `shouldBe` ("records" :: Text)
        Types.mlAnomalyResultTypeSegment Types.MlAnomalyResultOverallBuckets
          `shouldBe` ("overall_buckets" :: Text)

    describe "MlAnomalyJob JSON" $ do
      it "encodes the minimal body as analysis_config + data_description" $ do
        let job =
              Types.defaultMlAnomalyJob
                { Types.mlajAnalysisConfig = object ["detectors" .= ([] :: [Int])],
                  Types.mlajDataDescription = object ["time_field" .= ("timestamp" :: Text)]
                }
        decode (encode job) `shouldBe` Just job

      it "omits optional fields" $ do
        let job =
              Types.defaultMlAnomalyJob
                { Types.mlajAnalysisConfig = object []
                }
        case decode (encode job) :: Maybe Value of
          Just (Object o) -> KeyMap.member "description" o `shouldBe` False
          _ -> expectationFailure "expected a JSON object"

    describe "MlAnomalyJobUpdate JSON" $ do
      it "encodes to {} when empty" $
        encode Types.defaultMlAnomalyJobUpdate `shouldBe` "{}"

    describe "MlForecastOptions / MlFlushOptions JSON" $ do
      it "forecast round-trips a populated body" $ do
        let o =
              Types.defaultMlForecastOptions
                { Types.mlfoDuration = Just "1d"
                }
        decode (encode o) `shouldBe` Just o

      it "flush encodes to {} when empty" $
        encode Types.defaultMlFlushOptions `shouldBe` "{}"

    describe "MlAnomalyJobOptions params" $ do
      it "renders nothing by default" $
        Types.mlAnomalyJobOptionsParams Types.defaultMlAnomalyJobOptions
          `shouldBe` []

      it "renders force/timeout/from in order" $ do
        let opts =
              Types.defaultMlAnomalyJobOptions
                { Types.mlajoForce = Just True,
                  Types.mlajoTimeout = Just "30s",
                  Types.mlajoFrom = Just 10
                }
        Types.mlAnomalyJobOptionsParams opts
          `shouldBe` [ ("from", Just "10"),
                       ("force", Just "true"),
                       ("timeout", Just "30s")
                     ]

    describe "MlAnomalyJobDocument JSON" $ do
      it "decodes a stored job, folding extras" $ do
        let raw =
              LBS.pack
                "{\"job_id\":\"j\",\"state\":\"open\",\"job_type\":\"anomaly_detector\",\
                \\"create_time\":1700000000000,\"analysis_config\":{}}"
        case decode raw :: Maybe Types.MlAnomalyJobDocument of
          Just d -> do
            Types.mlajdId d `shouldBe` Just ("j" :: Text)
            Types.mlajdState d `shouldBe` Just Types.MlJobStateOpened
            extras <- pure (Types.mlajdExtras d)
            extras `shouldSatisfy` isJust
            forM_ extras $ \km ->
              KeyMap.lookup "analysis_config" km `shouldSatisfy` isJust
          Nothing -> expectationFailure "failed to decode document"

    describe "MlAnomalyJobsResponse JSON" $ do
      it "decodes a jobs envelope, defaulting count" $ do
        let raw = LBS.pack "{\"jobs\":[{\"job_id\":\"a\"},{\"job_id\":\"b\"}]}"
        case decode raw :: Maybe Types.MlAnomalyJobsResponse of
          Just r -> do
            length (Types.mlajsrJobs r) `shouldBe` 2
            Types.mlajsrCount r `shouldBe` Just 2
          Nothing -> expectationFailure "failed to decode jobs response"

    describe "MlAnomalyJobStatsResponse JSON" $ do
      it "decodes a stats envelope" $ do
        let raw =
              LBS.pack
                "{\"jobs\":[{\"job_id\":\"a\",\"state\":\"closed\",\"data_counts\":{}}]}"
        case decode raw :: Maybe Types.MlAnomalyJobStatsResponse of
          Just r -> do
            length (Types.mlajsrrStats r) `shouldBe` 1
            case Types.mlajsrrStats r of
              (s : _) -> Types.mlajssState s `shouldBe` Just Types.MlJobStateClosed
              [] -> expectationFailure "expected one stats entry"
          Nothing -> expectationFailure "failed to decode stats response"

    describe "forecast / flush responses" $ do
      it "decodes a forecast response, defaulting acknowledged" $ do
        let raw = LBS.pack "{\"forecast_id\":\"abc\"}"
        case decode raw :: Maybe Types.MlForecastResponse of
          Just r -> do
            Types.mlfrAcknowledged r `shouldBe` Just True
            Types.mlfrForecastId r `shouldBe` Just ("abc" :: Text)
          Nothing -> expectationFailure "failed to decode forecast response"

      it "decodes a flush response" $ do
        let raw = LBS.pack "{\"acknowledged\":true,\"last_finalized_bucket_end\":1234}"
        case decode raw :: Maybe Types.MlFlushResponse of
          Just r -> Types.mlflrLastFinalizedBucketEnd r `shouldBe` Just (1234 :: Scientific)
          Nothing -> expectationFailure "failed to decode flush response"

    describe "model snapshots" $ do
      it "decodes a snapshots envelope, defaulting count" $ do
        let raw =
              LBS.pack
                "{\"model_snapshots\":[{\"id\":\"s1\",\"description\":\"d\"}]}"
        case decode raw :: Maybe Types.MlModelSnapshotsResponse of
          Just r -> do
            length (Types.mlmsrSnapshots r) `shouldBe` 1
            Types.mlmsrCount r `shouldBe` Just 1
          Nothing -> expectationFailure "failed to decode snapshots response"

      it "decodes a revert response" $ do
        let raw =
              LBS.pack
                "{\"model\":{\"foo\":1},\"model_snapshot_id\":\"s1\"}"
        case decode raw :: Maybe Types.MlRevertSnapshotResponse of
          Just r -> do
            Types.mlrsrModelSnapshotId r `shouldBe` Just ("s1" :: Text)
            Types.mlrsrModel r `shouldSatisfy` isJust
          Nothing -> expectationFailure "failed to decode revert response"

    describe "MlDatafeed JSON" $ do
      it "encodes the minimal body with job_id only" $ do
        let df = Types.defaultMlDatafeed {Types.mldfJobId = "myjob"}
        case decode (encode df) :: Maybe Types.MlDatafeed of
          Just d -> Types.mldfJobId d `shouldBe` ("myjob" :: Text)
          Nothing -> expectationFailure "datafeed should round-trip"

      it "decodes a datafeed document" $ do
        let raw = LBS.pack "{\"datafeeds\":[{\"datafeed_id\":\"f\",\"job_id\":\"j\"}]}"
        case decode raw :: Maybe Types.MlDatafeedsResponse of
          Just r -> do
            length (Types.mldfrDatafeeds r) `shouldBe` 1
            case Types.mldfrDatafeeds r of
              (d : _) -> Types.mldfdJobId d `shouldBe` Just ("j" :: Text)
              [] -> expectationFailure "expected one datafeed"
          Nothing -> expectationFailure "failed to decode datafeeds response"

    describe "filters" $ do
      it "decodes a filters response" $ do
        let raw =
              LBS.pack
                "{\"filters\":[{\"id\":\"fi\",\"items\":[\"a\"]}]}"
        case decode raw :: Maybe Types.MlFiltersResponse of
          Just r -> length (Types.mlfsrFilters r) `shouldBe` 1
          Nothing -> expectationFailure "failed to decode filters response"

      it "round-trips a filter body" $ do
        let f =
              Types.defaultMlFilter
                { Types.mlfItems = Just ["a", "b"]
                }
        decode (encode f) `shouldBe` Just f

    describe "calendars and scheduled events" $ do
      it "decodes a calendars response" $ do
        let raw =
              LBS.pack
                "{\"calendars\":[{\"calendar_id\":\"c\",\"description\":\"d\"}]}"
        case decode raw :: Maybe Types.MlCalendarsResponse of
          Just r -> length (Types.mlcsrCalendars r) `shouldBe` 1
          Nothing -> expectationFailure "failed to decode calendars response"

      it "encodes scheduled events as {events: [...]}" $ do
        let e =
              Types.MlScheduledEvents
                { Types.mlseEvents = [object ["description" .= ("d" :: Text)]]
                }
        encode e `shouldBe` "{\"events\":[{\"description\":\"d\"}]}"

    describe "endpoint shape" $ do
      it "PUTs /_ml/anomaly_detectors/<id>" $ do
        let req = RequestsES9.putMlAnomalyJob "j" Types.defaultMlAnomalyJob
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "anomaly_detectors", "j"]
        bhRequestBody req `shouldSatisfy` isJust

      it "GETs /_ml/anomaly_detectors/<id>" $ do
        let req = RequestsES9.getMlAnomalyJobs (Just "j")
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "anomaly_detectors", "j"]

      it "POSTs /_ml/anomaly_detectors/<id>/_open (empty body)" $ do
        let req = RequestsES9.openMlAnomalyJob "j"
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "anomaly_detectors", "j", "_open"]
        bhRequestBody req `shouldSatisfy` isJust

      it "GETs /_ml/anomaly_detectors/<id>/_results/<kind>" $ do
        let req = RequestsES9.getMlAnomalyJobResults "j" Types.MlAnomalyResultBuckets
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "anomaly_detectors", "j", "_results", "buckets"]

      it "GETs /_ml/anomaly_detectors/<id>/_model_snapshots/<snap>" $ do
        let req = RequestsES9.getMlModelSnapshot "j" "snap"
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "anomaly_detectors", "j", "_model_snapshots", "snap"]

      it "POSTs /_ml/anomaly_detectors/<id>/_model_snapshots/<snap>/_revert" $ do
        let req = RequestsES9.revertMlModelSnapshot "j" "snap" (object [])
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "anomaly_detectors", "j", "_model_snapshots", "snap", "_revert"]
        bhRequestBody req `shouldSatisfy` isJust

      it "POSTs /_ml/anomaly_detectors/_validate" $ do
        let req = RequestsES9.validateMlAnomalyJob Types.defaultMlAnomalyJob
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "anomaly_detectors", "_validate"]

      it "PUTs /_ml/datafeeds/<id>" $ do
        let req = RequestsES9.putMlDatafeed "f" Types.defaultMlDatafeed
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "datafeeds", "f"]

      it "POSTs /_ml/datafeeds/<id>/_start with a body" $ do
        let req = RequestsES9.startMlDatafeed "f" (object [])
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "datafeeds", "f", "_start"]
        bhRequestBody req `shouldSatisfy` isJust

      it "GETs /_ml/datafeeds/<id>/_preview" $ do
        let req = RequestsES9.previewMlDatafeed "f"
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "datafeeds", "f", "_preview"]

      it "PUTs /_ml/filters/<id>" $ do
        let req = RequestsES9.putMlFilter "fi" Types.defaultMlFilter
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "filters", "fi"]

      it "GETs /_ml/filters/<id>" $ do
        let req = RequestsES9.getMlFilters (Just "fi")
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "filters", "fi"]

      it "PUTs /_ml/calendars/<id>" $ do
        let req = RequestsES9.putMlCalendar "c" Types.defaultMlCalendar
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "calendars", "c"]

      it "POSTs /_ml/calendars/<id>/events" $ do
        let req = RequestsES9.setMlScheduledEvents "c" (Types.MlScheduledEvents [])
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "calendars", "c", "events"]
        bhRequestBody req `shouldSatisfy` isJust

      it "GETs /_ml/calendars/<id>/events" $ do
        let req = RequestsES9.getMlScheduledEvents "c"
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "calendars", "c", "events"]