/tests/citation_handlers/test_streaming.py

"""
Comprehensive tests for LLM streaming functionality in citation handlers.
Tests stream callbacks, chunk handling, and streaming fallback behavior.
"""

from unittest.mock import Mock
from typing import List


class MockChunk:
    """Mock LLM chunk with content attribute."""

    def __init__(self, content: str):
        self.content = content


class MockLLMWithStreaming:
    """Mock LLM that supports streaming.

    The ``invoke`` Mock is kept as the public attribute so tests can use
    ``mock_llm.invoke.assert_called_once()``. A previous version of this
    class also defined ``def invoke(self, prompt)`` as a method — Python's
    attribute lookup means the instance Mock always shadowed it, so the
    method body (including ``self.invoke_called = True``) was unreachable
    dead code. The flag has been removed to match what actually works.
    """

    def __init__(self, chunks: List[str] = None, should_fail: bool = False):
        self.chunks = chunks or ["Hello", " ", "world", "!"]
        self.should_fail = should_fail
        self.invoke = Mock(return_value=Mock(content="".join(self.chunks)))
        self.stream_called = False

    def stream(self, prompt: str):
        """Yield chunks like a real streaming LLM."""
        self.stream_called = True
        if self.should_fail:
            raise Exception("Streaming failed")
        for chunk_text in self.chunks:
            yield MockChunk(chunk_text)


class TestStreamCallbackSetup:
    """Tests for stream callback configuration."""

    def test_set_stream_callback_stores_callback(self):
        """Test that set_stream_callback properly stores the callback."""
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        mock_llm = Mock()
        handler = StandardCitationHandler(llm=mock_llm)

        callback = Mock()
        handler.set_stream_callback(callback)

        assert handler.stream_callback == callback

    def test_set_stream_callback_to_none(self):
        """Test that stream callback can be cleared."""
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        mock_llm = Mock()
        handler = StandardCitationHandler(llm=mock_llm)

        # Set callback then clear it
        callback = Mock()
        handler.set_stream_callback(callback)
        handler.set_stream_callback(None)

        assert handler.stream_callback is None

    def test_citation_handler_passes_callback_to_underlying_handler(self):
        """Test that CitationHandler passes callback to the underlying handler."""
        from local_deep_research.citation_handler import CitationHandler

        mock_llm = Mock()
        mock_llm.invoke.return_value = Mock(content="Test response")

        handler = CitationHandler(mock_llm)
        callback = Mock()

        handler.set_stream_callback(callback)

        # The underlying handler should have the callback set
        assert handler._handler.stream_callback == callback


