MLTaskSpec.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.MLTaskSpec (spec) where

import Data.ByteString.Char8 qualified as BS
import Data.ByteString.Lazy.Char8 qualified as LBS
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

-- | A representative response body for @GET /_plugins/_ml/tasks/{task_id}@
-- after a successful model registration. Drawn from the shape documented at
-- <https://docs.opensearch.org/latest/ml-commons-plugin/api/model-apis/get-task/>.
sampleRegisterCompleted :: LBS.ByteString
sampleRegisterCompleted =
  "{\
  \  \"task_id\": \"ew8I44MBhyWuIwnfvDIH\",\
  \  \"model_id\": \"d-MCYeIBUTn_O1uORDWuDw\",\
  \  \"task_type\": \"REGISTER_MODEL\",\
  \  \"function_name\": \"TEXT_EMBEDDING\",\
  \  \"state\": \"COMPLETED\",\
  \  \"worker_node\": [\"KzwtsZ7ySdeg1GqYuISPjg\"],\
  \  \"create_time\": 1704110400000,\
  \  \"last_update_time\": 1704110460000,\
  \  \"is_async\": true\
  \}"

-- | A @CREATED@ task that has not started running yet: no @last_update_time@,
-- no @worker_node@. Exercises the all-'Maybe' record's tolerance for the
-- sparse fields the plugin emits early in the lifecycle.
sampleCreatedSparse :: LBS.ByteString
sampleCreatedSparse =
  "{\
  \  \"task_id\": \"ew8I44MBhyWuIwnfvDIH\",\
  \  \"model_id\": \"d-MCYeIBUTn_O1uORDWuDw\",\
  \  \"task_type\": \"DEPLOY_MODEL\",\
  \  \"state\": \"CREATED\",\
  \  \"create_time\": 1704110400000,\
  \  \"is_async\": true\
  \}"

-- | A @FAILED@ task carries an @error@ string instead of a @worker_node@
-- list. Exercises the @error@ field plus an unrelated field (@model_id@)
-- being absent.
sampleFailed :: LBS.ByteString
sampleFailed =
  "{\
  \  \"task_type\": \"DEPLOY_MODEL\",\
  \  \"state\": \"FAILED\",\
  \  \"error\": \"node memory constraints prevented deployment\"\
  \}"

-- | OS 1.x emits @worker_node@ as a bare string rather than an array. The
-- 'FromJSON' instance must coerce this into a one-element list so the same
-- client works against every backend version.
sampleWorkerNodeString :: LBS.ByteString
sampleWorkerNodeString =
  "{\
  \  \"task_type\": \"REGISTER_MODEL\",\
  \  \"state\": \"RUNNING\",\
  \  \"worker_node\": \"KzwtsZ7ySdeg1GqYuISPjg\"\
  \}"

-- | A response with @state: NOT_FOUND@ — the plugin emits this when the
-- supplied task id has been cleaned up. No @task_type@, no @model_id@.
sampleNotFound :: LBS.ByteString
sampleNotFound =
  "{\
  \  \"state\": \"NOT_FOUND\"\
  \}"

