baikai-openai-0.6.0.0: test/TransportSpec.hs

module TransportSpec (tests) where

import Baikai
import Baikai.Provider.OpenAI.Api (openaiChatStream)
import Baikai.Provider.OpenAI.Transport qualified as Transport
import Control.Concurrent (threadDelay)
import Control.Exception (bracket, try)
import Control.Lens ((&), (.~), (^.))
import Control.Monad (forM_)
import Data.CaseInsensitive qualified as CI
import Data.IORef (newIORef, readIORef, writeIORef)
import Data.Map.Strict qualified as Map
import Data.Text qualified as Text
import Data.Text.Encoding qualified as Text
import Network.HTTP.Types.Header (RequestHeaders)
import Servant.Client qualified as Client
import Streamly.Data.Stream qualified as Stream
import System.Environment (lookupEnv, setEnv, unsetEnv)
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=))

tests :: TestTree
tests =
  testGroup
    "Baikai.Provider.OpenAI.Transport"
    [ clientEnvCacheTest,
      requestHeadersTest,
      timeoutTest,
      nonPositiveTimeoutTest,
      unknownHostKeyTest,
      unusableBaseUrlTest
    ]

-- | One entry per target, and one notion of what a target is.
--
-- Both halves are asserted in a single case because the cache is
-- process-global and this suite runs in parallel: two cases each reading
-- a count and expecting it to move by exactly one would race each other.
--
-- The normalisation half is what makes the count meaningful. The key is
-- the canonical rendering of "Baikai.Url"'s parse rather than the
-- caller's text, so a trailing slash and a capitalised host do not each
-- open their own connection pool to the same host — and the two provider
-- packages, which now share one cache in @Baikai.Http@, cannot disagree
-- about which target a URL names.
clientEnvCacheTest :: TestTree
clientEnvCacheTest =
  testCase "the ClientEnv cache allocates once per normalised base URL" $ do
    let url = "https://cache-openai.test"
    before <- Transport.cachedClientEnvCount
    _ <- Transport.getClientEnvCached url
    afterFirst <- Transport.cachedClientEnvCount
    _ <- Transport.getClientEnvCached url
    afterSecond <- Transport.cachedClientEnvCount
    afterFirst @?= before + 1
    afterSecond @?= afterFirst
    -- A different spelling of the same target: capitalised host,
    -- trailing slash.
    env <- Transport.getClientEnvCached "https://Cache-openai.test/"
    afterVariant <- Transport.cachedClientEnvCount
    afterVariant @?= afterSecond
    Client.baseUrlHost (Client.baseUrl env) @?= "cache-openai.test"
    Client.baseUrlPath (Client.baseUrl env) @?= ""

requestHeadersTest :: TestTree
requestHeadersTest =
  testCase "model and option headers reach the wire, options winning case-insensitively" $ do
    let model =
          emptyModel
            & #headers .~ Map.fromList [("X-Trace", "model"), ("authorization", "model-auth")]
        opts =
          emptyOptions
            & #headers .~ Map.fromList [("x-trace", "option"), ("Authorization", "option-auth")]
        headers = Transport.requestHeaders "secret" model opts
    header "X-Trace" headers @?= Just "option"
    header "authorization" headers @?= Just "option-auth"
    header "Accept" headers @?= Just "text/event-stream"

timeoutTest :: TestTree
timeoutTest =
  testCase "runWithTimeout classifies an elapsed whole-call timeout as transient" $ do
    result <- Transport.runWithTimeout (Just 1) (threadDelay 100000)
    case result of
      Just be -> do
        be ^. #category @?= TransientError
        "timeoutMs=1" `Text.isInfixOf` (be ^. #message) @?= True
      Nothing -> assertFailure "expected timeout error"

nonPositiveTimeoutTest :: TestTree
nonPositiveTimeoutTest =
  testCase "runWithTimeout rejects a non-positive bound without running the action" $ do
    -- System.Timeout.timeout returns immediately at zero and runs
    -- unbounded below it, so both spellings used to fail instantly as a
    -- retryable TransientError, which a retry loop re-issues forever for
    -- what is a caller-side mistake.
    forM_ [0, -5] $ \ms -> do
      ran <- newIORef False
      result <- Transport.runWithTimeout (Just ms) (writeIORef ran True)
      case result of
        Just be -> do
          be ^. #category @?= InvalidRequest
          isRetryable be @?= False
        Nothing -> assertFailure ("expected an InvalidRequest for timeoutMs=" <> show ms)
      readIORef ran >>= (@?= False)

unknownHostKeyTest :: TestTree
unknownHostKeyTest =
  testCase "unknown hosts do not fall back to OPENAI_API_KEY" $
    withEnv "OPENAI_API_KEY" "openai-secret" $ do
      result <- try (Transport.resolveKey "https://unknown.example" emptyOptions) :: IO (Either BaikaiError Text.Text)
      case result of
        Left be -> be ^. #category @?= AuthError
        Right _ -> assertFailure "expected AuthError for unknown host"

header :: Text.Text -> RequestHeaders -> Maybe Text.Text
header name headers =
  Text.decodeUtf8 <$> lookup (CI.mk (Text.encodeUtf8 name)) headers

withEnv :: String -> String -> IO a -> IO a
withEnv name value =
  bracket
    (lookupEnv name <* setEnv name value)
    (maybe (unsetEnv name) (setEnv name))
    . const

-- | A base URL baikai will not send to is refused before a key is read.
--
-- The order matters as much as the refusal. These cases run with the
-- provider's own key variable *unset*, so an AuthError would prove the
-- check ran too late; an InvalidRequest proves nothing was looked up.
-- The messages also have to say what is wrong without echoing the part
-- of the URL that could be a credential.
unusableBaseUrlTest :: TestTree
unusableBaseUrlTest =
  testCase "an unusable base URL is refused before any key is read"
    $ withoutEnv "OPENAI_OpenAIChatCompletions_KEY"
    $ forM_
      [ ("https://h.test/v1?api-version=2024-01", "query string"),
        ("https://u:pw@h.test", "credentials"),
        ("h.test", "https://"),
        ("https://h.test/v1/chat/completions", "endpoint path")
      ]
    $ \(url, needle) -> do
      let model = emptyModel & #api .~ OpenAIChatCompletions & #baseUrl .~ url
      events <- Stream.toList (openaiChatStream model emptyContext emptyOptions)
      case events of
        [EventStart _, EventError payload] -> case payload ^. #errorInfo of
          Nothing -> assertFailure (Text.unpack url <> ": the error carried no errorInfo")
          Just err -> do
            let message = err ^. #message
            (url, err ^. #category) @?= (url, InvalidRequest)
            assertBool
              (Text.unpack (url <> " should name the problem: " <> message))
              (needle `Text.isInfixOf` message)
            assertBool
              (Text.unpack (url <> " must not echo the query: " <> message))
              (not ("api-version=2024-01" `Text.isInfixOf` message))
            assertBool
              (Text.unpack (url <> " must not echo the password: " <> message))
              (not ("pw@" `Text.isInfixOf` message))
        other ->
          assertFailure
            (Text.unpack url <> ": expected [EventStart, EventError], got: " <> show other)

withoutEnv :: String -> IO a -> IO a
withoutEnv name =
  bracket
    (lookupEnv name <* unsetEnv name)
    (maybe (unsetEnv name) (setEnv name))
    . const