MLModelGroupsSpec.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.MLModelGroupsSpec (spec) where

import Data.ByteString.Lazy.Char8 qualified as LBS
import Database.Bloodhound.OpenSearch3.Types
import TestsUtils.Import
import Prelude

-- | Register response: @{model_group_id, status}@.
sampleRegisterResponse :: LBS.ByteString
sampleRegisterResponse = "{\"model_group_id\":\"GDNmQ4gBYW0Qyy5ZcBcg\",\"status\":\"CREATED\"}"

-- | Get response. Note the response key is @access@ (not @access_mode@).
sampleGetResponse :: LBS.ByteString
sampleGetResponse =
  "{\
  \  \"name\": \"test_model_group\",\
  \  \"latest_version\": 0,\
  \  \"description\": \"This is a public model group\",\
  \  \"access\": \"public\",\
  \  \"created_time\": 1715112992748,\
  \  \"last_updated_time\": 1715112992748\
  \}"

spec :: Spec
spec = describe "ML Commons model group APIs" $ do
  describe "ModelGroupAccessMode" $ do
    it "encodes the documented access modes" $ do
      encode ModelGroupAccessModePublic `shouldBe` "\"public\""
      encode ModelGroupAccessModePrivate `shouldBe` "\"private\""
      encode ModelGroupAccessModeRestricted `shouldBe` "\"restricted\""
    it "decodes the documented access modes" $ do
      decode "\"public\"" `shouldBe` Just ModelGroupAccessModePublic
      decode "\"private\"" `shouldBe` Just ModelGroupAccessModePrivate
      decode "\"restricted\"" `shouldBe` Just ModelGroupAccessModeRestricted
    it "round-trips an unknown mode through the Other escape hatch" $ do
      let other = ModelGroupAccessModeOther "future"
      decode (encode other) `shouldBe` Just other

  describe "RegisterModelGroupResponse" $
    it "decodes the doc sample" $
      decode sampleRegisterResponse
        `shouldBe` Just
          RegisterModelGroupResponse
            { registerModelGroupResponseId = Just "GDNmQ4gBYW0Qyy5ZcBcg",
              registerModelGroupResponseStatus = Just "CREATED"
            }

  describe "RegisterModelGroupRequest" $
    it "encodes a public group, omitting the absent access-control fields" $ do
      let req =
            RegisterModelGroupRequest
              { registerModelGroupName = "test_model_group",
                registerModelGroupDescription = Just "This is a public model group",
                registerModelGroupAccessMode = Just ModelGroupAccessModePublic,
                registerModelGroupBackendRoles = Nothing,
                registerModelGroupAddAllBackendRoles = Nothing
              }
      encode req
        `shouldBe` "{\"access_mode\":\"public\",\"description\":\"This is a public model group\",\"name\":\"test_model_group\"}"

  describe "ModelGroupInfo" $ do
    it "decodes the doc get sample (response key is 'access')" $
      decode sampleGetResponse
        `shouldBe` Just
          ModelGroupInfo
            { modelGroupName = Just "test_model_group",
              modelGroupLatestVersion = Just 0,
              modelGroupDescription = Just "This is a public model group",
              modelGroupAccess = Just ModelGroupAccessModePublic,
              modelGroupBackendRoles = Nothing,
              modelGroupCreatedTime = Just 1715112992748,
              modelGroupLastUpdatedTime = Just 1715112992748
            }
    it "round-trips a sparse info" $ do
      let info =
            ModelGroupInfo
              { modelGroupName = Just "g",
                modelGroupLatestVersion = Nothing,
                modelGroupDescription = Nothing,
                modelGroupAccess = Just ModelGroupAccessModeRestricted,
                modelGroupBackendRoles = Just ["ops"],
                modelGroupCreatedTime = Nothing,
                modelGroupLastUpdatedTime = Nothing
              }
      decode (encode info) `shouldBe` Just info