74 lines
2.4 KiB
Python
74 lines
2.4 KiB
Python
|
import pytest
|
||
|
|
from service.sse import (
|
||
|
|
format_sse_event, create_token_event, create_tool_start_event,
|
||
|
|
create_tool_result_event, create_tool_error_event,
|
||
|
|
create_error_event
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def test_format_sse_event():
|
||
|
|
"""Test SSE event formatting"""
|
||
|
|
event = format_sse_event("test", {"message": "hello"})
|
||
|
|
expected = 'event: test\ndata: {"message": "hello"}\n\n'
|
||
|
|
assert event == expected
|
||
|
|
|
||
|
|
|
||
|
|
def test_create_token_event():
|
||
|
|
"""Test token event creation"""
|
||
|
|
event = create_token_event("hello", "tool_123")
|
||
|
|
assert "event: tokens" in event
|
||
|
|
assert '"delta": "hello"' in event
|
||
|
|
assert '"tool_call_id": "tool_123"' in event
|
||
|
|
|
||
|
|
|
||
|
|
def test_create_token_event_no_tool_id():
|
||
|
|
"""Test token event without tool call ID"""
|
||
|
|
event = create_token_event("hello")
|
||
|
|
assert "event: tokens" in event
|
||
|
|
assert '"delta": "hello"' in event
|
||
|
|
assert '"tool_call_id": null' in event
|
||
|
|
|
||
|
|
|
||
|
|
def test_create_tool_start_event():
|
||
|
|
"""Test tool start event"""
|
||
|
|
event = create_tool_start_event("tool_123", "retrieve_standard_regulation", {"query": "test"})
|
||
|
|
assert "event: tool_start" in event
|
||
|
|
assert '"id": "tool_123"' in event
|
||
|
|
assert '"name": "retrieve_standard_regulation"' in event
|
||
|
|
assert '"args": {"query": "test"}' in event
|
||
|
|
|
||
|
|
|
||
|
|
def test_create_tool_result_event():
|
||
|
|
"""Test tool result event"""
|
||
|
|
results = [{"id": "1", "title": "Test Standard"}]
|
||
|
|
event = create_tool_result_event("tool_123", "retrieve_standard_regulation", results, 500)
|
||
|
|
assert "event: tool_result" in event
|
||
|
|
assert '"id": "tool_123"' in event
|
||
|
|
assert '"took_ms": 500' in event
|
||
|
|
assert '"results"' in event
|
||
|
|
|
||
|
|
|
||
|
|
def test_create_tool_error_event():
|
||
|
|
"""Test tool error event"""
|
||
|
|
event = create_tool_error_event("tool_123", "retrieve_standard_regulation", "API timeout")
|
||
|
|
assert "event: tool_error" in event
|
||
|
|
assert '"id": "tool_123"' in event
|
||
|
|
assert '"error": "API timeout"' in event
|
||
|
|
|
||
|
|
|
||
|
|
def test_create_error_event():
|
||
|
|
"""Test error event"""
|
||
|
|
event = create_error_event("Something went wrong")
|
||
|
|
assert "event: error" in event
|
||
|
|
assert '"error": "Something went wrong"' in event
|
||
|
|
|
||
|
|
|
||
|
|
def test_create_error_event_with_details():
|
||
|
|
"""Test error event with details"""
|
||
|
|
details = {"code": 500, "source": "llm"}
|
||
|
|
event = create_error_event("Something went wrong", details)
|
||
|
|
assert "event: error" in event
|
||
|
|
assert '"error": "Something went wrong"' in event
|
||
|
|
assert '"details"' in event
|
||
|
|
assert '"code": 500' in event
|