IntegrationTest.hs

module Shibuya.Adapter.Kafka.IntegrationTest (tests) where

import Control.Exception (throwIO)
import Control.Monad (forM)
import Control.Monad.IO.Class (liftIO)
import Data.ByteString (ByteString)
import Data.ByteString.Char8 qualified as BS8
import Data.IORef (modifyIORef', newIORef, readIORef, writeIORef)
import Data.List (nub, sort)
import Data.Maybe (mapMaybe)
import Data.Text qualified as Text
import Effectful (runEff)
import Effectful.Error.Static (runError)
import Kafka.Consumer.Types (OffsetCommit (..), OffsetReset (..))
import Kafka.Effectful.Consumer (
    brokersList,
    groupId,
    noAutoOffsetStore,
    offsetReset,
    runKafkaConsumer,
    topics,
 )
import Kafka.Effectful.Consumer.Effect (commitAllOffsets, pollMessageBatch)
import Kafka.TestEnv (
    TestEnv (..),
    consumeN,
    createTopic,
    createTopicWithPartitions,
    produceKeyedMessages,
    produceMessages,
    withTestEnv,
 )
import Kafka.Types (
    BatchSize (..),
    KafkaError,
    Timeout (..),
    TopicName (..),
 )
import Shibuya.Adapter (Adapter (..))
import Shibuya.Adapter.Kafka (KafkaAdapterConfig (..), kafkaAdapter)
import Shibuya.App (ProcessorId (..), defaultAppConfig, mkProcessor, runApp, waitApp)
import Shibuya.Core.Ack (AckDecision (..), RetryDelay (..))
import Shibuya.Core.AckHandle (AckHandle (..))
import Shibuya.Core.Ingested (Ingested (..), Message (..))
import Shibuya.Core.Types (Cursor (..), Envelope (..), MessageId (..))
import Shibuya.Telemetry.Effect (runTracingNoop)
import Streamly.Data.Fold qualified as Fold
import Streamly.Data.Stream qualified as Stream
import System.Timeout (timeout)
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (assertBool, assertEqual, assertFailure, testCase)

tests :: TestTree
tests =
    testGroup
        "Integration"
        [ testCase "Basic produce-consume" testBasicProduceConsume
        , testCase "Offset commit verification" testOffsetCommit
        , testCase "Multi-partition distribution" testMultiPartition
        , testCase "Batch polling" testBatchPolling
        , testCase "Graceful shutdown" testGracefulShutdown
        , testCase "Idle graceful shutdown completes promptly" testIdleGracefulShutdown
        , testCase "AckRetry redelivers within the same session" testAckRetryRedelivery
        , testCase "AckRetry is not committed past when session exits" testAckRetryAbandonedSession
        , testCase "Handler exception redelivers instead of skipping" testHandlerExceptionRedelivery
        ]

testBasicProduceConsume :: IO ()
testBasicProduceConsume = withTestEnv $ \env -> do
    createTopic env
    let payloads = ["msg-1", "msg-2", "msg-3", "msg-4", "msg-5"]
    produceMessages env payloads
    envelopes <- consumeN env 5 AckOk

    -- Verify all 5 messages received
    assertEqual "message count" 5 (length envelopes)

    -- Verify payloads
    let receivedPayloads = map (\(Envelope{payload}) -> payload) envelopes
    assertEqual "payloads" (map Just payloads) receivedPayloads

    -- Verify messageId format: topic-partition-offset
    case envelopes of
        (Envelope{messageId = MessageId firstIdText} : _) ->
            assertBool "messageId contains topic" (Text.isPrefixOf (unTopicName env.testTopic) firstIdText)
        [] -> error "unreachable: already verified 5 envelopes"

    -- Verify cursor is populated
    assertBool "cursor is Just" (all (\(Envelope{cursor}) -> case cursor of Just (CursorInt _) -> True; _ -> False) envelopes)

    -- Verify partition is populated
    assertBool "partition is Just" (all (\(Envelope{partition}) -> case partition of Just _ -> True; _ -> False) envelopes)

testOffsetCommit :: IO ()
testOffsetCommit = withTestEnv $ \env -> do
    createTopic env
    let payloads = ["oc-1", "oc-2", "oc-3"]
    produceMessages env payloads

    -- Consume all 3, AckOk each (stores offsets), then commit
    _ <- consumeN env 3 AckOk

    -- Create new consumer in same group - should get no messages
    result <- runEff . runError @KafkaError $ do
        let props = brokersList [env.testBroker] <> groupId env.testGroupId <> noAutoOffsetStore
            sub = topics [env.testTopic] <> offsetReset Earliest
        runKafkaConsumer props sub $ do
            -- Poll a few times to allow group join + rebalance, verify no re-delivery
            allResults <- forM [1 .. 3 :: Int] $ \_ -> do
                results <- pollMessageBatch (Timeout 3000) (BatchSize 100)
                pure [cr | Right cr <- results]
            let totalMessages = concat allResults
            liftIO $ assertEqual "no re-delivery" 0 (length totalMessages)
    case result of
        Left err -> error $ "Failed: " <> show err
        Right () -> pure ()

testMultiPartition :: IO ()
testMultiPartition = withTestEnv $ \env -> do
    createTopicWithPartitions env 3
    let pairs =
            [ ("key-a", "msg-a")
            , ("key-b", "msg-b")
            , ("key-c", "msg-c")
            , ("key-d", "msg-d")
            , ("key-e", "msg-e")
            , ("key-f", "msg-f")
            ]
    produceKeyedMessages env pairs
    envelopes <- consumeN env 6 AckOk

    assertEqual "message count" 6 (length envelopes)

    let partitions = mapMaybe (\(Envelope{partition}) -> partition) envelopes
    assertEqual "all have partition" 6 (length partitions)

    let uniquePartitions = nub partitions
    assertBool
        ("expected multiple partitions, got: " <> show uniquePartitions)
        (length uniquePartitions >= 2)

testBatchPolling :: IO ()
testBatchPolling = withTestEnv $ \env -> do
    createTopic env
    let payloads = map (\i -> BS8.pack ("batch-" <> show i)) [1 .. 20 :: Int]
    produceMessages env payloads
    envelopes <- consumeN env 20 AckOk

    assertEqual "message count" 20 (length envelopes)

    let receivedPayloads = sort $ mapMaybe (\(Envelope{payload}) -> payload) envelopes
    assertEqual "payloads" (sort payloads) receivedPayloads

testGracefulShutdown :: IO ()
testGracefulShutdown = withTestEnv $ \env -> do
    createTopic env
    let payloads = ["sd-1", "sd-2", "sd-3"]
    produceMessages env payloads

    ref <- newIORef ([] :: [Envelope (Maybe ByteString)])
    result <- runEff . runError @KafkaError $ do
        let props = brokersList [env.testBroker] <> groupId env.testGroupId <> noAutoOffsetStore
            sub = topics [env.testTopic] <> offsetReset Earliest
        runKafkaConsumer props sub $ do
            let config =
                    KafkaAdapterConfig
                        { topics = [env.testTopic]
                        , pollTimeout = Timeout 5000
                        , batchSize = BatchSize 100
                        }
            Adapter{source, shutdown} <- kafkaAdapter config
            Stream.fold Fold.drain
                $ Stream.mapM
                    ( \(Ingested{envelope, ack = AckHandle finalize}) -> do
                        liftIO $ modifyIORef' ref (envelope :)
                        finalize AckOk
                    )
                $ Stream.take 3 source
            shutdown
    case result of
        Left err -> error $ "Failed: " <> show err
        Right () -> pure ()
    envelopes <- reverse <$> readIORef ref
    assertEqual "consumed 3 before shutdown" 3 (length envelopes)

testIdleGracefulShutdown :: IO ()
testIdleGracefulShutdown = withTestEnv $ \env -> do
    createTopic env

    timedResult <- timeout 3000000 $ runEff . runError @KafkaError $ do
        let props = brokersList [env.testBroker] <> groupId env.testGroupId <> noAutoOffsetStore
            sub = topics [env.testTopic] <> offsetReset Earliest
        runKafkaConsumer props sub $ do
            let config =
                    KafkaAdapterConfig
                        { topics = [env.testTopic]
                        , pollTimeout = Timeout 250
                        , batchSize = BatchSize 100
                        }
            Adapter{source, shutdown} <- kafkaAdapter config
            shutdown
            Stream.fold Fold.drain source
    case timedResult of
        Nothing -> assertFailure "idle shutdown did not terminate promptly"
        Just (Left (_cs, err)) -> assertFailure $ "idle shutdown failed: " <> show err
        Just (Right ()) -> pure ()

testAckRetryRedelivery :: IO ()
testAckRetryRedelivery = withTestEnv $ \env -> do
    createTopic env
    let payloads = ["r-1", "r-2", "r-3"]
    produceMessages env payloads

    retried <- newIORef False
    seen <- newIORef ([] :: [ByteString])
    result <- runEff . runError @KafkaError $ do
        let props = brokersList [env.testBroker] <> groupId env.testGroupId <> noAutoOffsetStore
            sub = topics [env.testTopic] <> offsetReset Earliest
        runKafkaConsumer props sub $ do
            Adapter{source} <- kafkaAdapter (testConfig env)
            Stream.fold Fold.drain
                $ Stream.mapM
                    ( \(Ingested{envelope, ack = AckHandle finalize}) -> do
                        let payload = maybe "" id envelope.payload
                        liftIO $ modifyIORef' seen (payload :)
                        hasRetried <- liftIO $ readIORef retried
                        if payload == "r-2" && not hasRetried
                            then do
                                liftIO $ writeIORef retried True
                                finalize (AckRetry (RetryDelay 0))
                            else finalize AckOk
                    )
                $ Stream.take 4 source
            commitAllOffsets OffsetCommit
    case result of
        Left (_cs, err) -> assertFailure $ "AckRetry redelivery failed: " <> show err
        Right () -> pure ()

    delivered <- reverse <$> readIORef seen
    assertBool ("expected r-2 at least twice, saw " <> show delivered) (countPayload "r-2" delivered >= 2)
    assertBool ("expected r-3 after retry, saw " <> show delivered) ("r-3" `elem` delivered)

    noRedelivery <- runEff . runError @KafkaError $ do
        let props = brokersList [env.testBroker] <> groupId env.testGroupId <> noAutoOffsetStore
            sub = topics [env.testTopic] <> offsetReset Earliest
        runKafkaConsumer props sub $ do
            batches <- forM [1 .. 3 :: Int] $ \_ ->
                pollMessageBatch (Timeout 500) (BatchSize 100)
            liftIO $ assertEqual "no redelivery after final AckOk commit" 0 (length [cr | Right cr <- concat batches])
    case noRedelivery of
        Left (_cs, err) -> assertFailure $ "post-commit verification failed: " <> show err
        Right () -> pure ()

testAckRetryAbandonedSession :: IO ()
testAckRetryAbandonedSession = withTestEnv $ \env -> do
    createTopic env
    let payloads = ["ab-1", "ab-2", "ab-3"]
    produceMessages env payloads

    firstSession <- runEff . runError @KafkaError $ do
        let props = brokersList [env.testBroker] <> groupId env.testGroupId <> noAutoOffsetStore
            sub = topics [env.testTopic] <> offsetReset Earliest
        runKafkaConsumer props sub $ do
            Adapter{source} <- kafkaAdapter (testConfig env)
            Stream.fold Fold.drain
                $ Stream.mapM
                    ( \(Ingested{envelope, ack = AckHandle finalize}) -> do
                        case envelope.payload of
                            Just "ab-2" -> finalize (AckRetry (RetryDelay 0))
                            _ -> finalize AckOk
                    )
                $ Stream.take 2 source
    case firstSession of
        Left (_cs, err) -> assertFailure $ "first session failed: " <> show err
        Right () -> pure ()

    redelivered <- newIORef ([] :: [ByteString])
    secondSession <- runEff . runError @KafkaError $ do
        let props = brokersList [env.testBroker] <> groupId env.testGroupId <> noAutoOffsetStore
            sub = topics [env.testTopic] <> offsetReset Earliest
        runKafkaConsumer props sub $ do
            Adapter{source} <- kafkaAdapter (testConfig env)
            Stream.fold Fold.drain
                $ Stream.mapM
                    ( \(Ingested{envelope, ack = AckHandle finalize}) -> do
                        maybe (pure ()) (liftIO . modifyIORef' redelivered . (:)) envelope.payload
                        finalize AckOk
                    )
                $ Stream.take 2 source
            commitAllOffsets OffsetCommit
    case secondSession of
        Left (_cs, err) -> assertFailure $ "second session failed: " <> show err
        Right () -> pure ()

    delivered <- readIORef redelivered
    assertBool ("expected ab-2 redelivery, saw " <> show delivered) ("ab-2" `elem` delivered)

testHandlerExceptionRedelivery :: IO ()
testHandlerExceptionRedelivery = withTestEnv $ \env -> do
    createTopic env
    let payloads = ["n-1", "n-2", "n-3"]
    produceMessages env payloads

    thrown <- newIORef False
    seen <- newIORef ([] :: [ByteString])
    result <- runEff . runError @KafkaError . runTracingNoop $ do
        let props = brokersList [env.testBroker] <> groupId env.testGroupId <> noAutoOffsetStore
            sub = topics [env.testTopic] <> offsetReset Earliest
        runKafkaConsumer props sub $ do
            upstream <- kafkaAdapter (testConfig env)
            let finiteAdapter = upstream{source = Stream.take 4 upstream.source}
                handler Message{envelope} = do
                    let payload = maybe "" id envelope.payload
                    liftIO $ modifyIORef' seen (payload :)
                    hasThrown <- liftIO $ readIORef thrown
                    if payload == "n-2" && not hasThrown
                        then liftIO $ do
                            writeIORef thrown True
                            throwIO (userError "planned handler exception")
                        else pure AckOk
            appResult <- runApp defaultAppConfig [(ProcessorId "handler-exception", mkProcessor finiteAdapter handler)]
            case appResult of
                Left appErr -> liftIO $ assertFailure $ "runApp failed: " <> show appErr
                Right appHandle -> waitApp appHandle
            commitAllOffsets OffsetCommit
    case result of
        Left (_cs, err) -> assertFailure $ "handler exception scenario failed: " <> show err
        Right () -> pure ()

    delivered <- reverse <$> readIORef seen
    assertBool ("expected n-2 at least twice, saw " <> show delivered) (countPayload "n-2" delivered >= 2)
    assertBool ("expected n-3 after handler exception retry, saw " <> show delivered) ("n-3" `elem` delivered)

testConfig :: TestEnv -> KafkaAdapterConfig
testConfig env =
    KafkaAdapterConfig
        { topics = [env.testTopic]
        , pollTimeout = Timeout 500
        , batchSize = BatchSize 100
        }

countPayload :: ByteString -> [ByteString] -> Int
countPayload target = length . filter (== target)