bloodhound-1.0.0.0: tests/Test/MLStatsSpec.hs
{-# LANGUAGE OverloadedStrings #-}
module Test.MLStatsSpec (spec) where
import Data.Aeson.KeyMap qualified as KM
import Data.ByteString.Char8 qualified as BS
import Data.ByteString.Lazy.Char8 qualified as LBS
import Data.Map.Strict qualified as Map
import Database.Bloodhound.OpenSearch1.Requests qualified as OS1Requests
import Database.Bloodhound.OpenSearch1.Types qualified as OS1Types
import Database.Bloodhound.OpenSearch2.Requests qualified as OS2Requests
import Database.Bloodhound.OpenSearch2.Types qualified as OS2Types
import Database.Bloodhound.OpenSearch3.Requests qualified as OS3Requests
import Database.Bloodhound.OpenSearch3.Types
import TestsUtils.Import
import Prelude
-- | The canonical ML Commons stats response: cluster-level scalars at the
-- top level (one key per cluster stat) plus a @nodes@ object whose value is
-- a map keyed by opaque node ID. Drawn from the documented example shape in
-- the ML Commons plugin source (RestMLStatsAction + MLNodeLevelStat +
-- MLClusterLevelStat). Stat values are arbitrary nested JSON; we
-- deliberately round-trip them as 'Value' rather than modelling individual
-- stats, since the stat-name set is plugin-version dependent.
sampleFullResponse :: LBS.ByteString
sampleFullResponse =
"{\
\ \"ml_task_index_status\": \"non-existent\",\
\ \"ml_connector_count\": 2,\
\ \"ml_config_index_status\": \"green\",\
\ \"ml_controller_index_status\": \"non-existent\",\
\ \"ml_model_index_status\": \"non-existent\",\
\ \"ml_connector_index_status\": \"non-existent\",\
\ \"ml_model_count\": 5,\
\ \"nodes\": {\
\ \"zbduvgCCSOeu6cfbQhTpnQ\": {\
\ \"ml_executing_task_count\": 0,\
\ \"ml_request_count\": 12,\
\ \"ml_failure_count\": 1,\
\ \"ml_jvm_heap_usage\": 0,\
\ \"ml_deployed_model_count\": 2,\
\ \"ml_circuit_breaker_trigger_count\": 0\
\ },\
\ \"54xOe0w8Qjyze00UuLDfdA\": {\
\ \"ml_executing_task_count\": 1,\
\ \"ml_request_count\": 30,\
\ \"ml_failure_count\": 0,\
\ \"ml_jvm_heap_usage\": 1048576,\
\ \"ml_deployed_model_count\": 2,\
\ \"ml_circuit_breaker_trigger_count\": 0\
\ }\
\ }\
\}"
-- | A nodes-only response, as returned when a node-level @stat@ filter is
-- supplied (e.g. @GET /_plugins/_ml/stats/ml_executing_task_count@). The
-- cluster-level stats are omitted because the filter restricts to node-level
-- values; the decoder must accept their absence.
sampleNodesOnlyResponse :: LBS.ByteString
sampleNodesOnlyResponse =
"{\
\ \"nodes\": {\
\ \"zbduvgCCSOeu6cfbQhTpnQ\": {\
\ \"ml_executing_task_count\": 0\
\ }\
\ }\
\}"
spec :: Spec
spec = describe "ML Stats API" $ do
describe "MLNodeId / MLStatName JSON" $ do
it "MLNodeId round-trips as a bare JSON string" $ do
encode (MLNodeId "node-id-22-chars-000") `shouldBe` "\"node-id-22-chars-000\""
decode "\"abc\"" `shouldBe` Just (MLNodeId "abc")
it "MLStatName round-trips as a bare JSON string" $ do
encode (MLStatName "ml_executing_task_count")
`shouldBe` "\"ml_executing_task_count\""
decode "\"ml_request_count\""
`shouldBe` Just (MLStatName "ml_request_count")
describe "MLStats JSON" $ do
it "decodes a full response with cluster stats and a nodes map" $ do
let Just decoded = decode sampleFullResponse :: Maybe MLStats
entries = mlStatsEntries decoded
-- 7 cluster-level keys + 1 "nodes" key.
Map.size entries `shouldBe` 8
Map.member "nodes" entries `shouldBe` True
Map.member "ml_model_count" entries `shouldBe` True
Map.member "ml_connector_count" entries `shouldBe` True
Map.member "ml_model_index_status" entries `shouldBe` True
it "preserves cluster-level scalars as bare aeson Values" $ do
let Just decoded = decode sampleFullResponse :: Maybe MLStats
entries = mlStatsEntries decoded
Map.lookup "ml_model_count" entries `shouldBe` Just (toJSON (5 :: Int))
Map.lookup "ml_connector_count" entries `shouldBe` Just (toJSON (2 :: Int))
Map.lookup "ml_model_index_status" entries
`shouldBe` Just (toJSON ("non-existent" :: Text))
Map.lookup "ml_config_index_status" entries
`shouldBe` Just (toJSON ("green" :: Text))
it "preserves the per-node map under the \"nodes\" key as an aeson Value" $ do
let Just decoded = decode sampleFullResponse :: Maybe MLStats
Just (Object nodesObj) = Map.lookup "nodes" (mlStatsEntries decoded)
-- The Value's KeyMap should contain both node IDs.
nodeIds = [k | (k, _) <- KM.toList nodesObj]
length nodeIds `shouldBe` 2
"zbduvgCCSOeu6cfbQhTpnQ" `elem` nodeIds `shouldBe` True
"54xOe0w8Qjyze00UuLDfdA" `elem` nodeIds `shouldBe` True
it "decodes the per-node body of a specific node" $ do
let Just decoded = decode sampleFullResponse :: Maybe MLStats
Just (Object nodesObj) = Map.lookup "nodes" (mlStatsEntries decoded)
Just nodeBody = KM.lookup "zbduvgCCSOeu6cfbQhTpnQ" nodesObj
nodeBody
`shouldBe` object
[ "ml_executing_task_count" .= (0 :: Int),
"ml_request_count" .= (12 :: Int),
"ml_failure_count" .= (1 :: Int),
"ml_jvm_heap_usage" .= (0 :: Int),
"ml_deployed_model_count" .= (2 :: Int),
"ml_circuit_breaker_trigger_count" .= (0 :: Int)
]
it "decodes a nodes-only response (cluster stats absent)" $ do
let Just decoded = decode sampleNodesOnlyResponse :: Maybe MLStats
entries = mlStatsEntries decoded
Map.size entries `shouldBe` 1
Map.member "nodes" entries `shouldBe` True
-- None of the cluster-level keys are present.
Map.member "ml_model_count" entries `shouldBe` False
it "decodes an empty object as an empty map" $ do
let Just decoded = decode "{}" :: Maybe MLStats
mlStatsEntries decoded `shouldBe` Map.empty
it "decodes an empty nodes map as an Object with no entries" $ do
let Just decoded = decode "{ \"nodes\": {} }" :: Maybe MLStats
Just (Object nodesObj) = Map.lookup "nodes" (mlStatsEntries decoded)
KM.null nodesObj `shouldBe` True
it "rejects a top-level null" $ do
let decoded = decode "null" :: Maybe MLStats
decoded `shouldBe` Nothing
it "rejects a top-level array" $ do
let decoded = decode "[]" :: Maybe MLStats
decoded `shouldBe` Nothing
it "rejects a top-level scalar" $ do
let decoded = decode "42" :: Maybe MLStats
decoded `shouldBe` Nothing
it "rejects malformed JSON" $ do
let decoded = decode "{ not json" :: Maybe MLStats
decoded `shouldBe` Nothing
-- Lenient contract: the flat-Map decoder does NOT validate the shape of
-- the @nodes@ value (it is just another entry in the map). This test
-- pins the current behaviour so any future tightening of the decoder is
-- a deliberate, breaking change rather than a silent one.
it "accepts a non-object nodes value (lenient contract — pinned)" $ do
let Just decoded = decode "{ \"nodes\": \"unexpected\" }" :: Maybe MLStats
Map.lookup "nodes" (mlStatsEntries decoded)
`shouldBe` Just (toJSON ("unexpected" :: Text))
it "round-trips a full response through ToJSON/FromJSON" $ do
let Just decoded = decode sampleFullResponse :: Maybe MLStats
(decode . encode) decoded `shouldBe` Just decoded
it "round-trips a nodes-only response through ToJSON/FromJSON" $ do
let Just decoded = decode sampleNodesOnlyResponse :: Maybe MLStats
(decode . encode) decoded `shouldBe` Just decoded
it "ToJSON emits both cluster and nodes keys" $ do
let Just decoded = decode sampleFullResponse :: Maybe MLStats
encoded = LBS.toStrict (encode decoded)
encoded `shouldSatisfy` BS.isInfixOf "\"nodes\""
encoded `shouldSatisfy` BS.isInfixOf "\"ml_model_count\""
encoded `shouldSatisfy` BS.isInfixOf "\"zbduvgCCSOeu6cfbQhTpnQ\""
describe "getMLStats endpoint shape (OS3)" $ do
-- Pure checks against the BHRequest shape; no live backend needed.
it "GETs /_plugins/_ml/stats when given Nothing/Nothing" $ do
let req = OS3Requests.getMLStats Nothing Nothing
getRawEndpoint (bhRequestEndpoint req)
`shouldBe` ["_plugins", "_ml", "stats"]
getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []
it "GETs /_plugins/_ml/stats/{stat} when only the stat name is set" $ do
let req =
OS3Requests.getMLStats
Nothing
(Just (MLStatName "ml_executing_task_count"))
getRawEndpoint (bhRequestEndpoint req)
`shouldBe` ["_plugins", "_ml", "stats", "ml_executing_task_count"]
getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []
it "GETs /_plugins/_ml/{nodeId}/stats when only the node id is set" $ do
let req =
OS3Requests.getMLStats
(Just (MLNodeId "node-id-22-chars-000"))
Nothing
getRawEndpoint (bhRequestEndpoint req)
`shouldBe` ["_plugins", "_ml", "node-id-22-chars-000", "stats"]
getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []
it "GETs /_plugins/_ml/{nodeId}/stats/{stat} when both are set" $ do
let req =
OS3Requests.getMLStats
(Just (MLNodeId "node-id-22-chars-000"))
(Just (MLStatName "ml_request_count"))
getRawEndpoint (bhRequestEndpoint req)
`shouldBe` [ "_plugins",
"_ml",
"node-id-22-chars-000",
"stats",
"ml_request_count"
]
getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []
it "uses the GET method" $ do
let req = OS3Requests.getMLStats Nothing Nothing
bhRequestMethod req `shouldBe` "GET"
it "does not attach a body (GET semantics)" $ do
let req = OS3Requests.getMLStats Nothing Nothing
bhRequestBody req `shouldBe` Nothing
describe "getMLStats cross-version parity (OS1/OS2/OS3)" $ do
-- All three backend versions share the same path shape; this guards
-- against accidental module-name typos when back-porting. Each version
-- has its own @MLNodeId@ \/ @MLStatName@ newtype (they are deliberately
-- distinct at the type level even though they share a name and shape),
-- so the per-version values are constructed against the matching module.
let expectedBoth =
[ "_plugins",
"_ml",
"node-id-22-chars-000",
"stats",
"ml_executing_task_count"
]
it "OS1 builds the same /_plugins/_ml/stats path as OS3" $ do
getRawEndpoint
(bhRequestEndpoint (OS1Requests.getMLStats Nothing Nothing))
`shouldBe` ["_plugins", "_ml", "stats"]
it "OS2 builds the same /_plugins/_ml/stats path as OS3" $ do
getRawEndpoint
(bhRequestEndpoint (OS2Requests.getMLStats Nothing Nothing))
`shouldBe` ["_plugins", "_ml", "stats"]
it "OS1 builds the same {nodeId}/stats/{stat} path as OS3" $ do
getRawEndpoint
( bhRequestEndpoint
( OS1Requests.getMLStats
(Just (OS1Types.MLNodeId "node-id-22-chars-000"))
(Just (OS1Types.MLStatName "ml_executing_task_count"))
)
)
`shouldBe` expectedBoth
it "OS2 builds the same {nodeId}/stats/{stat} path as OS3" $ do
getRawEndpoint
( bhRequestEndpoint
( OS2Requests.getMLStats
(Just (OS2Types.MLNodeId "node-id-22-chars-000"))
(Just (OS2Types.MLStatName "ml_executing_task_count"))
)
)
`shouldBe` expectedBoth