content: {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeFamilies #-}
  
  module Test.Langchain.Runnable.Core (tests) where
  
  import Data.IORef (modifyIORef, newIORef, readIORef, writeIORef)
  import Langchain.Error (LangchainError, llmError)
  import Langchain.Runnable.Core
  import Test.Tasty (TestTree, testGroup)
  import Test.Tasty.HUnit (assertEqual, testCase)
  
  newtype MockRunnable a b = MockRunnable
    { runMock :: a -> IO (Either LangchainError b)
    }
  
  instance Runnable (MockRunnable a b) where
    type RunnableInput (MockRunnable a b) = a
    type RunnableOutput (MockRunnable a b) = b
    invoke = runMock
  
  tests :: TestTree
  tests =
    testGroup
      "Runnable Tests"
      [ testCase "invoke success" $ do
          let mock = MockRunnable (\(s :: String) -> return $ Right (s ++ " processed"))
          result <- invoke mock "input"
          assertEqual "Should process input" (Right "input processed") result
      , testCase "invoke error" $ do
          let mock = MockRunnable (\(_ :: String) -> return $ Left (llmError "mock error" Nothing Nothing))
          result <- invoke mock "input"
          assertEqual
            "Should return error"
            (Left (llmError "mock error" Nothing Nothing) :: Either LangchainError String)
            result
      , testCase "batch success" $ do
          let mock = MockRunnable (\(s :: String) -> return $ Right (s ++ "!"))
          result <- batch mock ["a", "b", "c"]
          assertEqual "All inputs processed" (Right ["a!", "b!", "c!"]) result
      , testCase "batch with error" $ do
          let mock = MockRunnable $ \(s :: String) ->
                if s == "b"
                  then return (Left (llmError "error in batch" Nothing Nothing))
                  else return (Right (s ++ "!"))
          result <- batch mock ["a", "b", "c"]
          assertEqual "Should return first error" (Left (llmError "error in batch" Nothing Nothing)) result
      , testCase "stream success" $ do
          ref <- newIORef []
          let mock = MockRunnable (\(s :: String) -> return $ Right (s ++ "!"))
              callback x = modifyIORef ref (++ [x])
          result <- stream mock "test" callback
          readRef <- readIORef ref
          assertEqual "Stream should succeed" (Right ()) result
          assertEqual "Callback called with correct value" ["test!"] readRef
      , testCase "stream error" $ do
          ref <- newIORef []
          let mock = MockRunnable (\(_ :: String) -> return $ Left (llmError "stream error" Nothing Nothing))
              callback _ = writeIORef ref ["should not be called" :: String]
          result <- stream mock "test" callback
          readRef <- readIORef ref
          assertEqual "Stream should return error" (Left (llmError "stream error" Nothing Nothing)) result
          assertEqual "Callback not called" [] readRef
      ]
  