//! 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 { Some(SseItem::Data(s.to_string())) } /// Hands out one byte per call. struct OneByte(Cursor>); impl Read for OneByte { fn read(&mut self, buf: &mut [u8]) -> std::io::Result { let end = buf.len().min(1); self.0.read(&mut buf[..end]) } } /// Serves `bytes`, then fails every read with `kind`. struct ThenError(Cursor>, std::io::ErrorKind); impl Read for ThenError { fn read(&mut self, buf: &mut [u8]) -> std::io::Result { 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 = 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}"); } } } }