class TestStreamingInvocation:
    """Tests for _invoke_with_streaming method."""

    def test_invoke_with_streaming_uses_stream_when_callback_set(self):
        """Test that streaming is used when callback is set."""
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        mock_llm = MockLLMWithStreaming(chunks=["chunk1", "chunk2", "chunk3"])
        handler = StandardCitationHandler(llm=mock_llm)

        callback = Mock()
        handler.set_stream_callback(callback)

        result = handler._invoke_with_streaming("test prompt")

        # Verify streaming was used
        assert mock_llm.stream_called
        # Verify callback was called for each chunk
        assert callback.call_count == 3
        callback.assert_any_call("chunk1")
        callback.assert_any_call("chunk2")
        callback.assert_any_call("chunk3")
        # Verify complete response is returned
        assert result == "chunk1chunk2chunk3"

    def test_invoke_with_streaming_uses_invoke_when_no_callback(self):
        """Test that regular invoke is used when no callback is set."""
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        mock_llm = Mock()
        mock_llm.invoke.return_value = Mock(content="Regular response")

        handler = StandardCitationHandler(llm=mock_llm)
        # No callback set

        result = handler._invoke_with_streaming("test prompt")

        mock_llm.invoke.assert_called_once_with("test prompt")
        assert result == "Regular response"

    def test_invoke_with_streaming_falls_back_on_stream_error(self):
        """Test fallback to invoke when streaming fails."""
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        mock_llm = Mock()
        mock_llm.stream.side_effect = Exception("Stream error")
        mock_llm.invoke.return_value = Mock(content="Fallback response")

        handler = StandardCitationHandler(llm=mock_llm)
        handler.set_stream_callback(Mock())

        result = handler._invoke_with_streaming("test prompt")

        # Should have tried streaming first
        mock_llm.stream.assert_called_once()
        # Then fallen back to invoke
        mock_llm.invoke.assert_called_once()
        assert result == "Fallback response"

    def test_invoke_with_streaming_returns_partial_after_mid_stream_failure(
        self,
    ):
        """When streaming fails AFTER some chunks were already emitted to the
        client, the handler returns the partial content instead of falling
        back to ``invoke()``. Restarting via invoke would (a) double-bill
        the LLM and (b) make the frontend's accumulated streamed text
        diverge from the new full response (chat bubble would carry the
        streamed prefix while the DB row would carry the invoke result).
        See ``base_citation_handler._invoke_with_streaming`` comments."""
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        mock_llm = Mock()

        def partial_stream(prompt):
            yield Mock(content="partial ")
            yield Mock(content="content")
            raise Exception("connection dropped")

        mock_llm.stream.side_effect = partial_stream
        mock_llm.invoke.return_value = Mock(content="Full fallback response")

        handler = StandardCitationHandler(llm=mock_llm)
        received = []
        handler.set_stream_callback(lambda c: received.append(c))

        result = handler._invoke_with_streaming("test prompt")

        # Must NOT call invoke after partial chunks (double-bill / divergence).
        mock_llm.invoke.assert_not_called()
        # Returns the partial content the client already saw.
        assert result == "partial content"
        # Partial chunks were received before failure.
        assert "partial " in received
        assert "content" in received

    def test_invoke_with_streaming_handles_string_chunks(self):
        """Test handling of string chunks (no .content attribute)."""
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        # Some LLMs return plain strings instead of objects with .content
        mock_llm = Mock()
        mock_llm.stream.return_value = iter(["plain", " ", "string"])

        handler = StandardCitationHandler(llm=mock_llm)
        callback = Mock()
        handler.set_stream_callback(callback)

        result = handler._invoke_with_streaming("test prompt")

        assert callback.call_count == 3
        assert result == "plain string"

    def test_invoke_with_streaming_skips_empty_chunks(self):
        """Test that empty chunks don't trigger callback."""
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        mock_llm = Mock()
        mock_llm.stream.return_value = iter(
            [
                MockChunk("hello"),
                MockChunk(""),  # Empty chunk
                MockChunk(" world"),
            ]
        )

        handler = StandardCitationHandler(llm=mock_llm)
        callback = Mock()
        handler.set_stream_callback(callback)

        result = handler._invoke_with_streaming("test prompt")

        # Only non-empty chunks should trigger callback
        assert callback.call_count == 2
        callback.assert_any_call("hello")
        callback.assert_any_call(" world")
        assert result == "hello world"


class TestStreamingWithAnalyzeMethods:
    """Tests for streaming integration with analyze_initial and analyze_followup."""

    def test_analyze_initial_uses_streaming_when_callback_set(self):
        """Test that analyze_initial uses streaming when callback is set."""
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        mock_llm = Mock()
        mock_llm.stream.return_value = iter(
            [
                MockChunk("Analysis: "),
                MockChunk("The topic "),
                MockChunk("is interesting [1]."),
            ]
        )

        handler = StandardCitationHandler(llm=mock_llm)
        callback = Mock()
        handler.set_stream_callback(callback)

        search_results = [
            {
                "title": "Test",
                "link": "https://example.com",
                "snippet": "Test content",
            }
        ]

        result = handler.analyze_initial("What is the topic?", search_results)

        # Verify streaming was used
        mock_llm.stream.assert_called_once()
        # Verify callback was called for each chunk
        assert callback.call_count == 3
        # Verify content is assembled correctly
        assert result["content"] == "Analysis: The topic is interesting [1]."

    def test_analyze_followup_uses_streaming_when_callback_set(self):
        """Test that analyze_followup uses streaming when callback is set."""
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        mock_llm = Mock()
        mock_llm.stream.return_value = iter(
            [
                MockChunk("Follow-up: "),
                MockChunk("Additional info [2]."),
            ]
        )
        # Mock invoke for fact-checking (if enabled)
        mock_llm.invoke.return_value = Mock(content="Fact check passed")

        handler = StandardCitationHandler(
            llm=mock_llm,
            settings_snapshot={"general.enable_fact_checking": False},
        )
        callback = Mock()
        handler.set_stream_callback(callback)

        search_results = [
            {
                "title": "Test",
                "link": "https://example.com",
                "snippet": "More content",
            }
        ]

        result = handler.analyze_followup(
            "Follow-up question?",
            search_results,
            "Previous knowledge",
            nr_of_links=1,
        )

        # Verify streaming was used
        mock_llm.stream.assert_called_once()
        assert callback.call_count == 2
        assert result["content"] == "Follow-up: Additional info [2]."


