# Copyright (c) Microsoft. All rights reserved. """EventStream assertion helper for AG-UI regression tests.""" from __future__ import annotations from typing import Any class EventStream: """Wraps a list of AG-UI events with structured assertion methods. Usage: events = [event async for event in agent.run(payload)] stream = EventStream(events) stream.assert_bookends() stream.assert_text_messages_balanced() """ def __init__(self, events: list[Any]) -> None: self.events = events def __len__(self) -> int: return len(self.events) def __iter__(self): return iter(self.events) def types(self) -> list[str]: """Return ordered list of event type strings.""" return [self._type_str(e) for e in self.events] def get(self, event_type: str) -> list[Any]: """Filter events matching the given type string.""" return [e for e in self.events if self._type_str(e) == event_type] def first(self, event_type: str) -> Any: """Return the first event matching the given type, or raise.""" matches = self.get(event_type) if not matches: raise ValueError(f"No event of type {event_type!r} found. Available: {self.types()}") return matches[0] def last(self, event_type: str) -> Any: """Return the last event matching the given type, or raise.""" matches = self.get(event_type) if not matches: raise ValueError(f"No event of type {event_type!r} found. Available: {self.types()}") return matches[-1] def snapshot(self) -> dict[str, Any]: """Return the latest StateSnapshotEvent snapshot dict.""" return self.last("STATE_SNAPSHOT").snapshot def messages_snapshot(self) -> list[Any]: """Return the latest MessagesSnapshotEvent messages list.""" return self.last("MESSAGES_SNAPSHOT").messages def run_finished_interrupts(self, event: Any | None = None) -> list[dict[str, Any]]: """Return canonical interrupts from a RUN_FINISHED event.""" target = event or self.last("RUN_FINISHED") dumped = self._event_dump(target) assert "interrupt" not in dumped outcome = dumped.get("outcome") assert isinstance(outcome, dict), f"Expected RUN_FINISHED.outcome, got {dumped}" assert outcome.get("type") == "interrupt" interrupts = outcome.get("interrupts") assert isinstance(interrupts, list), f"Expected outcome.interrupts, got {outcome}" return interrupts @staticmethod def interrupt_metadata_value(interrupt: dict[str, Any]) -> dict[str, Any]: """Return Agent Framework interruption details from canonical interrupt metadata.""" metadata = interrupt.get("metadata") assert isinstance(metadata, dict) agent_framework_metadata = metadata.get("agent_framework") assert isinstance(agent_framework_metadata, dict) value = agent_framework_metadata.get("value") assert isinstance(value, dict) return value # ── Structural assertions ── def assert_bookends(self) -> None: """Assert first event is RUN_STARTED and last is RUN_FINISHED.""" types = self.types() assert types, "Event stream is empty" assert types[0] == "RUN_STARTED", f"Expected RUN_STARTED first, got {types[0]}" assert types[-1] == "RUN_FINISHED", f"Expected RUN_FINISHED last, got {types[-1]}" def assert_has_run_lifecycle(self) -> None: """Assert RUN_STARTED is first and RUN_FINISHED exists (may not be last). Use this instead of assert_bookends() for workflow resume streams where _drain_open_message() can emit TEXT_MESSAGE_END after RUN_FINISHED. """ types = self.types() assert types, "Event stream is empty" assert types[0] == "RUN_STARTED", f"Expected RUN_STARTED first, got {types[0]}" assert "RUN_FINISHED" in types, f"Expected RUN_FINISHED in stream. Types: {types}" def assert_strict_types(self, expected: list[str]) -> None: """Assert exact type sequence match.""" actual = self.types() assert actual == expected, f"Event type mismatch.\nExpected: {expected}\nActual: {actual}" def assert_ordered_types(self, expected: list[str]) -> None: """Assert expected types appear as a subsequence (in order, not necessarily contiguous).""" actual = self.types() actual_idx = 0 for expected_type in expected: found = False while actual_idx < len(actual): if actual[actual_idx] == expected_type: actual_idx += 1 found = True break actual_idx += 1 if not found: raise AssertionError( f"Expected subsequence type {expected_type!r} not found after index {actual_idx}.\n" f"Expected subsequence: {expected}\n" f"Actual types: {actual}" ) def assert_text_messages_balanced(self) -> None: """Assert every TEXT_MESSAGE_START has a matching TEXT_MESSAGE_END with the same message_id.""" starts: dict[str, int] = {} ends: set[str] = set() for i, event in enumerate(self.events): t = self._type_str(event) if t == "TEXT_MESSAGE_START": mid = event.message_id assert mid not in starts, f"Duplicate TEXT_MESSAGE_START for message_id={mid}" starts[mid] = i elif t == "TEXT_MESSAGE_END": mid = event.message_id assert mid in starts, f"TEXT_MESSAGE_END for unknown message_id={mid}" assert mid not in ends, f"Duplicate TEXT_MESSAGE_END for message_id={mid}" ends.add(mid) unclosed = set(starts.keys()) - ends assert not unclosed, f"Unclosed text messages: {unclosed}" def assert_tool_calls_balanced(self) -> None: """Assert every TOOL_CALL_START has a matching TOOL_CALL_END with the same tool_call_id.""" starts: dict[str, int] = {} ends: set[str] = set() for i, event in enumerate(self.events): t = self._type_str(event) if t == "TOOL_CALL_START": tid = event.tool_call_id assert tid not in starts, f"Duplicate TOOL_CALL_START for tool_call_id={tid}" starts[tid] = i elif t == "TOOL_CALL_END": tid = event.tool_call_id assert tid in starts, f"TOOL_CALL_END for unknown tool_call_id={tid}" assert tid not in ends, f"Duplicate TOOL_CALL_END for tool_call_id={tid}" ends.add(tid) unclosed = set(starts.keys()) - ends assert not unclosed, f"Unclosed tool calls: {unclosed}" def assert_no_run_error(self) -> None: """Assert no RUN_ERROR events exist.""" errors = self.get("RUN_ERROR") if errors: messages = [getattr(e, "message", str(e)) for e in errors] raise AssertionError(f"Found {len(errors)} RUN_ERROR event(s): {messages}") def assert_has_type(self, event_type: str) -> None: """Assert at least one event of the given type exists.""" assert event_type in self.types(), f"Expected {event_type!r} in stream. Available: {self.types()}" def assert_message_ids_consistent(self) -> None: """Assert TEXT_MESSAGE_CONTENT events reference valid, open message_ids.""" open_messages: set[str] = set() for event in self.events: t = self._type_str(event) if t == "TEXT_MESSAGE_START": open_messages.add(event.message_id) elif t == "TEXT_MESSAGE_END": open_messages.discard(event.message_id) elif t == "TEXT_MESSAGE_CONTENT": mid = event.message_id assert mid in open_messages, f"TEXT_MESSAGE_CONTENT references message_id={mid} which is not open" # ── Internal ── @staticmethod def _type_str(event: Any) -> str: """Extract event type as a plain string.""" t = getattr(event, "type", None) if t is None: return type(event).__name__ if isinstance(t, str): return t return getattr(t, "value", str(t)) @staticmethod def _event_dump(event: Any) -> dict[str, Any]: """Serialize an event object or return a raw event dict.""" if isinstance(event, dict): return event if hasattr(event, "model_dump"): return event.model_dump(by_alias=True, exclude_none=True) raw = getattr(event, "raw", None) if isinstance(raw, dict): return raw raise TypeError(f"Unsupported event type: {type(event).__name__}")