stream-rs 0.1.0

//! Tests for the `OpenAI`, Anthropic, and Gemini delta accumulators.

use stream_rs::accumulators::anthropic::{AnthropicAccumulator, BlockKind};
use stream_rs::accumulators::gemini::GeminiAccumulator;
use stream_rs::accumulators::openai::OpenAiAccumulator;
use stream_rs::error::AccumulateError;

#[test]
fn openai_accumulates_content_and_role() {
    let mut acc = OpenAiAccumulator::new();
    acc.push_role(0, "assistant");
    acc.push_content(0, "Hel");
    acc.push_content(0, "lo, ");
    acc.push_content(0, "world");
    acc.set_finish_reason(0, "stop");

    let c = acc.choice(0).unwrap();
    assert_eq!(c.role.as_deref(), Some("assistant"));
    assert_eq!(c.content, "Hello, world");
    assert_eq!(c.finish_reason.as_deref(), Some("stop"));
}

#[test]
fn openai_role_only_set_once() {
    let mut acc = OpenAiAccumulator::new();
    acc.push_role(0, "assistant");
    acc.push_role(0, "system");
    assert_eq!(acc.choice(0).unwrap().role.as_deref(), Some("assistant"));
}

#[test]
fn openai_multiple_choices() {
    let mut acc = OpenAiAccumulator::new();
    acc.push_content(0, "a");
    acc.push_content(2, "c");

    let collected: Vec<_> = acc.choices().map(|(i, c)| (i, c.content.clone())).collect();
    assert_eq!(collected, vec![(0, "a".to_string()), (2, "c".to_string())]);
    assert!(acc.choice(1).is_none());
}

#[test]
fn openai_tool_call_assembly() {
    let mut acc = OpenAiAccumulator::new();
    acc.push_tool_call(0, 0, Some("call_1"), Some("get_weather"), Some("{\"loc"));
    acc.push_tool_call(0, 0, None, None, Some("ation\":\"NYC\"}"));

    let tc = &acc.choice(0).unwrap().tool_calls[&0];
    assert_eq!(tc.id.as_deref(), Some("call_1"));
    assert_eq!(tc.name.as_deref(), Some("get_weather"));
    assert_eq!(tc.arguments, "{\"location\":\"NYC\"}");
}

#[test]
fn anthropic_text_block_assembly() {
    let mut acc = AnthropicAccumulator::new();
    acc.message_start();
    acc.content_block_start(0, BlockKind::Text).unwrap();
    acc.text_delta(0, "Hel").unwrap();
    acc.text_delta(0, "lo").unwrap();
    acc.content_block_stop(0).unwrap();
    acc.message_delta(Some("end_turn"));
    acc.message_stop();

    let b = acc.block(0).unwrap();
    assert_eq!(b.kind, BlockKind::Text);
    assert_eq!(b.text, "Hello");
    assert!(b.stopped);
    assert_eq!(acc.stop_reason(), Some("end_turn"));
}

#[test]
fn anthropic_tool_use_json_assembly() {
    let mut acc = AnthropicAccumulator::new();
    acc.message_start();
    acc.content_block_start(0, BlockKind::ToolUse).unwrap();
    acc.input_json_delta(0, "{\"x\":").unwrap();
    acc.input_json_delta(0, "1}").unwrap();
    acc.content_block_stop(0).unwrap();

    assert_eq!(acc.block(0).unwrap().text, "{\"x\":1}");
    assert_eq!(acc.block(0).unwrap().kind, BlockKind::ToolUse);
}

#[test]
fn anthropic_delta_before_start_is_rejected() {
    let mut acc = AnthropicAccumulator::new();
    acc.message_start();
    let err = acc.text_delta(0, "oops").unwrap_err();
    assert!(matches!(err, AccumulateError::UnexpectedEvent { .. }));
}

#[test]
fn anthropic_stop_before_start_is_rejected() {
    let mut acc = AnthropicAccumulator::new();
    let err = acc.content_block_stop(3).unwrap_err();
    assert!(matches!(err, AccumulateError::UnexpectedEvent { .. }));
}

#[test]
fn anthropic_block_event_before_message_start_is_rejected() {
    let mut acc = AnthropicAccumulator::new();
    let err = acc.content_block_start(0, BlockKind::Text).unwrap_err();
    assert!(matches!(err, AccumulateError::UnexpectedEvent { .. }));
}

#[test]
fn anthropic_block_event_after_message_stop_is_rejected() {
    let mut acc = AnthropicAccumulator::new();
    acc.message_start();
    acc.content_block_start(0, BlockKind::Text).unwrap();
    acc.content_block_stop(0).unwrap();
    acc.message_stop();
    let err = acc.content_block_start(1, BlockKind::Text).unwrap_err();
    assert!(matches!(err, AccumulateError::UnexpectedEvent { .. }));
}

#[test]
fn anthropic_duplicate_open_start_is_rejected() {
    let mut acc = AnthropicAccumulator::new();
    acc.message_start();
    acc.content_block_start(0, BlockKind::Text).unwrap();
    // Starting the same still-open index again is illegal.
    let err = acc.content_block_start(0, BlockKind::Text).unwrap_err();
    assert!(matches!(err, AccumulateError::UnexpectedEvent { .. }));
    // After stopping it, restarting the index is allowed again.
    acc.content_block_stop(0).unwrap();
    acc.content_block_start(0, BlockKind::ToolUse).unwrap();
    assert_eq!(acc.block(0).unwrap().kind, BlockKind::ToolUse);
}

#[test]
fn anthropic_delta_kind_mismatch_is_rejected() {
    let mut acc = AnthropicAccumulator::new();
    acc.message_start();
    acc.content_block_start(0, BlockKind::ToolUse).unwrap();
    // A text_delta into a tool_use block is a kind mismatch.
    let err = acc.text_delta(0, "nope").unwrap_err();
    assert!(matches!(
        err,
        AccumulateError::BlockKindMismatch {
            index: 0,
            expected: "text",
            actual: "tool_use"
        }
    ));
}

#[test]
fn gemini_accumulates_text_and_finish_reason() {
    let mut acc = GeminiAccumulator::new();
    acc.push_text(0, "The weather");
    acc.push_text(0, " in ");
    acc.push_text(0, "Paris");
    acc.set_finish_reason(0, "STOP");

    let c = acc.candidate(0).unwrap();
    assert_eq!(c.text, "The weather in Paris");
    assert_eq!(c.finish_reason.as_deref(), Some("STOP"));
}

#[test]
fn gemini_function_calls_recorded_in_order() {
    let mut acc = GeminiAccumulator::new();
    acc.push_function_call(0, "get_weather", r#"{"city":"Paris"}"#);
    acc.push_function_call(0, "get_time", r#"{"tz":"UTC"}"#);

    let calls = &acc.candidate(0).unwrap().function_calls;
    assert_eq!(calls.len(), 2);
    assert_eq!(calls[0].name, "get_weather");
    assert_eq!(calls[0].args, r#"{"city":"Paris"}"#);
    assert_eq!(calls[1].name, "get_time");
    assert_eq!(calls[1].args, r#"{"tz":"UTC"}"#);
}

#[test]
fn gemini_multiple_candidates() {
    let mut acc = GeminiAccumulator::new();
    acc.push_text(0, "a");
    acc.push_text(2, "c");

    let collected: Vec<_> = acc.candidates().map(|(i, c)| (i, c.text.clone())).collect();
    assert_eq!(collected, vec![(0, "a".to_string()), (2, "c".to_string())]);
    assert!(acc.candidate(1).is_none());
}