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

"""Extended tests for thread_context module.

Focuses on thread-local context propagation correctness,
thread isolation guarantees, and shallow-copy edge cases.
"""

import threading
from unittest.mock import patch

import pytest

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


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


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


# ---------------------------------------------------------------------------
# set_search_context
# ---------------------------------------------------------------------------


class TestSetSearchContextExtended:
    """Extended tests for set_search_context."""

    def setup_method(self):
        _cleanup()

    def teardown_method(self):
        _cleanup()

    def test_sets_context_in_thread_local_storage(self):
        """Setting context stores it for retrieval via get_search_context."""
        ctx = {"research_id": "abc", "query": "test query"}
        set_search_context(ctx)
        stored = get_search_context()
        assert stored is not None
        assert stored == {"research_id": "abc", "query": "test query"}

    def test_copies_context_so_original_mutation_has_no_effect(self):
        """set_search_context copies the dict -- mutating the original after
        the call must not change the stored value."""
        original = {"key": "original_value"}
        set_search_context(original)

        # Mutate the original
        original["key"] = "mutated"
        original["extra"] = "added"

        stored = get_search_context()
        assert stored["key"] == "original_value"
        assert "extra" not in stored

    def test_overwrites_existing_context_and_logs_debug(self):
        """When context already exists the old value is replaced and a
        debug message is emitted."""
        set_search_context({"round": 1})

        with patch(
            "local_deep_research.utilities.thread_context.logger"
        ) as mock_logger:
            set_search_context({"round": 2})
            mock_logger.debug.assert_called_once()
            assert (
                "overwritten" in mock_logger.debug.call_args[0][0].lower()
                or "already set" in mock_logger.debug.call_args[0][0].lower()
            )

        assert get_search_context() == {"round": 2}


# ---------------------------------------------------------------------------
# clear_search_context
# ---------------------------------------------------------------------------


class TestClearSearchContextExtended:
    """Extended tests for clear_search_context."""

    def setup_method(self):
        _cleanup()

    def teardown_method(self):
        _cleanup()

    def test_removes_context_from_thread_local_storage(self):
        """After clearing, get_search_context returns None."""
        set_search_context({"id": "will-be-cleared"})
        assert get_search_context() is not None

        clear_search_context()
        assert get_search_context() is None

    def test_noop_when_no_context_set(self):
        """Calling clear when there is nothing stored must not raise."""
        assert get_search_context() is None
        clear_search_context()  # should not raise
        assert get_search_context() is None


# ---------------------------------------------------------------------------
# get_search_context
# ---------------------------------------------------------------------------


class TestGetSearchContextExtended:
    """Extended tests for get_search_context."""

    def setup_method(self):
        _cleanup()

    def teardown_method(self):
        _cleanup()

    def test_returns_none_when_no_context_set(self):
        """Without prior set_search_context the return value is None."""
        assert get_search_context() is None

    def test_returns_copy_so_mutation_does_not_affect_stored(self):
        """Modifying the returned dict must not alter the stored context."""
        set_search_context({"research_id": "immutable-check"})

        returned = get_search_context()
        returned["research_id"] = "mutated"
        returned["injected"] = True

        # The stored context must remain untouched
        stored = get_search_context()
        assert stored["research_id"] == "immutable-check"
        assert "injected" not in stored

    def test_returns_correct_context_values(self):
        """The returned dict contains exactly the same key-value pairs that
        were originally stored."""
        ctx = {
            "research_id": "r-42",
            "user": "alice",
            "max_depth": 3,
            "tags": ["science", "ai"],
        }
        set_search_context(ctx)
        result = get_search_context()
        assert result == ctx


# ---------------------------------------------------------------------------
# search_context (context manager)
# ---------------------------------------------------------------------------


class TestSearchContextManagerExtended:
    """Extended tests for the search_context context manager."""

    def setup_method(self):
        _cleanup()

    def teardown_method(self):
        _cleanup()

    def test_sets_context_within_the_block(self):
        """Inside the with-block the context must be available."""
        with search_context({"research_id": "inside-check"}):
            ctx = get_search_context()
            assert ctx is not None
            assert ctx["research_id"] == "inside-check"

    def test_clears_context_after_normal_exit(self):
        """After the with-block exits normally the context is cleared."""
        with search_context({"research_id": "will-clear"}):
            pass  # normal exit
        assert get_search_context() is None

    def test_clears_context_when_exception_occurs(self):
        """Context is cleared even when an exception propagates out of the
        with-block."""
        with pytest.raises(RuntimeError, match="boom"):
            with search_context({"research_id": "exception-path"}):
                assert get_search_context() is not None
                raise RuntimeError("boom")

        assert get_search_context() is None

    def test_re_raises_exception_from_inside_block(self):
        """The context manager does not swallow exceptions."""
        with pytest.raises(TypeError, match="bad type"):
            with search_context({"research_id": "re-raise"}):
                raise TypeError("bad type")


# ---------------------------------------------------------------------------
# preserve_research_context (decorator)
# ---------------------------------------------------------------------------


