init

vw-agentic-rag/tests/unit/test_sse.py

@@ -0,0 +1,73 @@
+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