spec :: Spec
spec = describe "ML Commons Get Task API" $ do
  describe "MLTaskId JSON" $ do
    it "MLTaskId round-trips as a bare JSON string" $ do
      encode (MLTaskId "task-xyz-123") `shouldBe` "\"task-xyz-123\""
      decode "\"abc\"" `shouldBe` Just (MLTaskId "abc")

  describe "MLTaskType JSON" $ do
    it "decodes known task_type values" $ do
      decode "\"REGISTER_MODEL\"" `shouldBe` Just MLTaskTypeRegisterModel
      decode "\"DEPLOY_MODEL\"" `shouldBe` Just MLTaskTypeDeployModel
      decode "\"TRAINING\"" `shouldBe` Just MLTaskTypeTraining
      decode "\"REGISTER_REMOTE_MODEL\"" `shouldBe` Just MLTaskTypeRegisterRemoteModel
      decode "\"DEPLOY_REMOTE_MODEL\"" `shouldBe` Just MLTaskTypeDeployRemoteModel

    it "round-trips known values through ToJSON/FromJSON" $ do
      encode MLTaskTypeRegisterModel `shouldBe` "\"REGISTER_MODEL\""
      encode MLTaskTypeDeployModel `shouldBe` "\"DEPLOY_MODEL\""

    it "decodes unknown values to MLTaskTypeOther without failing" $ do
      let Just decoded = decode "\"SOMETHING_NEW_IN_OS4\"" :: Maybe MLTaskType
      decoded `shouldBe` MLTaskTypeOther "SOMETHING_NEW_IN_OS4"
      (decode . encode) decoded `shouldBe` Just decoded

    it "rejects non-string task_type values" $ do
      (decode "42" :: Maybe MLTaskType) `shouldBe` Nothing
      (decode "{}" :: Maybe MLTaskType) `shouldBe` Nothing

  describe "MLTaskState JSON" $ do
    it "decodes known state values" $ do
      decode "\"CREATED\"" `shouldBe` Just MLTaskStateCreated
      decode "\"RUNNING\"" `shouldBe` Just MLTaskStateRunning
      decode "\"COMPLETED\"" `shouldBe` Just MLTaskStateCompleted
      decode "\"FAILED\"" `shouldBe` Just MLTaskStateFailed
      decode "\"NOT_FOUND\"" `shouldBe` Just MLTaskStateNotFound

    it "round-trips known values through ToJSON/FromJSON" $ do
      encode MLTaskStateCreated `shouldBe` "\"CREATED\""
      encode MLTaskStateFailed `shouldBe` "\"FAILED\""

    it "decodes unknown values to MLTaskStateOther without failing" $ do
      let Just decoded = decode "\"CANCELLED\"" :: Maybe MLTaskState
      decoded `shouldBe` MLTaskStateOther "CANCELLED"
      (decode . encode) decoded `shouldBe` Just decoded

    it "rejects non-string state values" $ do
      (decode "42" :: Maybe MLTaskState) `shouldBe` Nothing
      (decode "[]" :: Maybe MLTaskState) `shouldBe` Nothing

  describe "MLTaskInfo decode" $ do
    it "decodes a fully-populated REGISTER_MODEL/COMPLETED response" $ do
      let Just decoded = decode sampleRegisterCompleted :: Maybe MLTaskInfo
      mlTaskInfoTaskId decoded `shouldBe` Just "ew8I44MBhyWuIwnfvDIH"
      mlTaskInfoModelId decoded `shouldBe` Just "d-MCYeIBUTn_O1uORDWuDw"
      mlTaskInfoTaskType decoded `shouldBe` Just MLTaskTypeRegisterModel
      mlTaskInfoFunctionName decoded `shouldBe` Just "TEXT_EMBEDDING"
      mlTaskInfoState decoded `shouldBe` Just MLTaskStateCompleted
      mlTaskInfoWorkerNode decoded `shouldBe` Just ["KzwtsZ7ySdeg1GqYuISPjg"]
      mlTaskInfoCreateTime decoded `shouldBe` Just 1704110400000
      mlTaskInfoLastUpdateTime decoded `shouldBe` Just 1704110460000
      mlTaskInfoIsAsync decoded `shouldBe` Just True
      mlTaskInfoError decoded `shouldBe` Nothing

    it "decodes a sparse CREATED response (absent fields become Nothing)" $ do
      let Just decoded = decode sampleCreatedSparse :: Maybe MLTaskInfo
      mlTaskInfoTaskId decoded `shouldBe` Just "ew8I44MBhyWuIwnfvDIH"
      mlTaskInfoCreateTime decoded `shouldBe` Just 1704110400000
      mlTaskInfoLastUpdateTime decoded `shouldBe` Nothing
      mlTaskInfoWorkerNode decoded `shouldBe` Nothing
      mlTaskInfoError decoded `shouldBe` Nothing
      mlTaskInfoState decoded `shouldBe` Just MLTaskStateCreated

    it "decodes a FAILED response carrying an error string" $ do
      let Just decoded = decode sampleFailed :: Maybe MLTaskInfo
      mlTaskInfoState decoded `shouldBe` Just MLTaskStateFailed
      mlTaskInfoError decoded `shouldBe` Just "node memory constraints prevented deployment"
      mlTaskInfoModelId decoded `shouldBe` Nothing

    it "coerces a single-string worker_node (OS 1.x shape) to a one-element list" $ do
      let Just decoded = decode sampleWorkerNodeString :: Maybe MLTaskInfo
      mlTaskInfoWorkerNode decoded `shouldBe` Just ["KzwtsZ7ySdeg1GqYuISPjg"]

    -- The tolerant worker_node parser silently shapes unexpected inputs into
    -- the closest safe value rather than failing the record decode. These
    -- three tests pin the chosen behaviour so a future tightening (or
    -- loosening) of the parser is a deliberate decision rather than an
    -- accidental side effect of a refactor.
    it "decodes an empty worker_node array as Just [] (empty list, not Nothing)" $ do
      let Just decoded = decode "{\"worker_node\":[]}" :: Maybe MLTaskInfo
      mlTaskInfoWorkerNode decoded `shouldBe` Just ([] :: [Text])

    it "silently drops non-String elements inside a worker_node array" $ do
      -- OpenSearch never emits this, but the parser uses a list comprehension
      -- @String t <- toList xs@ that filters rather than fails. Pin the
      -- behaviour: the malformed element disappears and the rest survive.
      let Just decoded = decode "{\"worker_node\":[\"ok\", 42, \"also-ok\"]}" :: Maybe MLTaskInfo
      mlTaskInfoWorkerNode decoded `shouldBe` Just ["ok", "also-ok"]

    it "coerces an unexpected worker_node shape (object) to Nothing" $ do
      -- The catch-all @Just _ -> Nothing@ arm. If OpenSearch ever ships a
      -- third shape, the field is silently absent rather than failing the
      -- whole decode.
      let Just decoded = decode "{\"worker_node\":{\"id\":\"x\"}}" :: Maybe MLTaskInfo
      mlTaskInfoWorkerNode decoded `shouldBe` Nothing

    it "decodes a NOT_FOUND response with only the state field" $ do
      let Just decoded = decode sampleNotFound :: Maybe MLTaskInfo
      mlTaskInfoState decoded `shouldBe` Just MLTaskStateNotFound
      mlTaskInfoTaskType decoded `shouldBe` Nothing

    it "decodes an empty object as all-Nothing fields" $ do
      let Just decoded = decode "{}" :: Maybe MLTaskInfo
      mlTaskInfoModelId decoded `shouldBe` Nothing
      mlTaskInfoState decoded `shouldBe` Nothing

    it "rejects a top-level array" $ do
      (decode "[]" :: Maybe MLTaskInfo) `shouldBe` Nothing

    it "rejects a field with the wrong type (state as a number)" $ do
      let decoded = decode "{\"state\": 42}" :: Maybe MLTaskInfo
      decoded `shouldBe` Nothing

    it "rejects malformed JSON" $ do
      (decode "{ not json" :: Maybe MLTaskInfo) `shouldBe` Nothing

  describe "MLTaskInfo ToJSON round-trip" $ do
    it "round-trips a fully-populated body through ToJSON/FromJSON" $ do
      let Just decoded = decode sampleRegisterCompleted :: Maybe MLTaskInfo
      (decode . encode) decoded `shouldBe` Just decoded

    it "round-trips a sparse CREATED body through ToJSON/FromJSON" $ do
      let Just decoded = decode sampleCreatedSparse :: Maybe MLTaskInfo
      (decode . encode) decoded `shouldBe` Just decoded

    it "ToJSON omits Nothing fields (no null leakage)" $ do
      let Just decoded = decode sampleCreatedSparse :: Maybe MLTaskInfo
          encoded = LBS.toStrict (encode decoded)
      encoded `shouldSatisfy` not . BS.isInfixOf "\"error\""
      encoded `shouldSatisfy` not . BS.isInfixOf "\"worker_node\""
      encoded `shouldSatisfy` BS.isInfixOf "\"state\""

  describe "getMLTask endpoint shape (OS1)" $ do
    it "GETs /_plugins/_ml/tasks/{task_id}" $ do
      let req = OS1Requests.getMLTask (OS1Types.MLTaskId "task-abc-123")
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_ml", "tasks", "task-abc-123"]
      getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []
    it "uses the GET method" $ do
      let req = OS1Requests.getMLTask (OS1Types.MLTaskId "task-abc-123")
      bhRequestMethod req `shouldBe` "GET"
    it "does not attach a body" $ do
      let req = OS1Requests.getMLTask (OS1Types.MLTaskId "task-abc-123")
      bhRequestBody req `shouldBe` Nothing

  describe "getMLTask endpoint shape (OS2)" $ do
    it "GETs /_plugins/_ml/tasks/{task_id}" $ do
      let req = OS2Requests.getMLTask (OS2Types.MLTaskId "task-abc-123")
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_ml", "tasks", "task-abc-123"]
      getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []
    it "uses the GET method" $ do
      let req = OS2Requests.getMLTask (OS2Types.MLTaskId "task-abc-123")
      bhRequestMethod req `shouldBe` "GET"
    it "does not attach a body" $ do
      let req = OS2Requests.getMLTask (OS2Types.MLTaskId "task-abc-123")
      bhRequestBody req `shouldBe` Nothing

  describe "getMLTask endpoint shape (OS3)" $ do
    it "GETs /_plugins/_ml/tasks/{task_id}" $ do
      let req = OS3Requests.getMLTask (MLTaskId "task-abc-123")
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_ml", "tasks", "task-abc-123"]
      getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []
    it "uses the GET method" $ do
      let req = OS3Requests.getMLTask (MLTaskId "task-abc-123")
      bhRequestMethod req `shouldBe` "GET"
    it "does not attach a body" $ do
      let req = OS3Requests.getMLTask (MLTaskId "task-abc-123")
      bhRequestBody req `shouldBe` Nothing