learningcircuit--local-deep-research/tests/utilities/test_thread_context_coverage.py

"""Coverage tests for thread_context.py.

Focuses on untested paths:
- preserve_research_context: cleanup_current_thread called on success,
  cleanup_current_thread exception suppressed
"""

import sys
from types import ModuleType
from unittest.mock import MagicMock, patch

import pytest

from local_deep_research.utilities.thread_context import (
    clear_search_context,
    get_search_context,
    preserve_research_context,
    set_search_context,
)


def _cleanup():
    """Remove any leftover context on the current thread."""
    clear_search_context()


# ---------------------------------------------------------------------------
# preserve_research_context — cleanup paths
# ---------------------------------------------------------------------------


class TestPreserveResearchContextCleanup:
    """Tests for the cleanup logic inside preserve_research_context wrapper."""

    def setup_method(self):
        _cleanup()
        # Evict the module so each test exercises a fresh lazy import; keep
        # the original so teardown can put it back instead of leaving the
        # NEXT importer to build a new singleton distinct from the one
        # earlier-bound consumers hold.
        self._saved_tls_module = sys.modules.pop(
            "local_deep_research.database.thread_local_session", None
        )

    def teardown_method(self):
        _cleanup()
        sys.modules.pop(
            "local_deep_research.database.thread_local_session", None
        )
        db_pkg = sys.modules.get("local_deep_research.database")
        if self._saved_tls_module is not None:
            sys.modules["local_deep_research.database.thread_local_session"] = (
                self._saved_tls_module
            )
            if db_pkg is not None:
                db_pkg.thread_local_session = self._saved_tls_module
        elif db_pkg is not None and hasattr(db_pkg, "thread_local_session"):
            # The test's fresh import bound a throwaway module on the parent
            # package; drop it so attribute lookups can't resolve to it.
            delattr(db_pkg, "thread_local_session")

    def test_cleanup_current_thread_called_on_success(self):
        """When context is set and the function succeeds,
        cleanup_current_thread is called in the finally block."""
        set_search_context({"research_id": "cleanup-success"})

        @preserve_research_context
        def task():
            return 42

        clear_search_context()

        mock_cleanup_fn = MagicMock()
        fake_module = ModuleType(
            "local_deep_research.database.thread_local_session"
        )
        fake_module.cleanup_current_thread = mock_cleanup_fn

        with patch.dict(
            sys.modules,
            {
                "local_deep_research.database.thread_local_session": fake_module,
            },
        ):
            result = task()

        assert result == 42
        mock_cleanup_fn.assert_called_once()

    def test_cleanup_current_thread_called_on_exception(self):
        """When context is set and the function raises,
        cleanup_current_thread is still called in the finally block."""
        set_search_context({"research_id": "cleanup-exception"})

        @preserve_research_context
        def failing_task():
            raise ValueError("boom")

        clear_search_context()

        mock_cleanup_fn = MagicMock()
        fake_module = ModuleType(
            "local_deep_research.database.thread_local_session"
        )
        fake_module.cleanup_current_thread = mock_cleanup_fn

        with patch.dict(
            sys.modules,
            {
                "local_deep_research.database.thread_local_session": fake_module,
            },
        ):
            with pytest.raises(ValueError, match="boom"):
                failing_task()

        mock_cleanup_fn.assert_called_once()

    def test_cleanup_exception_is_suppressed(self):
        """If cleanup_current_thread raises, the exception is suppressed
        and does not propagate to the caller."""
        set_search_context({"research_id": "cleanup-suppressed"})

        @preserve_research_context
        def task():
            return "ok"

        clear_search_context()

        fake_module = ModuleType(
            "local_deep_research.database.thread_local_session"
        )
        fake_module.cleanup_current_thread = MagicMock(
            side_effect=RuntimeError("db engine error")
        )

        with patch.dict(
            sys.modules,
            {
                "local_deep_research.database.thread_local_session": fake_module,
            },
        ):
            result = task()

        assert result == "ok"
        assert get_search_context() is None

    def test_cleanup_import_failure_is_suppressed(self):
        """If importing thread_local_session itself fails, the exception
        is suppressed (caught by the bare except)."""
        set_search_context({"research_id": "import-fail"})

        @preserve_research_context
        def task():
            return "fine"

        clear_search_context()

        with patch.dict(
            sys.modules,
            {"local_deep_research.database.thread_local_session": None},
        ):
            result = task()

        assert result == "fine"
        assert get_search_context() is None

    def test_no_cleanup_when_no_context(self):
        """When no context was captured at decoration time, the cleanup
        block is never entered (context is None branch)."""
        assert get_search_context() is None

        @preserve_research_context
        def task():
            return "no-context"

        with patch.dict(
            sys.modules,
            {"local_deep_research.database.thread_local_session": None},
        ):
            result = task()

        assert result == "no-context"