sub-agents-skills/tests/test_stream.py at 3c1d1861044407a3067a20f54fdce586ff86e971 · shinpr/sub-agents-skills · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
"""Tests for _stream.StreamProcessor — NDJSON parsing for each backend's stream."""

from __future__ import annotations

from _stream import StreamProcessor


class TestStreamProcessor:
    def test_claude_result(self):
        processor = StreamProcessor()
        assert processor.process_line('{"type": "result", "result": "hello"}')
        result = processor.get_result()
        assert result["result"] == "hello"

    def test_gemini_stream(self):
        processor = StreamProcessor()
        assert not processor.process_line('{"type": "init"}')
        assert not processor.process_line(
            '{"type": "message", "role": "assistant", "content": "part1"}'
        )
        assert not processor.process_line(
            '{"type": "message", "role": "assistant", "content": "part2"}'
        )
        assert processor.process_line('{"type": "result", "status": "success"}')
        result = processor.get_result()
        assert result["result"] == "part1part2"

    def test_gemini_stream_with_banner_prefix(self):
        # Gemini CLI can prepend a non-JSON banner to the first event line, e.g.
        # "MCP issues detected. Run /mcp list for status." glued onto `init`.
        # The processor must recover the JSON so is_gemini is set and assistant
        # text is captured (otherwise the result is silently empty).
        processor = StreamProcessor()
        assert not processor.process_line(
            'MCP issues detected. Run /mcp list for status.{"type": "init"}'
        )
        assert not processor.process_line(
            '{"type": "message", "role": "assistant", "content": "hi"}'
        )
        assert processor.process_line('{"type": "result", "status": "success"}')
        assert processor.get_result()["result"] == "hi"

    def test_codex_stream(self):
        processor = StreamProcessor()
        assert not processor.process_line('{"type": "thread.started"}')
        assert not processor.process_line(
            '{"type": "item.completed", "item": {"type": "agent_message", "text": "msg1"}}'
        )
        assert not processor.process_line(
            '{"type": "item.completed", "item": {"type": "agent_message", "text": "msg2"}}'
        )
        assert processor.process_line('{"type": "turn.completed"}')
        result = processor.get_result()
        assert result["result"] == "msg1\nmsg2"