class TestForcedAnswerHandlerStreaming:
    """Tests for streaming in ForcedAnswerCitationHandler."""

    def test_forced_answer_handler_supports_streaming(self):
        """Test that ForcedAnswerCitationHandler supports streaming."""
        from local_deep_research.citation_handlers.forced_answer_citation_handler import (
            ForcedAnswerCitationHandler,
        )

        mock_llm = Mock()
        mock_llm.stream.return_value = iter(
            [
                MockChunk("Forced answer: "),
                MockChunk("specific result."),
            ]
        )

        handler = ForcedAnswerCitationHandler(llm=mock_llm)
        callback = Mock()
        handler.set_stream_callback(callback)

        assert handler.stream_callback == callback

    def test_forced_answer_handler_streams_during_analyze(self):
        """Test ForcedAnswerCitationHandler uses streaming during analysis."""
        from local_deep_research.citation_handlers.forced_answer_citation_handler import (
            ForcedAnswerCitationHandler,
        )

        mock_llm = Mock()
        mock_llm.stream.return_value = iter(
            [
                MockChunk("Answer: test"),
            ]
        )

        handler = ForcedAnswerCitationHandler(llm=mock_llm)
        callback = Mock()
        handler.set_stream_callback(callback)

        search_results = [
            {
                "title": "Test",
                "link": "https://example.com",
                "snippet": "Content",
            }
        ]

        _ = handler.analyze_initial("Question?", search_results)

        mock_llm.stream.assert_called()
        callback.assert_called()


