MLControllersSpec.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.MLControllersSpec (spec) where

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

-- | Create controller response: @{model_id, status}@.
sampleCreateResponse :: LBS.ByteString
sampleCreateResponse = "{\"model_id\":\"1\",\"status\":\"CREATED\"}"

-- | Get controller response. Note @limit@ comes back as a JSON string
-- even though it is sent as a number.
sampleGetResponse :: LBS.ByteString
sampleGetResponse =
  "{\
  \  \"model_id\": \"1\",\
  \  \"user_rate_limiter\": {\
  \    \"alice\": {\"limit\": \"10\", \"unit\": \"MINUTES\"},\
  \    \"bob\": {\"limit\": 5, \"unit\": \"HOURS\"}\
  \  }\
  \}"

spec :: Spec
spec = describe "ML Commons controller APIs" $ do
  describe "ControllerCreateResponse" $
    it "decodes the doc sample" $
      decode sampleCreateResponse
        `shouldBe` Just
          ControllerCreateResponse
            { controllerCreateResponseModelId = Just "1",
              controllerCreateResponseStatus = Just "CREATED"
            }

  describe "ControllerConfig" $ do
    it "decodes the doc get sample (limit as string AND number; reuses ModelRateLimiter)" $ do
      let decoded = decode sampleGetResponse :: Maybe ControllerConfig
      case decoded of
        Just cfg -> do
          controllerConfigModelId cfg `shouldBe` Just "1"
          -- The per-user map should contain both alice and bob entries,
          -- tolerating the stringified-number limit quirk.
          controllerConfigUserRateLimiter cfg `shouldSatisfy` (\m -> length m == 2)
        Nothing -> expectationFailure "expected ControllerConfig decode"
    it "round-trips a config built from a single user rate limiter" $ do
      let cfg =
            ControllerConfig
              { controllerConfigModelId = Nothing,
                controllerConfigUserRateLimiter =
                  M.fromList
                    [ ( "alice",
                        ModelRateLimiter
                          { modelRateLimiterLimit = Just (ModelRateLimit 10),
                            modelRateLimiterUnit = Just ModelRateLimiterUnitMinutes
                          }
                      )
                    ]
              }
      decode (encode cfg) `shouldBe` Just cfg