SSESpec.hs

{-# LANGUAGE OverloadedStrings #-}

-- | Pure unit tests for the SSE parser
-- ('Database.Bloodhound.Internal.Utils.SSE'). The parser is exercised
-- against literal byte streams copied from the OpenSearch ML Commons
-- Execute Stream Agent documentation (both the conversational-agent
-- @data:@ framing and the AG-UI event framing), so these tests pin the
-- wire shape against future drift. No live cluster is required: a
-- 'Network.HTTP.Client.BodyReader' is synthesised from an in-memory
-- 'ByteString' via a small helper.
module Test.SSESpec (spec) where

import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.IORef (newIORef, readIORef, writeIORef)
import Data.Text (Text)
import Data.Text qualified as T
import Database.Bloodhound.Internal.Utils.SSE
import Network.HTTP.Client (BodyReader)
import Test.Hspec
import Prelude hiding (head)

-- | Build a 'BodyReader' that yields the given bytes in one shot then
-- EOF. Mirrors how @http-client@ delivers a fully-buffered chunk for
-- small responses; sufficient for these parser tests.
mkBodyReader :: ByteString -> IO BodyReader
mkBodyReader bytes = do
  ref <- newIORef bytes
  pure $
    readIORef ref >>= \remaining ->
      if BS.null remaining
        then pure BS.empty
        else do
          writeIORef ref BS.empty
          pure remaining

spec :: Spec
spec = do
  describe "SSE parser" $ do
    it "parses a single data: event" $ do
      br <- mkBodyReader "data: {\"hello\":1}\n\n"
      st <- newSSEState br
      ev <- readSSEEvent st
      sseData <$> ev `shouldBe` Just "{\"hello\":1}"
      sseEvent <$> ev `shouldBe` Just Nothing

    it "parses multiple consecutive events" $ do
      br <- mkBodyReader "data: a\n\ndata: b\n\ndata: c\n\n"
      st <- newSSEState br
      e1 <- readSSEEvent st
      e2 <- readSSEEvent st
      e3 <- readSSEEvent st
      e4 <- readSSEEvent st
      sseData <$> e1 `shouldBe` Just "a"
      sseData <$> e2 `shouldBe` Just "b"
      sseData <$> e3 `shouldBe` Just "c"
      e4 `shouldBe` Nothing

    it "joins multi-line data: fields with \\n (per spec)" $ do
      br <- mkBodyReader "data: line1\ndata: line2\n\n"
      st <- newSSEState br
      ev <- readSSEEvent st
      sseData <$> ev `shouldBe` Just "line1\nline2"

    it "strips a single leading space after the colon" $ do
      br <- mkBodyReader "data:    padded\n\n"
      st <- newSSEState br
      ev <- readSSEEvent st
      -- Only the FIRST space is stripped; the rest are preserved.
      sseData <$> ev `shouldBe` Just "   padded"

    it "tolerates CRLF line endings" $ do
      br <- mkBodyReader "data: crlf\r\n\r\n"
      st <- newSSEState br
      ev <- readSSEEvent st
      sseData <$> ev `shouldBe` Just "crlf"

    it "ignores comment lines starting with :" $ do
      br <- mkBodyReader ": this is a comment\ndata: real\n\n"
      st <- newSSEState br
      ev <- readSSEEvent st
      sseData <$> ev `shouldBe` Just "real"

    it "parses event:/id:/retry: fields" $ do
      br <- mkBodyReader "event: add\ndata: x\nid: 7\nretry: 3000\n\n"
      st <- newSSEState br
      ev <- readSSEEvent st
      ev
        `shouldBe` Just
          (SSEEvent {sseEvent = Just "add", sseData = "x", sseId = Just "7", sseRetry = Just 3000})

    it "handles incremental byte delivery (chunks smaller than a frame)" $ do
      -- A BodyReader that yields 3 bytes at a time, simulating a slow
      -- network stream. The parser must reassemble across reads.
      ref <- newIORef ("data: hello\n\n" :: ByteString)
      let br :: BodyReader
          br = do
            remaining <- readIORef ref
            let (chunk, rest) = BS.splitAt 3 remaining
            writeIORef ref rest
            pure chunk
      st <- newSSEState br
      ev <- readSSEEvent st
      sseData <$> ev `shouldBe` Just "hello"

    it "returns Nothing at end of stream" $ do
      br <- mkBodyReader ""
      st <- newSSEState br
      ev <- readSSEEvent st
      ev `shouldBe` Nothing

    it "dispatches a trailing partial event at EOF" $ do
      br <- mkBodyReader "data: tail-no-separator"
      st <- newSSEState br
      ev <- readSSEEvent st
      sseData <$> ev `shouldBe` Just "tail-no-separator"

    it "parses a realistic conversational-agent chunk sequence" $ do
      let stream :: ByteString =
            "data: {\"inference_results\":[{\"output\":[{\"name\":\"memory_id\",\"result\":\"m1\"}]}]}\n\n\
            \data: {\"inference_results\":[{\"output\":[{\"name\":\"response\",\"dataAsMap\":{\"content\":\"There\",\"is_last\":false}}]}]}\n\n\
            \data: {\"inference_results\":[{\"output\":[{\"name\":\"response\",\"dataAsMap\":{\"content\":\" are\",\"is_last\":false}}]}]}\n\n\
            \data: {\"inference_results\":[{\"output\":[{\"name\":\"response\",\"dataAsMap\":{\"content\":\"\",\"is_last\":true}}]}]}\n\n"
      br <- mkBodyReader stream
      st <- newSSEState br
      chunks <- collect st 10
      length chunks `shouldBe` 4
      case chunks of
        firstEv : _ -> sseData firstEv `shouldSatisfy` T.isInfixOf "memory_id"
        [] -> expectationFailure "expected at least one chunk"

-- | Collect up to @n@ events from the parser, then stop.
collect :: SSEState -> Int -> IO [SSEEvent]
collect _ 0 = pure []
collect st n = do
  next <- readSSEEvent st
  case next of
    Just ev -> (ev :) <$> collect st (n - 1)
    Nothing -> pure []