44This test suite validates context propagation behavior across various concurrency patterns.
55
66TEST ISOLATION STRATEGY:
7- - Tests use pytest-forked to run each test in an isolated process
8- - This ensures setup_threads() patches don't leak between tests
7+ - Tests that call setup_threads() run in isolated subprocesses (via subprocess.run)
8+ - This ensures setup_threads() monkey- patches don't leak between tests
99- Use unpatched(scenario) for xfail tests (documents context loss)
1010- Use patched(scenario) for tests that prove setup_threads() fixes it
1111
@@ -16,12 +16,12 @@ def _threadpool_scenario(test_logger, with_memory_logger):
1616 test_threadpool_loses_context = unpatched(_threadpool_scenario)
1717 test_threadpool_with_patch = patched(_threadpool_scenario)
1818
19- Run with: pytest --forked src/braintrust/test_context.py
19+ Run with: pytest src/braintrust/test_context.py
2020"""
2121
2222import asyncio
2323import concurrent .futures
24- import functools
24+ import subprocess
2525import sys
2626import threading
2727from typing import AsyncGenerator , Callable , Generator , TypeVar
@@ -30,44 +30,84 @@ def _threadpool_scenario(test_logger, with_memory_logger):
3030import pytest
3131from braintrust import current_span , start_span
3232from braintrust .test_helpers import init_test_logger , with_memory_logger # noqa: F401
33- from braintrust .wrappers .threads import setup_threads
3433
3534
3635F = TypeVar ("F" , bound = Callable )
3736
37+ # ---------------------------------------------------------------------------
38+ # Subprocess isolation helpers
39+ # ---------------------------------------------------------------------------
40+ # Tests that call setup_threads() must run in isolated subprocesses so that
41+ # the monkey-patches applied to threading.Thread / ThreadPoolExecutor do not
42+ # leak between tests. Each isolated test spawns a fresh ``python -c``
43+ # subprocess.
44+ # ---------------------------------------------------------------------------
45+
46+ _SCENARIO_TEMPLATE = """\
47+ import os, inspect, asyncio
48+ os.environ["BRAINTRUST_APP_URL"] = "https://www.braintrust.dev"
49+ os.environ.setdefault("OPENAI_API_KEY", "sk-test-dummy-api-key-for-vcr-tests")
50+ os.environ.setdefault("ANTHROPIC_API_KEY", "sk-ant-test-dummy-api-key-for-vcr-tests")
51+ os.environ.setdefault("MISTRAL_API_KEY", "mistral-test-dummy-api-key-for-vcr-tests")
52+ os.environ.setdefault("GOOGLE_API_KEY", os.environ.get("GEMINI_API_KEY", "your_google_api_key_here"))
53+ from braintrust import logger as _logger
54+ from braintrust.test_helpers import init_test_logger
55+ from braintrust.test_context import {fn_name} as _fn
56+ _logger._state.reset_parent_state()
57+ with _logger._internal_with_memory_background_logger() as _bgl:
58+ _tl = init_test_logger("test-context-project")
59+ if {instrument}:
60+ from braintrust.wrappers.threads import setup_threads
61+ setup_threads()
62+ if inspect.iscoroutinefunction(_fn):
63+ asyncio.run(_fn(_tl, _bgl))
64+ else:
65+ _fn(_tl, _bgl)
66+ _logger._state.reset_parent_state()
67+ """
68+
69+
70+ def _run_in_subprocess (code : str , label : str , timeout : int = 60 ) -> None :
71+ """Execute *code* in a fresh ``python -c`` subprocess and assert it exits cleanly."""
72+ try :
73+ result = subprocess .run (
74+ [sys .executable , "-c" , code ],
75+ capture_output = True ,
76+ text = True ,
77+ timeout = timeout ,
78+ )
79+ except subprocess .TimeoutExpired as exc :
80+ raise AssertionError (f"Isolated test { label } timed out after { timeout } s" ) from exc
81+ if result .returncode != 0 :
82+ raise AssertionError (
83+ f"Isolated test { label } failed (exit code { result .returncode } ):\n { result .stderr } \n { result .stdout } "
84+ )
85+
3886
3987def isolate (instrument : bool ) -> Callable [[F ], F ]:
4088 """
4189 Decorator for isolated context propagation tests.
4290
43- - Always runs in forked process (pytest-forked)
44- - If instrument=True: calls setup_threads() before test
91+ Runs each test in a separate subprocess for full process isolation,
92+ ensuring setup_threads() patches don't leak between tests.
93+
94+ - If instrument=True: calls setup_threads() before the test
4595 - If instrument=False: marks test as xfail (context loss expected)
4696 """
4797
4898 def decorator (fn : F ) -> F :
49- if asyncio .iscoroutinefunction (fn ):
50-
51- @functools .wraps (fn )
52- async def async_wrapper (* args , ** kwargs ):
53- if instrument :
54- setup_threads ()
55- return await fn (* args , ** kwargs )
99+ fn_name = fn .__name__
56100
57- wrapped = pytest .mark .forked (async_wrapper )
58- else :
101+ def isolated_test ():
102+ code = _SCENARIO_TEMPLATE .format (fn_name = fn_name , instrument = instrument )
103+ _run_in_subprocess (code , fn_name )
59104
60- @functools .wraps (fn )
61- def wrapper (* args , ** kwargs ):
62- if instrument :
63- setup_threads ()
64- return fn (* args , ** kwargs )
65-
66- wrapped = pytest .mark .forked (wrapper )
105+ isolated_test .__doc__ = fn .__doc__
67106
68107 if not instrument :
69- wrapped = pytest .mark .xfail (reason = "context lost without patch" )(wrapped )
70- return wrapped # type: ignore
108+ isolated_test = pytest .mark .xfail (reason = "context lost without patch" )(isolated_test )
109+
110+ return isolated_test # type: ignore
71111
72112 return decorator
73113
@@ -1266,21 +1306,39 @@ def failing_function():
12661306# ============================================================================
12671307
12681308
1269- @pytest .mark .forked
1270- def test_setup_threads_returns_true ():
1271- """setup_threads() returns True on success."""
1309+ def _setup_threads_returns_true_check ():
1310+ """Subprocess helper: verify setup_threads() returns True."""
1311+ from braintrust .wrappers .threads import setup_threads
1312+
12721313 result = setup_threads ()
12731314 assert result is True
12741315
12751316
1276- @pytest .mark .forked
1277- def test_setup_threads_idempotent ():
1278- """Calling setup_threads() multiple times is safe."""
1317+ def _setup_threads_idempotent_check ():
1318+ """Subprocess helper: verify setup_threads() is idempotent."""
1319+ from braintrust .wrappers .threads import setup_threads
1320+
12791321 result1 = setup_threads ()
12801322 result2 = setup_threads ()
12811323 assert result1 is True
12821324 assert result2 is True
12831325
12841326
1327+ def test_setup_threads_returns_true ():
1328+ """setup_threads() returns True on success."""
1329+ _run_in_subprocess (
1330+ "from braintrust.test_context import _setup_threads_returns_true_check; _setup_threads_returns_true_check()" ,
1331+ "test_setup_threads_returns_true" ,
1332+ )
1333+
1334+
1335+ def test_setup_threads_idempotent ():
1336+ """Calling setup_threads() multiple times is safe."""
1337+ _run_in_subprocess (
1338+ "from braintrust.test_context import _setup_threads_idempotent_check; _setup_threads_idempotent_check()" ,
1339+ "test_setup_threads_idempotent" ,
1340+ )
1341+
1342+
12851343if __name__ == "__main__" :
12861344 pytest .main ([__file__ , "-v" , "-s" ])
0 commit comments