async def test_streaming_agent_citation_extraction(mock_streaming_agent):
    """Test that streaming agent extracts citations from streamed content."""
    # Run the agent
    messages = [Message(role="user", content="Test query")]

    # We need to run this in a coroutine
    stream = mock_streaming_agent.arun(messages=messages)
    output = await collect_stream_output(stream)

    # Look for citation events in the output
    citation_events = [
        line for line in output
        if 'event: citation' in line
    ]

    assert len(citation_events) > 0, "Citation events should be emitted"

    # Check citation IDs in events
    citation_abc = any('abc1234' in event for event in citation_events)
    citation_def = any('def5678' in event for event in citation_events)

    assert citation_abc, "Citation abc1234 should be found in stream output"
    assert citation_def, "Citation def5678 should be found in stream output"
Test that streaming agent extracts citations from streamed content.
py/tests/unit/agent/test_agent_citations.py