packages feed

bloodhound-1.0.0.0: tests/Test/MlDataFrameAnalyticsSpec.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.MlDataFrameAnalyticsSpec (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

hasKey :: Key -> LBS.ByteString -> Bool
hasKey k bs = case decode bs of
  Just (Object obj) -> k `KeyMap.member` obj
  _ -> False

spec :: Spec
spec =
  describe "ML data frame analytics APIs (/_ml/data_frame/analytics/*)" $ do
    describe "MlDataFrameAnalyticsId JSON" $ do
      it "round-trips an id" $ do
        let i = Types.MlDataFrameAnalyticsId "house_price_regression"
        decode (encode i) `shouldBe` Just i

      it "unMlDataFrameAnalyticsId extracts the underlying Text" $
        Types.unMlDataFrameAnalyticsId (Types.MlDataFrameAnalyticsId "x")
          `shouldBe` ("x" :: Text)

    describe "MlDataFrameAnalysisType JSON" $ do
      it "encodes the documented values" $ do
        encode Types.MlDataFrameAnalysisTypeClassification
          `shouldBe` "\"classification\""
        encode Types.MlDataFrameAnalysisTypeRegression
          `shouldBe` "\"regression\""
        encode Types.MlDataFrameAnalysisTypeOutlierDetection
          `shouldBe` "\"outlier_detection\""

      it "round-trips a custom value" $ do
        let t = Types.MlDataFrameAnalysisTypeCustom "future_kind"
        decode (encode t) `shouldBe` Just t

      it "decodes documented values back" $
        decode "\"regression\"" `shouldBe` Just Types.MlDataFrameAnalysisTypeRegression

    describe "MlDataFrameAnalyticsState JSON" $ do
      it "round-trips documented states" $ do
        mapM_
          (\s -> decode (encode s) `shouldBe` Just s)
          [ Types.MlDataFrameAnalyticsStateStopped,
            Types.MlDataFrameAnalyticsStateStarting,
            Types.MlDataFrameAnalyticsStateStarted,
            Types.MlDataFrameAnalyticsStateStopping,
            Types.MlDataFrameAnalyticsStateReverting,
            Types.MlDataFrameAnalyticsStateFailed
          ]

      it "absorbs an unknown state" $ do
        decode "\"awaiting_input\""
          `shouldBe` Just (Types.MlDataFrameAnalyticsStateCustom "awaiting_input")

    describe "MlDataFrameAnalyticsSource JSON" $ do
      it "omits Nothing fields" $ do
        let s =
              Types.defaultMlDataFrameAnalyticsSource
                { Types.mldasIndex = ["source-*"]
                }
        hasKey "query" (encode s) `shouldBe` False
        hasKey "_source" (encode s) `shouldBe` False
        decode (encode s) `shouldBe` Just s

      it "decodes a documented source" $ do
        let raw =
              LBS.pack
                "{\"index\":[\"source-*\"],\"_source\":[\"price\",\"sqft\"]}"
        case decode raw :: Maybe Types.MlDataFrameAnalyticsSource of
          Just s -> do
            Types.mldasIndex s `shouldBe` ["source-*"]
            Types.mldasSource s `shouldBe` Just ["price" :: Text, "sqft"]
          Nothing -> expectationFailure "failed to decode source"

    describe "MlDataFrameAnalyticsDest JSON" $ do
      it "encodes only the index when results_field is Nothing" $ do
        let d = Types.defaultMlDataFrameAnalyticsDest {Types.mldadIndex = Just "dest"}
        hasKey "results_field" (encode d) `shouldBe` False
        decode (encode d) `shouldBe` Just d

    describe "MlDataFrameAnalyticsPutBody JSON" $ do
      it "omits every optional field for the minimal body" $ do
        let body =
              Types.defaultMlDataFrameAnalyticsPutBody
                { Types.mldapbSource =
                    Types.defaultMlDataFrameAnalyticsSource
                      { Types.mldasIndex = ["source-*"]
                      },
                  Types.mldapbDest =
                    Types.defaultMlDataFrameAnalyticsDest
                      { Types.mldadIndex = Just "dest"
                      }
                }
        hasKey "description" (encode body) `shouldBe` False
        hasKey "analysis" (encode body) `shouldBe` False
        decode (encode body) `shouldBe` Just body

      it "encodes an analysis sub-object verbatim" $ do
        let cfg = object ["outlier_detection" .= object ["n_neighbors" .= (4 :: Int)]]
            analysis = Types.MlDataFrameAnalyticsAnalysis cfg
            body =
              Types.defaultMlDataFrameAnalyticsPutBody
                { Types.mldapbAnalysis = Just analysis
                }
        hasKey "analysis" (encode body) `shouldBe` True
        Types.mldaanConfig
          <$> Types.mldapbAnalysis body
            `shouldBe` Just cfg

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

      it "round-trips with a description" $ do
        let b =
              Types.defaultMlDataFrameAnalyticsUpdateBody
                { Types.mldaubDescription = Just "updated"
                }
        decode (encode b) `shouldBe` Just b

    describe "MlDataFrameAnalyticsOptions params" $ do
      it "renders nothing by default" $
        Types.mlDataFrameAnalyticsOptionsParams Types.defaultMlDataFrameAnalyticsOptions
          `shouldBe` []

      it "renders from/size/force/verbose/timeout" $ do
        let opts =
              Types.defaultMlDataFrameAnalyticsOptions
                { Types.mldaoFrom = Just 0,
                  Types.mldaoSize = Just 100,
                  Types.mldaoForce = Just True,
                  Types.mldaoVerbose = Just True,
                  Types.mldaoTimeout = Just "30s"
                }
        Types.mlDataFrameAnalyticsOptionsParams opts
          `shouldBe` [ ("force", Just "true"),
                       ("from", Just "0"),
                       ("size", Just "100"),
                       ("verbose", Just "true"),
                       ("timeout", Just "30s")
                     ]

    describe "MlDataFrameAnalyticsDocument JSON" $ do
      it "decodes a documented stored job verbatim" $ do
        let raw =
              LBS.pack
                "{\"id\":\"house_price_regression\",\"state\":\"stopped\",\
                \\"analysis_type\":\"regression\",\"create_time\":1700000000000,\
                \\"version\":\"17.0.0\",\"source\":{\"index\":[\"source-*\"]},\
                \\"dest\":{\"index\":\"dest\"},\"description\":\"regress\"}"
        case decode raw :: Maybe Types.MlDataFrameAnalyticsDocument of
          Just d -> do
            Types.mldadId d `shouldBe` Just ("house_price_regression" :: Text)
            Types.mldadState d `shouldBe` Just Types.MlDataFrameAnalyticsStateStopped
            Types.mldadAnalysisType d `shouldBe` Just Types.MlDataFrameAnalysisTypeRegression
            Types.mldadCreateTime d `shouldBe` Just (1700000000000 :: Scientific)
            -- description is not a typed Document field, so it lands in extras
            let extras = Types.mldadExtras d
            extras `shouldSatisfy` isJust
            forM_ extras $ \km ->
              KeyMap.lookup "description" km `shouldBe` Just (String "regress")
          Nothing -> expectationFailure "failed to decode document"

      it "folds unknown sibling fields into extras" $ do
        let raw =
              LBS.pack
                "{\"id\":\"x\",\"create_time\":1,\"model_memory_limit\":\"11mb\"}"
        case decode raw :: Maybe Types.MlDataFrameAnalyticsDocument of
          Just d -> do
            Types.mldadId d `shouldBe` Just ("x" :: Text)
            extras <- pure (Types.mldadExtras d)
            extras `shouldSatisfy` isJust
            let Just km = extras
            KeyMap.lookup "model_memory_limit" km `shouldSatisfy` isJust
          Nothing -> expectationFailure "failed to decode document with extras"

      it "round-trips including extras" $ do
        let raw =
              LBS.pack
                "{\"id\":\"x\",\"state\":\"started\",\"future_field\":42}"
        case decode raw :: Maybe Types.MlDataFrameAnalyticsDocument of
          Just d -> encode d `shouldSatisfy` (\bs -> decode bs == Just d)
          Nothing -> expectationFailure "failed to decode"

    describe "MlDataFrameAnalyticsGetResponse JSON" $ do
      it "decodes a list envelope, defaulting count" $ do
        let raw =
              LBS.pack
                "{\"data_frame_analytics\":[{\"id\":\"a\"},{\"id\":\"b\"}]}"
        case decode raw :: Maybe Types.MlDataFrameAnalyticsGetResponse of
          Just r -> do
            length (Types.mldagrAnalytics r) `shouldBe` 2
            Types.mldagrCount r `shouldBe` Just 2
          Nothing -> expectationFailure "failed to decode get response"

      it "decodes an explicit count" $ do
        let raw = LBS.pack "{\"count\":5,\"data_frame_analytics\":[]}"
        case decode raw :: Maybe Types.MlDataFrameAnalyticsGetResponse of
          Just r -> Types.mldagrCount r `shouldBe` Just 5
          Nothing -> expectationFailure "failed to decode get response"

      it "decodes an empty object to count 0, []" $
        case decode (LBS.pack "{}") :: Maybe Types.MlDataFrameAnalyticsGetResponse of
          Just r -> do
            Types.mldagrCount r `shouldBe` Just 0
            Types.mldagrAnalytics r `shouldBe` []
          Nothing -> expectationFailure "failed to decode empty get response"

    describe "MlDataFrameAnalyticsStatsResponse JSON" $ do
      it "decodes a stats envelope, defaulting count" $ do
        let raw =
              LBS.pack
                "{\"data_frame_analytics\":[{\"id\":\"a\",\"state\":\"started\",\
                \\"analysis_type\":\"outlier_detection\",\"progress\":[0.9]}]}"
        case decode raw :: Maybe Types.MlDataFrameAnalyticsStatsResponse of
          Just r -> do
            length (Types.mldasrStats r) `shouldBe` 1
            Types.mldasrCount r `shouldBe` Just 1
            case Types.mldasrStats r of
              (s : _) -> do
                Types.mldasState s `shouldBe` Just Types.MlDataFrameAnalyticsStateStarted
                Types.mldasAnalysisType s `shouldBe` Just Types.MlDataFrameAnalysisTypeOutlierDetection
                let Just km = Types.mldasExtras s
                KeyMap.size km `shouldBe` 1
              [] -> expectationFailure "expected one stats entry"
          Nothing -> expectationFailure "failed to decode stats response"

    describe "MlDataFrameAnalyticsPreviewResponse JSON" $ do
      it "decodes an array of row objects" $ do
        let raw =
              LBS.pack
                "[{\"a\":1},{\"a\":2}]"
        case decode raw :: Maybe Types.MlDataFrameAnalyticsPreviewResponse of
          Just r -> length (Types.mldaprrRows r) `shouldBe` 2
          Nothing -> expectationFailure "failed to decode preview rows"

      it "decodes a non-array body to []" $
        case decode (LBS.pack "{}") :: Maybe Types.MlDataFrameAnalyticsPreviewResponse of
          Just r -> Types.mldaprrRows r `shouldBe` []
          Nothing -> expectationFailure "failed to decode non-array preview"

    describe "endpoint shape" $ do
      it "PUTs /_ml/data_frame/analytics/<id> with a body" $ do
        let body =
              Types.defaultMlDataFrameAnalyticsPutBody
                { Types.mldapbSource =
                    Types.defaultMlDataFrameAnalyticsSource
                      { Types.mldasIndex = ["source-*"]
                      }
                }
            req = RequestsES9.putMlDataFrameAnalytics "myjob" body
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "data_frame", "analytics", "myjob"]
        bhRequestBody req `shouldSatisfy` isJust

      it "PUT-With forwards options as query params" $ do
        let opts =
              Types.defaultMlDataFrameAnalyticsOptions
                { Types.mldaoAllowLazyStart = Just True
                }
            req =
              RequestsES9.putMlDataFrameAnalyticsWith
                "myjob"
                Types.defaultMlDataFrameAnalyticsPutBody
                opts
        getRawEndpointQueries (bhRequestEndpoint req) `shouldSatisfy` (not . null)

      it "GETs /_ml/data_frame/analytics (all)" $ do
        let req = RequestsES9.getMlDataFrameAnalytics Nothing
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "data_frame", "analytics"]

      it "GETs /_ml/data_frame/analytics/<id>" $ do
        let req = RequestsES9.getMlDataFrameAnalytics (Just "myjob")
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "data_frame", "analytics", "myjob"]

      it "GETs /_ml/data_frame/analytics/<id>/_stats" $ do
        let req = RequestsES9.getMlDataFrameAnalyticsStats (Just "myjob")
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "data_frame", "analytics", "myjob", "_stats"]

      it "GETs /_ml/data_frame/analytics/_stats (all)" $ do
        let req = RequestsES9.getMlDataFrameAnalyticsStats Nothing
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "data_frame", "analytics", "_stats"]

      it "POSTs /_ml/data_frame/analytics/<id>/_update" $ do
        let req =
              RequestsES9.updateMlDataFrameAnalytics
                "myjob"
                Types.defaultMlDataFrameAnalyticsUpdateBody
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "data_frame", "analytics", "myjob", "_update"]
        bhRequestBody req `shouldSatisfy` isJust

      it "DELETEs /_ml/data_frame/analytics/<id>" $ do
        let req = RequestsES9.deleteMlDataFrameAnalytics "myjob"
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "data_frame", "analytics", "myjob"]
        bhRequestBody req `shouldSatisfy` isNothing

      it "POSTs /_ml/data_frame/analytics/<id>/_start" $ do
        let req = RequestsES9.startMlDataFrameAnalytics "myjob"
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "data_frame", "analytics", "myjob", "_start"]
        bhRequestBody req `shouldSatisfy` isJust

      it "POSTs /_ml/data_frame/analytics/<id>/_stop" $ do
        let req = RequestsES9.stopMlDataFrameAnalytics "myjob"
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "data_frame", "analytics", "myjob", "_stop"]

      it "POSTs /_ml/data_frame/analytics/_preview with a body" $ do
        let req =
              RequestsES9.previewMlDataFrameAnalytics $
                Types.MlDataFrameAnalyticsPreviewRequest (object ["id" .= ("x" :: Text)])
        getRawEndpoint (bhRequestEndpoint req)
          `shouldBe` ["_ml", "data_frame", "analytics", "_preview"]
        bhRequestBody req `shouldSatisfy` isJust