class TestPreserveResearchContextExtended:
    """Extended tests for the preserve_research_context decorator."""

    def setup_method(self):
        _cleanup()

    def teardown_method(self):
        _cleanup()

    def test_no_context_function_runs_normally(self):
        """When no context exists at decoration time the wrapped function
        executes without setting or clearing any context."""
        # Ensure no context is set
        assert get_search_context() is None

        @preserve_research_context
        def plain_task(x, y):
            """Add two numbers."""
            # Context should remain unset inside the function
            assert get_search_context() is None
            return x + y

        result = plain_task(3, 7)
        assert result == 10
        # Still no context after the call
        assert get_search_context() is None

    def test_with_context_sets_before_and_clears_after(self):
        """When context is set at decoration time the decorator sets it
        before the function body and clears it in the finally block."""
        set_search_context({"research_id": "decorator-lifecycle"})

        observed_inside = []

        @preserve_research_context
        def worker():
            observed_inside.append(get_search_context())

        # Clear the main-thread context so we can verify the decorator
        # independently restores it inside the wrapper.
        clear_search_context()

        worker()

        # During execution the context was set
        assert observed_inside[0] is not None
        assert observed_inside[0]["research_id"] == "decorator-lifecycle"

        # After execution the context is cleared
        assert get_search_context() is None

    def test_preserves_function_name_and_docstring(self):
        """functools.wraps should keep __name__ and __doc__ intact."""
        set_search_context({"research_id": "wraps-test"})

        @preserve_research_context
        def documented_function(a, b):
            """This is my docstring."""
            return a * b

        assert documented_function.__name__ == "documented_function"
        assert documented_function.__doc__ == "This is my docstring."

    def test_clears_context_when_decorated_function_raises(self):
        """If the wrapped function raises, context is still cleaned up via
        the finally clause."""
        set_search_context({"research_id": "error-cleanup"})

        @preserve_research_context
        def failing():
            raise ValueError("intentional failure")

        clear_search_context()

        with pytest.raises(ValueError, match="intentional failure"):
            failing()

        assert get_search_context() is None


# ---------------------------------------------------------------------------
# Thread isolation
# ---------------------------------------------------------------------------


class TestThreadIsolationExtended:
    """Thread isolation guarantees for thread-local storage."""

    def setup_method(self):
        _cleanup()

    def teardown_method(self):
        _cleanup()

    def test_context_set_in_one_thread_not_visible_in_another(self):
        """A context set on the main thread must not be visible in a freshly
        spawned thread."""
        set_search_context({"owner": "main-thread"})

        child_result = [None]
        error_holder = [None]

        def child():
            try:
                child_result[0] = get_search_context()
            except Exception as exc:
                error_holder[0] = exc

        t = threading.Thread(target=child)
        t.start()
        t.join()

        assert error_holder[0] is None, (
            f"Child thread raised: {error_holder[0]}"
        )
        assert child_result[0] is None, (
            "Child thread should not see main thread context"
        )

    def test_clearing_in_one_thread_does_not_affect_another(self):
        """Clearing context on the main thread must not affect a context
        set independently in a child thread."""
        barrier = threading.Barrier(2, timeout=5)
        child_context_after_main_clear = [None]
        error_holder = [None]

        def child():
            try:
                set_search_context({"owner": "child-thread"})
                # Signal that child has set its context
                barrier.wait()
                # Wait for main thread to clear its context
                barrier.wait()
                child_context_after_main_clear[0] = get_search_context()
            except Exception as exc:
                error_holder[0] = exc

        set_search_context({"owner": "main-thread"})
        t = threading.Thread(target=child)
        t.start()

        # Wait until child has set its context
        barrier.wait()
        # Main clears its own context
        clear_search_context()
        # Signal child to read its context
        barrier.wait()

        t.join()

        assert error_holder[0] is None, (
            f"Child thread raised: {error_holder[0]}"
        )
        assert child_context_after_main_clear[0] is not None
        assert child_context_after_main_clear[0]["owner"] == "child-thread"

    def test_search_context_manager_isolates_across_threads(self):
        """Using the search_context context manager in two threads
        concurrently must not cause cross-thread leakage."""
        barrier = threading.Barrier(2, timeout=5)
        results = {}
        errors = []

        def worker(name, ctx_value):
            try:
                with search_context({"worker": ctx_value}):
                    # Synchronise so both threads are inside their
                    # context managers at the same time.
                    barrier.wait()
                    observed = get_search_context()
                    results[name] = observed["worker"] if observed else None
            except Exception as exc:
                errors.append(exc)

        t1 = threading.Thread(target=worker, args=("t1", "value-from-t1"))
        t2 = threading.Thread(target=worker, args=("t2", "value-from-t2"))

        t1.start()
        t2.start()
        t1.join()
        t2.join()

        assert not errors, f"Threads raised errors: {errors}"
        assert results["t1"] == "value-from-t1"
        assert results["t2"] == "value-from-t2"


# ---------------------------------------------------------------------------
# Edge cases
# ---------------------------------------------------------------------------


class TestEdgeCasesExtended:
    """Edge-case scenarios for thread_context."""

    def setup_method(self):
        _cleanup()

    def teardown_method(self):
        _cleanup()

    def test_empty_dict_context_is_valid(self):
        """An empty dict is a legitimate context -- it should be stored and
        retrieved without being confused with None."""
        set_search_context({})
        result = get_search_context()
        assert result is not None
        assert result == {}

    def test_shallow_copy_semantics_for_nested_dicts(self):
        """set_search_context and get_search_context use dict.copy() which is
        a shallow copy.  Nested mutable objects are shared references.

        This test documents the expected (but potentially surprising)
        behaviour: mutating a nested dict through a retrieved copy *does*
        affect the stored context because dict.copy() is shallow."""
        nested = {"inner_key": "original"}
        set_search_context({"nested": nested})

        retrieved = get_search_context()
        # Top-level key mutation does NOT leak back
        retrieved["top_new"] = "added"
        assert "top_new" not in get_search_context()

        # Nested mutation DOES leak back (shallow copy)
        retrieved["nested"]["inner_key"] = "mutated"
        assert get_search_context()["nested"]["inner_key"] == "mutated"