learningcircuit--local-deep-research/tests/fix_tests/test_milestone_logging_fix.py

"""
Test milestone logging functionality
"""

from datetime import datetime
from unittest.mock import MagicMock, Mock, patch


from local_deep_research.utilities.log_utils import (
    _get_research_id,
    frontend_progress_sink,
)


class TestMilestoneLogging:
    """Test milestone logging functionality"""

    def test_get_research_id_from_bound_logger(self):
        """Test that _get_research_id extracts research_id from bound logger"""
        # Create a mock record with bound research_id
        record = {"extra": {"research_id": 123}}

        research_id = _get_research_id(record)
        assert research_id == 123

    def test_get_research_id_from_flask_context(self):
        """Test that _get_research_id falls back to Flask g object"""
        record = {}

        with patch(
            "local_deep_research.utilities.log_utils.has_app_context",
            return_value=True,
        ):
            mock_g = Mock()
            mock_g.get.return_value = 456
            with patch(
                "local_deep_research.utilities.log_utils.g",
                mock_g,
            ):
                research_id = _get_research_id(record)
                assert research_id == 456

    def test_get_research_id_returns_none_when_not_found(self):
        """Test that _get_research_id returns None when no research_id is available"""
        record = {}

        with patch(
            "local_deep_research.utilities.log_utils.has_app_context",
            return_value=False,
        ):
            research_id = _get_research_id(record)
            assert research_id is None

    @patch("local_deep_research.utilities.log_utils.SocketIOService")
    def test_frontend_progress_sink_handles_milestone(
        self, mock_socket_service
    ):
        """Test that frontend_progress_sink properly handles MILESTONE logs"""
        # Mock the SocketIOService instance
        mock_socket_instance = Mock()
        mock_socket_service.return_value = mock_socket_instance

        # Create a mock level object - Mock(name=...) sets the display name,
        # not the .name attribute, so we must set it explicitly
        level_mock = MagicMock()
        level_mock.name = "MILESTONE"

        time_mock = MagicMock()
        time_mock.isoformat.return_value = "2023-01-01T00:00:00"

        # Create a mock message with MILESTONE level
        message = Mock()
        message.record = {
            "level": level_mock,
            "message": "Generating search questions for iteration 1",
            "time": time_mock,
            "extra": {"research_id": 789},
        }

        # Call the sink
        frontend_progress_sink(message)

        # Verify emit_to_subscribers was called correctly
        mock_socket_instance.emit_to_subscribers.assert_called_once_with(
            "progress",
            789,
            {
                "log_entry": {
                    "message": "Generating search questions for iteration 1",
                    "type": "MILESTONE",
                    "time": "2023-01-01T00:00:00",
                }
            },
            enable_logging=False,
        )

    def test_milestone_logging_in_research_service(self):
        """Test that milestone logging uses logger.bind() in research service"""
        # This test verifies the pattern used in research_service.py
        # where logger.bind(research_id=research_id) is called before logging
        from loguru import logger

        # Mock the logger
        with patch.object(logger, "bind") as mock_bind:
            mock_bound_logger = Mock()
            mock_bind.return_value = mock_bound_logger

            # Simulate the pattern used in research_service.py
            research_id = 999
            message = "Test milestone message"

            # This is how it's done in the actual code
            bound_logger = logger.bind(research_id=research_id)
            bound_logger.log("MILESTONE", message)

            # Verify
            mock_bind.assert_called_once_with(research_id=999)
            mock_bound_logger.log.assert_called_once_with("MILESTONE", message)

    def test_get_research_status_includes_milestone(self):
        """Test that get_research_status endpoint includes latest milestone.

        The research_routes.get_research_status is a Flask route handler that
        requires full request context. We verify the pattern works by testing
        the milestone log entry structure directly.
        """
        # Test that the expected log entry structure is correct
        milestone_message = "Current milestone message"
        log_entry = {
            "message": milestone_message,
            "type": "MILESTONE",
            "time": datetime.utcnow().isoformat(),
        }

        assert log_entry["type"] == "MILESTONE"
        assert log_entry["message"] == "Current milestone message"
        assert "time" in log_entry

    def test_milestone_logs_thread_safety(self):
        """Test that milestone logging works correctly across threads"""
        from concurrent.futures import ThreadPoolExecutor

        from loguru import logger

        results = []

        def thread_task(research_id):
            # Simulate binding research_id and logging
            with patch.object(logger, "bind") as mock_bind:
                mock_bound_logger = Mock()
                mock_bind.return_value = mock_bound_logger

                # This is how it's done in the actual code
                bound_logger = logger.bind(research_id=research_id)
                bound_logger.log("MILESTONE", f"Thread {research_id} milestone")

                # Store the call for verification
                results.append(
                    {
                        "research_id": research_id,
                        "bind_call": mock_bind.call_args,
                        "log_call": mock_bound_logger.log.call_args,
                    }
                )

        # Run multiple threads
        with ThreadPoolExecutor(max_workers=3) as executor:
            futures = [executor.submit(thread_task, i) for i in range(3)]
            for future in futures:
                future.result()

        # Verify each thread had its own research_id binding (order-independent)
        assert len(results) == 3
        seen_ids = set()
        for result in results:
            rid = result["research_id"]
            assert result["bind_call"][1]["research_id"] == rid
            assert result["log_call"][0][1] == f"Thread {rid} milestone"
            seen_ids.add(rid)
        assert seen_ids == {0, 1, 2}