boxmaker/crates/loopd/tests/sse.rs

//! Tests for the SSE reader. Do not edit.

mod support;

use loopd::http::{Connection, Request};
use loopd::sse::{Events, SseError, SseItem};
use std::io::{Cursor, Read};
use std::time::Duration;
use support::{FakeServer, Reply};

fn data(s: &str) -> Option<SseItem> {
    Some(SseItem::Data(s.to_string()))
}

/// Hands out one byte per call.
struct OneByte(Cursor<Vec<u8>>);

impl Read for OneByte {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        let end = buf.len().min(1);
        self.0.read(&mut buf[..end])
    }
}

/// Serves `bytes`, then fails every read with `kind`.
struct ThenError(Cursor<Vec<u8>>, std::io::ErrorKind);

impl Read for ThenError {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        match self.0.read(buf)? {
            0 => Err(std::io::Error::from(self.1)),
            n => Ok(n),
        }
    }
}

const STREAM: &str =
    "data: {\"a\":1}\n\n: a comment\ndata:{\"b\":2}\r\n\r\nevent: x\nid: 7\n\ndata: [DONE]\n\n";

#[test]
fn yields_data_lines_and_done() {
    let mut ev = Events::new(Cursor::new(STREAM), 1024);
    assert_eq!(ev.next_item().unwrap(), data("{\"a\":1}"));
    assert_eq!(
        ev.next_item().unwrap(),
        data("{\"b\":2}"),
        "no space after the colon, and CRLF"
    );
    assert_eq!(ev.next_item().unwrap(), Some(SseItem::Done));
    assert_eq!(ev.next_item().unwrap(), None);
    assert_eq!(
        ev.next_item().unwrap(),
        None,
        "asking again after the end is harmless"
    );
}

#[test]
fn the_result_does_not_depend_on_how_the_bytes_arrive() {
    let mut ev = Events::new(OneByte(Cursor::new(STREAM.as_bytes().to_vec())), 1024);
    assert_eq!(ev.next_item().unwrap(), data("{\"a\":1}"));
    assert_eq!(ev.next_item().unwrap(), data("{\"b\":2}"));
    assert_eq!(ev.next_item().unwrap(), Some(SseItem::Done));
    assert_eq!(ev.next_item().unwrap(), None);
}

#[test]
fn only_a_leading_space_is_stripped_and_multibyte_text_survives() {
    let mut ev = Events::new(
        OneByte(Cursor::new(
            "data:  two spaces\ndata: caf\u{e9} \u{1f4e6}\n".into(),
        )),
        1024,
    );
    assert_eq!(ev.next_item().unwrap(), data(" two spaces"));
    assert_eq!(ev.next_item().unwrap(), data("caf\u{e9} \u{1f4e6}"));
}

#[test]
fn a_stream_that_ends_mid_line_is_truncated() {
    let mut ev = Events::new(Cursor::new("data: {\"a\":1}\n\ndata: {\"b\""), 1024);
    assert_eq!(ev.next_item().unwrap(), data("{\"a\":1}"));
    assert!(matches!(ev.next_item(), Err(SseError::Truncated)));
}

#[test]
fn reader_errors_are_classified() {
    use std::io::ErrorKind::{ConnectionReset, TimedOut, UnexpectedEof, WouldBlock};
    let first = "data: 1\n";
    for (kind, want) in [
        (WouldBlock, "timeout"),
        (TimedOut, "timeout"),
        (UnexpectedEof, "truncated"),
        (ConnectionReset, "io"),
    ] {
        let mut ev = Events::new(ThenError(Cursor::new(first.into()), kind), 1024);
        assert_eq!(ev.next_item().unwrap(), data("1"));
        let got = match ev.next_item() {
            Err(SseError::Timeout) => "timeout",
            Err(SseError::Truncated) => "truncated",
            Err(SseError::Io(_)) => "io",
            other => panic!("{kind:?}: {other:?}"),
        };
        assert_eq!(got, want, "{kind:?}");
    }
}

#[test]
fn long_lines_and_bad_text_are_errors() {
    let long = format!("data: {}\n", "x".repeat(5000));
    let mut ev = Events::new(Cursor::new(long.clone()), 100);
    assert!(matches!(ev.next_item(), Err(SseError::LineTooLong)));
    let mut ev = Events::new(OneByte(Cursor::new("x".repeat(5000).into())), 100);
    assert!(
        matches!(ev.next_item(), Err(SseError::LineTooLong)),
        "with no newline at all"
    );
    let mut ev = Events::new(Cursor::new(long), 10_000);
    assert!(
        matches!(ev.next_item(), Ok(Some(SseItem::Data(_)))),
        "under the limit it is fine"
    );

    let mut ev = Events::new(Cursor::new(b"data: \xff\xfe\n".to_vec()), 1024);
    assert!(matches!(ev.next_item(), Err(SseError::NotUtf8)));
    let e: Box<dyn std::error::Error> = Box::new(SseError::Truncated);
    assert!(!e.to_string().is_empty());
}

/// The recorded streams, read through the HTTP client as `loopd` will read them.
#[test]
fn recorded_streams_have_the_expected_shape() {
    for (name, want_data) in [("plain", 21), ("tool_call", 27), ("turn1", 6), ("turn2", 7)] {
        for piece in [usize::MAX, 1, 13] {
            let server = FakeServer::start();
            server.route(
                "/v1/chat/completions",
                vec![Reply::fixture(name).trickle(piece, 0)],
            );
            let mut conn = Connection::open(&server.socket).unwrap();
            conn.set_read_timeout(Duration::from_secs(5)).unwrap();
            conn.send(&Request {
                method: "POST",
                path: "/v1/chat/completions",
                body: Some(b"{}"),
            })
            .unwrap();
            let head = conn.read_head().unwrap();
            let mut ev = Events::new(conn.body(&head).unwrap(), 1024 * 1024);
            let mut items = Vec::new();
            while let Some(item) = ev.next_item().unwrap() {
                items.push(item);
            }
            assert_eq!(items.last(), Some(&SseItem::Done), "{name}");
            assert_eq!(items.len() - 1, want_data, "{name}: data items");
            for item in &items[..items.len() - 1] {
                let SseItem::Data(text) = item else {
                    panic!("{name}: Done in the middle")
                };
                let json: serde_json::Value = serde_json::from_str(text).unwrap();
                assert!(json["choices"].is_array(), "{name}");
            }
        }
    }
}