From eef4cbf7d673b2b1d5bc2f82f2653fbf898cdf3b Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Thu, 15 Oct 2026 20:15:26 -0400 Subject: [PATCH] fix(tests): prevent background uploads from racing VCR replay --- py/src/braintrust/test_context.py | 17 ++++++----------- py/src/braintrust/test_helpers.py | 5 ++++- py/src/braintrust/test_logger.py | 7 +++++-- py/src/braintrust/test_span_components.py | 5 +++-- 4 files changed, 18 insertions(+), 16 deletions(-) diff --git a/py/src/braintrust/test_context.py b/py/src/braintrust/test_context.py index 8ee55f7be..156c61ab2 100644 --- a/py/src/braintrust/test_context.py +++ b/py/src/braintrust/test_context.py @@ -308,17 +308,12 @@ def blocking_work(): test_logger.flush() logs = with_memory_logger.pop() - # SURPRISING: Even to_thread() loses logger context (logger is a ContextVar too!) - # Only parent span is logged - # However, to_thread() DOES preserve span parent context - assert len(logs) >= 1 - - # If both spans logged (logger context preserved), verify parent chain - if len(logs) == 2: - parent_log = next(l for l in logs if l["span_attributes"]["name"] == "parent") - worker_log = next(l for l in logs if l["span_attributes"]["name"] == "to_thread_worker") - assert worker_log["root_span_id"] == parent_log["root_span_id"] - assert parent_log["span_id"] in worker_log.get("span_parents", []) + # The memory logger must capture worker-thread exports as well as the parent. + assert len(logs) == 2 + parent_log = next(l for l in logs if l["span_attributes"]["name"] == "parent") + worker_log = next(l for l in logs if l["span_attributes"]["name"] == "to_thread_worker") + assert worker_log["root_span_id"] == parent_log["root_span_id"] + assert parent_log["span_id"] in worker_log.get("span_parents", []) # ============================================================================ diff --git a/py/src/braintrust/test_helpers.py b/py/src/braintrust/test_helpers.py index f592677a8..109099986 100644 --- a/py/src/braintrust/test_helpers.py +++ b/py/src/braintrust/test_helpers.py @@ -60,9 +60,12 @@ def with_simulate_login(): @pytest.fixture -def with_memory_logger(): +def with_memory_logger(monkeypatch): logger._state.reset_parent_state() with logger._internal_with_memory_background_logger() as bgl: + # The override is thread-local. Worker threads need the same fallback + # so they cannot start HTTP uploads that outlive the test and race VCR. + monkeypatch.setattr(logger._state, "_global_bg_logger", LazyValue(lambda: bgl, use_mutex=False)) yield bgl # Clean up global state to prevent test contamination logger._state.reset_parent_state() diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index bd064c162..b2e7ea2a1 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -2484,7 +2484,7 @@ def reset_id_generator_state(monkeypatch): monkeypatch.delenv("BRAINTRUST_LEGACY_IDS", raising=False) -def test_span_with_otel_ids_export_import(reset_id_generator_state, monkeypatch): +def test_span_with_otel_ids_export_import(with_memory_logger, reset_id_generator_state, monkeypatch): """Test that actual Span objects with OTEL IDs can export and be used as parent context.""" init_test_logger(__name__) monkeypatch.setenv("BRAINTRUST_OTEL_COMPAT", "true") @@ -2502,7 +2502,7 @@ def test_span_with_otel_ids_export_import(reset_id_generator_state, monkeypatch) assert imported.root_span_id == span.root_span_id -def test_span_with_uuid_ids_share_root_span_id(reset_id_generator_state, monkeypatch): +def test_span_with_uuid_ids_share_root_span_id(with_memory_logger, reset_id_generator_state, monkeypatch): """Test that legacy UUID generators share span_id as root_span_id for backwards compatibility.""" monkeypatch.setenv("BRAINTRUST_LEGACY_IDS", "true") init_test_logger(__name__) @@ -2908,6 +2908,9 @@ def setUp(self): metadata_patch.start() self.addCleanup(metadata_patch.stop) self.state = logger.login_to_state(api_key=logger.TEST_API_KEY) + # These tests use a separate state, so the global memory override does + # not apply. Keep their span exports local too. + self.state._global_bg_logger = LazyValue(logger._MemoryBackgroundLogger, use_mutex=True) self.state.org_name = "test-org" api_client = MagicMock() client_patch = patch.object(self.state, "api_client", return_value=api_client) diff --git a/py/src/braintrust/test_span_components.py b/py/src/braintrust/test_span_components.py index 1d654cb56..fa12a94e1 100644 --- a/py/src/braintrust/test_span_components.py +++ b/py/src/braintrust/test_span_components.py @@ -10,6 +10,7 @@ from braintrust.id_gen import OTELIDGenerator from braintrust.span_identifier_v3 import SpanComponentsV3, SpanObjectTypeV3 from braintrust.span_identifier_v4 import SpanComponentsV4 +from braintrust.test_helpers import with_memory_logger # noqa: F401 # type: ignore[reportUnusedImport] class TestSpanComponentsV3: @@ -346,7 +347,7 @@ class TestExportFormatSelection: BRAINTRUST_LEGACY_IDS) serialize as V3. """ - def test_export_format_based_on_env_variable(self): + def test_export_format_based_on_env_variable(self, with_memory_logger): """Test that export format follows the legacy-UUID opt-out flag.""" from braintrust.test_helpers import init_test_logger, preserve_env_vars @@ -374,7 +375,7 @@ def test_export_format_based_on_env_variable(self): assert SpanComponentsV4.from_str(export_v3_mode) is not None assert SpanComponentsV4.from_str(export_v4_mode) is not None - def test_export_uses_v4_by_default(self): + def test_export_uses_v4_by_default(self, with_memory_logger): """Test that export uses V4 format by default (no env vars set).""" from braintrust.test_helpers import init_test_logger, preserve_env_vars