"""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"