"""Tests for AG-UI event creation utilities.""" from pydantic import BaseModel from haiku.rag.graph.agui.events import ( emit_activity, emit_run_error, emit_run_finished, emit_run_started, emit_state_snapshot, emit_step_finished, emit_step_started, emit_text_message, ) class TestState(BaseModel): """Test state model.""" value: int class TestResult(BaseModel): """Test result model.""" status: str def test_emit_run_started(): """Test RUN_STARTED event creation.""" event = emit_run_started("thread-1", "run-1") assert event["type"] == "RUN_STARTED" assert event["threadId"] == "thread-1" assert event["runId"] == "run-1" assert "input" not in event def test_emit_run_started_with_input(): """Test RUN_STARTED event with input data.""" event = emit_run_started("thread-1", "run-1", input_data="test input") assert event["type"] == "RUN_STARTED" assert event["input"] == "test input" def test_emit_run_finished(): """Test RUN_FINISHED event creation.""" result = TestResult(status="complete") event = emit_run_finished("thread-1", "run-1", result) assert event["type"] == "RUN_FINISHED" assert event["threadId"] == "thread-1" assert event["runId"] == "run-1" assert event["result"] == {"status": "complete"} def test_emit_run_finished_with_dict(): """Test RUN_FINISHED event with dict result.""" result = {"status": "complete", "count": 42} event = emit_run_finished("thread-1", "run-1", result) assert event["type"] == "RUN_FINISHED" assert event["result"] == result def test_emit_run_error(): """Test RUN_ERROR event creation.""" event = emit_run_error("Something went wrong") assert event["type"] == "RUN_ERROR" assert event["message"] == "Something went wrong" assert "code" not in event def test_emit_run_error_with_code(): """Test RUN_ERROR event with error code.""" event = emit_run_error("Something went wrong", code="ERR_001") assert event["type"] == "RUN_ERROR" assert event["message"] == "Something went wrong" assert event["code"] == "ERR_001" def test_emit_step_started(): """Test STEP_STARTED event creation.""" event = emit_step_started("plan") assert event["type"] == "STEP_STARTED" assert event["stepName"] == "plan" def test_emit_step_finished(): """Test STEP_FINISHED event creation.""" event = emit_step_finished("plan") assert event["type"] == "STEP_FINISHED" assert event["stepName"] == "plan" def test_emit_text_message(): """Test TEXT_MESSAGE_CHUNK event creation.""" event = emit_text_message("Hello world") assert event["type"] == "TEXT_MESSAGE_CHUNK" assert event["delta"] == "Hello world" assert event["role"] == "assistant" assert "messageId" in event def test_emit_text_message_with_role(): """Test TEXT_MESSAGE_CHUNK event with custom role.""" event = emit_text_message("Hello", role="user") assert event["type"] == "TEXT_MESSAGE_CHUNK" assert event["role"] == "user" def test_emit_state_snapshot(): """Test STATE_SNAPSHOT event creation.""" state = TestState(value=42) event = emit_state_snapshot(state) assert event["type"] == "STATE_SNAPSHOT" assert event["snapshot"] == {"value": 42} def test_emit_activity(): """Test ACTIVITY_SNAPSHOT event creation.""" event = emit_activity("msg-1", "processing", "Working on task") assert event["type"] == "ACTIVITY_SNAPSHOT" assert event["messageId"] == "msg-1" assert event["activityType"] == "processing" assert event["content"] == "Working on task" def test_event_structure_consistency(): """Test that all events have consistent structure.""" events = [ emit_run_started("t1", "r1"), emit_run_finished("t1", "r1", {"result": "done"}), emit_run_error("error"), emit_step_started("step1"), emit_step_finished("step1"), emit_text_message("text"), emit_state_snapshot(TestState(value=1)), emit_activity("m1", "type", "content"), ] for event in events: assert isinstance(event, dict) assert "type" in event assert isinstance(event["type"], str) assert event["type"].isupper() # Event types are uppercase