class TestPrecisionExtractionHandlerStreaming:
    """Tests for streaming in PrecisionExtractionHandler."""

    def test_precision_handler_supports_streaming(self):
        """Test that PrecisionExtractionHandler supports streaming."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        mock_llm = Mock()
        handler = PrecisionExtractionHandler(llm=mock_llm)
        callback = Mock()
        handler.set_stream_callback(callback)

        assert handler.stream_callback == callback


class TestStreamingEdgeCases:
    """Tests for edge cases in streaming behavior."""

    def test_streaming_with_very_long_response(self):
        """Test streaming handles very long responses."""
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        # Generate many chunks
        chunks = [f"chunk{i} " for i in range(1000)]
        mock_llm = Mock()
        mock_llm.stream.return_value = iter([MockChunk(c) for c in chunks])

        handler = StandardCitationHandler(llm=mock_llm)
        callback = Mock()
        handler.set_stream_callback(callback)

        result = handler._invoke_with_streaming("test")

        assert callback.call_count == 1000
        # The chunks reach the client verbatim (1000 callbacks), but the
        # returned/persisted value is normalized exactly like the
        # non-streaming invoke() path (get_llm_response_text -> strip()).
        assert result == "".join(chunks).strip()

    def test_streaming_strips_think_tags_from_result(self):
        """Reasoning models' <think> blocks must not leak into the
        returned/persisted synthesis — the streaming path must normalize
        the joined chunks the same way the non-streaming invoke() path does.

        Regression for the streaming citation path returning raw joined
        chunks without <think> stripping (PR #2953).
        """
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        chunks = [
            "<think>",
            "internal reasoning ",
            "</think>",
            "Final answer.",
        ]
        mock_llm = Mock()
        mock_llm.stream.return_value = iter([MockChunk(c) for c in chunks])

        handler = StandardCitationHandler(llm=mock_llm)
        callback = Mock()
        handler.set_stream_callback(callback)

        result = handler._invoke_with_streaming("test")

        # The live stream still receives every chunk (stripping the live
        # token stream is a separate concern); the returned value, however,
        # is think-stripped so the persisted answer matches invoke().
        assert callback.call_count == len(chunks)
        assert "<think>" not in result
        assert "internal reasoning" not in result
        assert result == "Final answer."

    def test_streaming_with_special_characters(self):
        """Test streaming handles special characters correctly."""
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        special_chunks = [
            "Hello\n",
            "World\t",
            "Special: €£¥",
            "Emoji: 🎉",
            "Unicode: 日本語",
        ]
        mock_llm = Mock()
        mock_llm.stream.return_value = iter(
            [MockChunk(c) for c in special_chunks]
        )

        handler = StandardCitationHandler(llm=mock_llm)
        callback = Mock()
        handler.set_stream_callback(callback)

        result = handler._invoke_with_streaming("test")

        assert callback.call_count == 5
        assert "€£¥" in result
        assert "🎉" in result
        assert "日本語" in result

    def test_streaming_callback_exception_does_not_stop_streaming(self):
        """Test that exception in callback doesn't stop streaming."""
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        mock_llm = Mock()
        mock_llm.stream.return_value = iter(
            [
                MockChunk("chunk1"),
                MockChunk("chunk2"),
                MockChunk("chunk3"),
            ]
        )

        handler = StandardCitationHandler(llm=mock_llm)

        # Callback that tracks calls but we test if all chunks are processed
        received_chunks = []

        def callback(chunk):
            received_chunks.append(chunk)

        handler.set_stream_callback(callback)

        result = handler._invoke_with_streaming("test")

        # All chunks should be received despite any issues
        assert len(received_chunks) == 3
        assert result == "chunk1chunk2chunk3"

    def test_multiple_callback_changes_during_session(self):
        """Test changing callbacks between invocations."""
        from local_deep_research.citation_handlers.standard_citation_handler import (
            StandardCitationHandler,
        )

        mock_llm = Mock()
        mock_llm.stream.return_value = iter([MockChunk("test")])
        mock_llm.invoke.return_value = Mock(content="invoke response")

        handler = StandardCitationHandler(llm=mock_llm)

        # First invocation with callback
        callback1 = Mock()
        handler.set_stream_callback(callback1)
        mock_llm.stream.return_value = iter([MockChunk("first")])
        handler._invoke_with_streaming("test1")
        assert callback1.call_count == 1

        # Second invocation with different callback
        callback2 = Mock()
        handler.set_stream_callback(callback2)
        mock_llm.stream.return_value = iter([MockChunk("second")])
        handler._invoke_with_streaming("test2")
        assert callback2.call_count == 1
        assert callback1.call_count == 1  # First callback not called again

        # Third invocation without callback
        handler.set_stream_callback(None)
        handler._invoke_with_streaming("test3")
        mock_llm.invoke.assert_called()


class TestStreamingIntegrationWithCitationHandler:
    """Integration tests for streaming through the main CitationHandler."""

    def test_citation_handler_streaming_end_to_end(self):
        """Test end-to-end streaming through CitationHandler."""
        from local_deep_research.citation_handler import CitationHandler

        mock_llm = Mock()
        mock_llm.stream.return_value = iter(
            [
                MockChunk("Research shows "),
                MockChunk("that the topic "),
                MockChunk("is important [1]."),
            ]
        )

        handler = CitationHandler(mock_llm, handler_type="standard")

        received_chunks = []

        def callback(chunk):
            received_chunks.append(chunk)

        handler.set_stream_callback(callback)

        search_results = [
            {
                "title": "Test",
                "link": "https://example.com",
                "full_content": "Test content",
            }
        ]

        result = handler.analyze_initial("What is the topic?", search_results)

        assert len(received_chunks) == 3
        assert received_chunks[0] == "Research shows "
        assert "important [1]" in result["content"]

    def test_different_handler_types_support_streaming(self):
        """Test that all handler types support streaming."""
        from local_deep_research.citation_handler import CitationHandler

        handler_types = ["standard", "forced", "precision"]

        for handler_type in handler_types:
            mock_llm = Mock()
            mock_llm.stream.return_value = iter([MockChunk("test")])
            mock_llm.invoke.return_value = Mock(content="test")

            handler = CitationHandler(mock_llm, handler_type=handler_type)
            callback = Mock()
            handler.set_stream_callback(callback)

            # Verify callback was set on the underlying handler
            assert handler._handler.stream_callback == callback, (
                f"Handler type '{handler_type}' should support streaming"
            )