test_streaming.py

"""Tests for StreamReader and AsyncStreamReader."""

from __future__ import annotations

import json
from unittest.mock import AsyncMock, MagicMock

import httpx
import pytest

from hawk.streaming import AsyncStreamReader, StreamReader


class TestStreamReader:
    """Tests for synchronous StreamReader."""

    def _make_stream(self, lines: list[str]) -> StreamReader:
        """Create a StreamReader from a list of raw lines."""
        response = MagicMock(spec=httpx.Response)
        response.iter_lines.return_value = iter(lines)
        response.close = MagicMock()
        return StreamReader(response)

    def test_basic_events(self) -> None:
        lines = [
            "data: Hello",
            "",
            "data: World",
            "",
            "event: done",
            "data: {}",
            "",
        ]
        reader = self._make_stream(lines)
        events = list(reader.events())

        assert len(events) == 3
        assert events[0].data == "Hello"
        assert events[0].event is None
        assert events[1].data == "World"
        assert events[2].event == "done"
        assert events[2].data == "{}"

    def test_collect_text(self) -> None:
        lines = [
            "data: Hello ",
            "",
            "data: World",
            "",
            "event: done",
            "data: {}",
            "",
        ]
        reader = self._make_stream(lines)
        text = reader.collect_text()

        assert text == "Hello World"

    def test_collect_tool_calls(self) -> None:
        tool_data = json.dumps(
            {
                "id": "tc-1",
                "name": "get_weather",
                "arguments": {"location": "NYC"},
            }
        )
        lines = [
            "event: tool_call",
            f"data: {tool_data}",
            "",
            "event: done",
            "data: {}",
            "",
        ]
        reader = self._make_stream(lines)
        calls = reader.collect_tool_calls()

        assert len(calls) == 1
        assert calls[0].name == "get_weather"
        assert calls[0].arguments == {"location": "NYC"}

    def test_multi_line_data(self) -> None:
        """SSE spec: multiple data: lines in one event are joined with newlines."""
        lines = [
            "data: line1",
            "data: line2",
            "",
        ]
        reader = self._make_stream(lines)
        events = list(reader.events())

        assert len(events) == 1
        assert events[0].data == "line1\nline2"

    def test_empty_data_field(self) -> None:
        lines = [
            "data:",
            "",
        ]
        reader = self._make_stream(lines)
        events = list(reader.events())

        assert len(events) == 1
        assert events[0].data == ""

    def test_context_manager(self) -> None:
        response = MagicMock(spec=httpx.Response)
        response.iter_lines.return_value = iter(["data: hi", ""])
        response.close = MagicMock()

        with StreamReader(response) as reader:
            events = list(reader.events())
            assert len(events) == 1

        response.close.assert_called_once()

    def test_multiple_empty_lines_ignored(self) -> None:
        lines = [
            "",
            "",
            "data: content",
            "",
            "",
            "",
        ]
        reader = self._make_stream(lines)
        events = list(reader.events())

        assert len(events) == 1
        assert events[0].data == "content"


@pytest.mark.asyncio
class TestAsyncStreamReader:
    """Tests for async stream reader."""

    async def test_basic_events(self) -> None:
        lines = [
            "data: Hello",
            "",
            "data: World",
            "",
        ]

        response = MagicMock(spec=httpx.Response)

        async def _aiter_lines():
            for line in lines:
                yield line

        response.aiter_lines.return_value = _aiter_lines()
        response.aclose = AsyncMock()

        reader = AsyncStreamReader(response)
        events = []
        async for event in reader:
            events.append(event)

        assert len(events) == 2
        assert events[0].data == "Hello"
        assert events[1].data == "World"

    async def test_multi_line_data(self) -> None:
        """SSE spec: multiple data: lines in one event are joined with newlines."""
        lines = [
            "data: line1",
            "data: line2",
            "",
        ]

        response = MagicMock(spec=httpx.Response)

        async def _aiter_lines():
            for line in lines:
                yield line

        response.aiter_lines.return_value = _aiter_lines()
        response.aclose = AsyncMock()

        reader = AsyncStreamReader(response)
        events = []
        async for event in reader:
            events.append(event)

        assert len(events) == 1
        assert events[0].data == "line1\nline2"

    async def test_collect_text(self) -> None:
        lines = [
            "data: Async ",
            "",
            "data: text",
            "",
            "event: done",
            "data: {}",
            "",
        ]

        response = MagicMock(spec=httpx.Response)

        async def _aiter_lines():
            for line in lines:
                yield line

        response.aiter_lines.return_value = _aiter_lines()
        response.aclose = AsyncMock()

        reader = AsyncStreamReader(response)
        text = await reader.collect_text()

        assert text == "Async text"