tatlock/tests/integration/test_steward_streaming.py

"""
Integration tests for Steward + Tatlock streaming.

Tests the complete streaming flow with Steward preprocessing.
"""
import pytest
from unittest.mock import AsyncMock, MagicMock, patch

from src.responses.schemas import ResponseRequest
from src.responses.streaming import StreamingCoordinator, StreamEventType
from src.core.startup import initialize_application


@pytest.fixture(scope="module", autouse=True)
def setup_household_registry():
    """Initialize household registry before running tests."""
    initialize_application()


class TestStewardStreaming:
    """Test Steward + Tatlock streaming integration."""

    @pytest.mark.asyncio
    async def test_stream_with_steward_basic(self):
        """Test basic streaming with Steward preprocessing."""
        request = ResponseRequest(
            model="tatlock",
            input=[{"role": "user", "content": "What's 2 + 2?"}],
            stream=True,
        )

        # Mock the Steward analysis
        with patch("src.core.preprocessing.analyze_request") as mock_steward:
            # Mock the streaming method (async generator)
            with patch("src.agents.tatlock.TatlockAgent.run_with_scoped_tools_stream") as mock_tatlock_stream:
                from src.agents.steward.schemas import ConversationContext, StewardRecommendation

                # Mock Steward recommendation
                mock_steward.return_value = StewardRecommendation(
                    recommended_capabilities=["tatlock_core"],
                    reasoning="Math calculation requires tatlock_core",
                    estimated_complexity="simple",
                    conversation_context=ConversationContext(has_previous_context=False),
                )

                # Mock Tatlock streaming response as async generator
                async def mock_stream(*args, **kwargs):
                    yield "Certainly, sir. "
                    yield "2 + 2 equals 4."

                mock_tatlock_stream.return_value = mock_stream()

                # Execute streaming
                coordinator = StreamingCoordinator()
                events = []

                async for event in coordinator.stream_response_with_steward(request):
                    events.append(event)

                # Verify event sequence
                event_types = [e.event for e in events]

                # Should have reasoning summary deltas
                assert StreamEventType.REASONING_SUMMARY_DELTA in event_types
                assert StreamEventType.REASONING_SUMMARY_DONE in event_types

                # Should have output text deltas
                assert StreamEventType.OUTPUT_TEXT_DELTA in event_types
                assert StreamEventType.OUTPUT_TEXT_DONE in event_types

                # Should end with response.done
                assert events[-1].event == StreamEventType.RESPONSE_DONE

                # Verify Steward and Tatlock were called
                assert mock_steward.called
                assert mock_tatlock_stream.called

    @pytest.mark.asyncio
    async def test_stream_with_conversation_history(self):
        """Test streaming with conversation history."""
        request = ResponseRequest(
            model="tatlock",
            input=[
                {"role": "user", "content": "What's 5 times 3?"},
                {"role": "assistant", "content": "That equals 15, sir."},
                {"role": "user", "content": "And divided by 3?"},
            ],
            stream=True,
        )

        with patch("src.core.preprocessing.analyze_request") as mock_steward:
            with patch("src.agents.tatlock.TatlockAgent.run_with_scoped_tools_stream") as mock_tatlock_stream:
                from src.agents.steward.schemas import ConversationContext, StewardRecommendation

                mock_steward.return_value = StewardRecommendation(
                    recommended_capabilities=["tatlock_core"],
                    reasoning="Follow-up calculation based on previous result of 15",
                    estimated_complexity="simple",
                    conversation_context=ConversationContext(
                        has_previous_context=True,
                        relevant_turns=[0],
                        context_summary="Previous calculation in turn 0"
                    ),
                )

                async def mock_stream(*args, **kwargs):
                    yield "15 divided by 3 equals 5, sir."

                mock_tatlock_stream.return_value = mock_stream()

                coordinator = StreamingCoordinator()
                events = []

                async for event in coordinator.stream_response_with_steward(request):
                    events.append(event)

                # Verify conversation history was passed to Steward
                call_kwargs = mock_steward.call_args[1]
                assert "conversation_history" in call_kwargs
                assert len(call_kwargs["conversation_history"]) == 2  # First Q&A pair

                # Verify final response includes both reasoning and message
                final_event = events[-1]
                assert final_event.event == StreamEventType.RESPONSE_DONE
                assert len(final_event.response.output) == 2  # Reasoning + Message

    @pytest.mark.asyncio
    async def test_stream_reasoning_contains_steward_analysis(self):
        """Test that reasoning summary contains Steward's analysis."""
        request = ResponseRequest(
            model="tatlock",
            input=[{"role": "user", "content": "Test request"}],
            stream=True,
        )

        with patch("src.core.preprocessing.analyze_request") as mock_steward:
            with patch("src.agents.tatlock.TatlockAgent.run_with_scoped_tools_stream") as mock_tatlock_stream:
                from src.agents.steward.schemas import ConversationContext, StewardRecommendation

                mock_steward.return_value = StewardRecommendation(
                    recommended_capabilities=["tatlock_core"],
                    reasoning="This is a test analysis with specific markers",
                    estimated_complexity="simple",
                    conversation_context=ConversationContext(has_previous_context=False),
                )

                async def mock_stream(*args, **kwargs):
                    yield "Test response"

                mock_tatlock_stream.return_value = mock_stream()

                coordinator = StreamingCoordinator()
                reasoning_deltas = []

                async for event in coordinator.stream_response_with_steward(request):
                    if event.event == StreamEventType.REASONING_SUMMARY_DELTA:
                        reasoning_deltas.append(event.delta)

                # Combine all reasoning deltas
                full_reasoning = "".join(reasoning_deltas)

                # Should contain Steward's analysis
                assert "test analysis" in full_reasoning.lower()
                assert len(reasoning_deltas) > 0, "Should have streamed reasoning deltas"

    @pytest.mark.asyncio
    async def test_stream_with_missing_capabilities(self):
        """Test streaming when Steward detects missing capabilities."""
        request = ResponseRequest(
            model="tatlock",
            input=[{"role": "user", "content": "Generate an image of a sunset"}],
            stream=True,
        )

        with patch("src.core.preprocessing.analyze_request") as mock_steward:
            with patch("src.agents.tatlock.TatlockAgent.run_with_scoped_tools_stream") as mock_tatlock_stream:
                from src.agents.steward.schemas import ConversationContext, StewardRecommendation

                mock_steward.return_value = StewardRecommendation(
                    recommended_capabilities=[],
                    reasoning="Image generation not available in current toolset",
                    estimated_complexity="simple",
                    conversation_context=ConversationContext(has_previous_context=False),
                    missing_capabilities="Image generation capability would be needed",
                )

                async def mock_stream(*args, **kwargs):
                    yield "I'm afraid I don't have image generation capabilities, sir."

                mock_tatlock_stream.return_value = mock_stream()

                coordinator = StreamingCoordinator()
                events = []

                async for event in coordinator.stream_response_with_steward(request):
                    events.append(event)

                # Should complete successfully even with missing capabilities
                assert events[-1].event == StreamEventType.RESPONSE_DONE

                # Verify empty scoped tools were passed to stream method
                tatlock_kwargs = mock_tatlock_stream.call_args[1]
                assert "scoped_tools" in tatlock_kwargs
                assert tatlock_kwargs["scoped_tools"] == []