braintrust 0.4.3__py3-none-any.whl → 0.5.2__py3-none-any.whl

braintrust/test_logger.py

@@ -59,6 +59,33 @@ class TestInit(TestCase):
 
         assert str(cm.exception) == "Cannot open an experiment without specifying its name"
 
+    def test_init_with_dataset_id_only(self):
+        """Test that init accepts dataset={'id': '...'} parameter"""
+        # Test the logic that extracts dataset_id from the dict
+        from braintrust.logger import Dataset
+
+        # Test 1: dict with only id
+        dataset_dict = {"id": "dataset-id-123"}
+        assert isinstance(dataset_dict, dict)
+        assert not isinstance(dataset_dict, Dataset)
+        assert dataset_dict["id"] == "dataset-id-123"
+
+        # Test 2: full Dataset object has different behavior
+        # (We can't easily instantiate a Dataset here, but we can verify
+        # that the isinstance check distinguishes them)
+
+    def test_init_with_dataset_id_and_version(self):
+        """Test that init accepts dataset={'id': '...', 'version': '...'} parameter"""
+        # Test the logic that extracts both dataset_id and dataset_version from the dict
+        from braintrust.logger import Dataset
+
+        # Test: dict with id and version
+        dataset_dict = {"id": "dataset-id-123", "version": "v2"}
+        assert isinstance(dataset_dict, dict)
+        assert not isinstance(dataset_dict, Dataset)
+        assert dataset_dict["id"] == "dataset-id-123"
+        assert dataset_dict["version"] == "v2"
+
 
 class TestLogger(TestCase):
     def test_extract_attachments_no_op(self):
@@ -806,6 +833,16 @@ def test_span_project_id_logged_in(with_memory_logger, with_simulate_login):
     )
 
 
+def test_span_export_disables_cache(with_memory_logger):
+    """Test that span.export() disables the span cache."""
+    logger = init_test_logger(__name__)
+
+    with logger.start_span(name="test_span") as span:
+        # Exporting should disable the span cache
+        span.export()
+        assert logger.state.span_cache.disabled
+
+
 def test_span_project_name_logged_in(with_simulate_login, with_memory_logger):
     init_logger(project="test-project")
     span = logger.start_span(name="test-span")
@@ -902,11 +939,7 @@ def test_permalink_with_valid_span_logged_in(with_simulate_login, with_memory_lo
 
 @pytest.mark.asyncio
 async def test_span_link_in_async_context(with_simulate_login, with_memory_logger):
-    """Test that span.link() works correctly when called from within an async function.
-
-    This tests the bug where current_logger was a plain attribute instead of a ContextVar,
-    causing span.link() to return a noop link in async contexts even though the span was valid.
-    """
+    """Test that span.link() works correctly when called from within an async function."""
     import asyncio
 
     logger = init_logger(
@@ -939,6 +972,174 @@ async def test_span_link_in_async_context(with_simulate_login, with_memory_logge
     assert "test-project-id" in link
 
 
+@pytest.mark.asyncio
+async def test_current_logger_after_multiple_awaits(with_simulate_login, with_memory_logger):
+    """Test that current_logger() works after multiple await points."""
+    import asyncio
+
+    logger = init_logger(project="test-project", project_id="test-project-id")
+
+    async def check_logger_after_awaits():
+        assert braintrust.current_logger() is logger
+        await asyncio.sleep(0.01)
+        assert braintrust.current_logger() is logger
+        await asyncio.sleep(0.01)
+        assert braintrust.current_logger() is logger
+        return braintrust.current_logger()
+
+    result = await check_logger_after_awaits()
+    assert result is logger
+
+
+@pytest.mark.asyncio
+async def test_current_logger_in_async_generator(with_simulate_login, with_memory_logger):
+    """Test that current_logger() works within an async generator (yield)."""
+    import asyncio
+
+    logger = init_logger(project="test-project", project_id="test-project-id")
+
+    async def logger_generator():
+        for i in range(3):
+            await asyncio.sleep(0.01)
+            yield braintrust.current_logger()
+
+    results = []
+    async for log in logger_generator():
+        results.append(log)
+
+    assert len(results) == 3
+    assert all(r is logger for r in results)
+
+
+@pytest.mark.asyncio
+async def test_current_logger_in_separate_task(with_simulate_login, with_memory_logger):
+    """Test that current_logger() works in a separately created asyncio task."""
+    import asyncio
+
+    logger = init_logger(project="test-project", project_id="test-project-id")
+
+    async def get_logger_in_task():
+        await asyncio.sleep(0.01)
+        return braintrust.current_logger()
+
+    # Create a separate task
+    task = asyncio.create_task(get_logger_in_task())
+    result = await task
+
+    assert result is logger
+
+
+@pytest.mark.asyncio
+async def test_span_link_in_nested_async(with_simulate_login, with_memory_logger):
+    """Test that span.link() works in deeply nested async calls."""
+    import asyncio
+
+    logger = init_logger(project="test-project", project_id="test-project-id")
+    span = logger.start_span(name="test-span")
+
+    async def level3():
+        await asyncio.sleep(0.01)
+        return span.link()
+
+    async def level2():
+        await asyncio.sleep(0.01)
+        return await level3()
+
+    async def level1():
+        await asyncio.sleep(0.01)
+        return await level2()
+
+    link = await level1()
+    span.end()
+
+    assert link != "https://www.braintrust.dev/noop-span"
+    assert span._id in link
+
+
+def test_current_logger_in_thread(with_simulate_login, with_memory_logger):
+    """Test that current_logger() works correctly when called from a new thread.
+
+    Regression test: ContextVar values don't propagate to new threads,
+    so current_logger must be a plain attribute for thread access.
+    """
+    import threading
+
+    logger = init_logger(project="test-project", project_id="test-project-id")
+    assert braintrust.current_logger() is logger
+
+    thread_result = {}
+
+    def check_logger_in_thread():
+        thread_result["logger"] = braintrust.current_logger()
+
+    thread = threading.Thread(target=check_logger_in_thread)
+    thread.start()
+    thread.join()
+
+    assert thread_result["logger"] is logger
+
+
+def test_span_link_in_thread(with_simulate_login, with_memory_logger):
+    """Test that span.link() works correctly when called from a new thread.
+
+    The span should be able to generate a valid link even when link() is called
+    from a different thread than where the span was created.
+    """
+    import threading
+
+    logger = init_logger(project="test-project", project_id="test-project-id")
+    span = logger.start_span(name="test-span")
+
+    thread_result = {}
+
+    def get_link_in_thread():
+        # Call link() on the span directly (not via current_span() which uses ContextVar)
+        thread_result["link"] = span.link()
+
+    thread = threading.Thread(target=get_link_in_thread)
+    thread.start()
+    thread.join()
+    span.end()
+
+    # The link should NOT be the noop link
+    assert thread_result["link"] != "https://www.braintrust.dev/noop-span"
+    # The link should contain the span ID
+    assert span._id in thread_result["link"]
+
+
+@pytest.mark.asyncio
+async def test_current_logger_async_context_isolation(with_simulate_login, with_memory_logger):
+    """Test that different async contexts can have different loggers.
+
+    When a child task sets its own logger, it should not affect the parent context.
+    This ensures async context isolation via ContextVar.
+    """
+    import asyncio
+
+    parent_logger = init_logger(project="parent-project", project_id="parent-project-id")
+    assert braintrust.current_logger() is parent_logger
+
+    child_result = {}
+
+    async def child_task():
+        # Child initially inherits parent's logger
+        assert braintrust.current_logger() is parent_logger
+
+        # Child sets its own logger
+        child_logger = init_logger(project="child-project", project_id="child-project-id")
+        child_result["logger"] = braintrust.current_logger()
+        return child_logger
+
+    # Run child task
+    child_logger = await asyncio.create_task(child_task())
+
+    # Child should have seen its own logger
+    assert child_result["logger"] is child_logger
+
+    # Parent should still see parent logger (not affected by child)
+    assert braintrust.current_logger() is parent_logger
+
+
 def test_span_set_current(with_memory_logger):
     """Test that span.set_current() makes the span accessible via current_span()."""
     init_test_logger(__name__)
@@ -2434,6 +2635,95 @@ def test_logger_export_respects_otel_compat_enabled():
         assert version == 4, f"Expected V4 encoding (version=4), got version={version}"
 
 
+def test_register_otel_flush_callback():
+    """Test that register_otel_flush registers a callback correctly."""
+    import asyncio
+
+    from braintrust import register_otel_flush
+    from braintrust.logger import _internal_get_global_state
+    from braintrust.test_helpers import init_test_logger
+
+    init_test_logger(__name__)
+    state = _internal_get_global_state()
+
+    # Track if callback was invoked
+    callback_invoked = False
+
+    async def mock_flush():
+        nonlocal callback_invoked
+        callback_invoked = True
+
+    # Register the callback
+    register_otel_flush(mock_flush)
+
+    # Calling flush_otel should invoke the registered callback
+    asyncio.run(state.flush_otel())
+
+    assert callback_invoked is True
+
+
+def test_register_otel_flush_disables_span_cache():
+    """Test that register_otel_flush disables the span cache."""
+    from braintrust import register_otel_flush
+    from braintrust.logger import _internal_get_global_state
+    from braintrust.test_helpers import init_test_logger
+
+    init_test_logger(__name__)
+    state = _internal_get_global_state()
+
+    # Enable the cache (simulating what happens during eval)
+    state.span_cache.start()
+    assert state.span_cache.disabled is False
+
+    async def mock_flush():
+        pass
+
+    # Register OTEL flush
+    register_otel_flush(mock_flush)
+
+    # Cache should now be disabled
+    assert state.span_cache.disabled is True
+
+
+def test_flush_otel_noop_when_no_callback():
+    """Test that flush_otel is a no-op when no callback is registered."""
+    import asyncio
+
+    from braintrust.logger import _internal_get_global_state
+    from braintrust.test_helpers import init_test_logger
+
+    init_test_logger(__name__)
+    state = _internal_get_global_state()
+
+    # Should not throw even with no callback registered
+    asyncio.run(state.flush_otel())
+
+
+def test_register_otel_flush_permanently_disables_cache():
+    """Test that register_otel_flush permanently disables the cache."""
+    from braintrust import register_otel_flush
+    from braintrust.logger import _internal_get_global_state
+    from braintrust.test_helpers import init_test_logger
+
+    init_test_logger(__name__)
+    state = _internal_get_global_state()
+
+    # Enable the cache
+    state.span_cache.start()
+    assert state.span_cache.disabled is False
+
+    async def mock_flush():
+        pass
+
+    # Register OTEL flush
+    register_otel_flush(mock_flush)
+    assert state.span_cache.disabled is True
+
+    # Try to start again - should still be disabled because of explicit disable
+    state.span_cache.start()
+    assert state.span_cache.disabled is True
+
+
 class TestJSONAttachment(TestCase):
     def test_create_attachment_from_json_data(self):
         """Test creating an attachment from JSON data."""

braintrust/test_span_cache.py

@@ -0,0 +1,344 @@
+"""Tests for SpanCache (disk-based cache)."""
+
+
+from braintrust.span_cache import CachedSpan, SpanCache
+
+
+def test_span_cache_write_and_read():
+    """Test storing and retrieving spans by rootSpanId."""
+    cache = SpanCache()
+    cache.start()  # Start for testing (cache is disabled by default)
+
+    root_span_id = "root-123"
+    span1 = CachedSpan(
+        span_id="span-1",
+        input={"text": "hello"},
+        output={"response": "world"},
+    )
+    span2 = CachedSpan(
+        span_id="span-2",
+        input={"text": "foo"},
+        output={"response": "bar"},
+    )
+
+    cache.queue_write(root_span_id, span1.span_id, span1)
+    cache.queue_write(root_span_id, span2.span_id, span2)
+
+    spans = cache.get_by_root_span_id(root_span_id)
+    assert spans is not None
+    assert len(spans) == 2
+
+    span_ids = {s.span_id for s in spans}
+    assert "span-1" in span_ids
+    assert "span-2" in span_ids
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_return_none_for_unknown():
+    """Test that unknown rootSpanId returns None."""
+    cache = SpanCache()
+    cache.start()
+
+    spans = cache.get_by_root_span_id("nonexistent")
+    assert spans is None
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_merge_on_duplicate_writes():
+    """Test that subsequent writes to same spanId merge data."""
+    cache = SpanCache()
+    cache.start()
+
+    root_span_id = "root-123"
+    span_id = "span-1"
+
+    cache.queue_write(
+        root_span_id,
+        span_id,
+        CachedSpan(span_id=span_id, input={"text": "hello"}),
+    )
+
+    cache.queue_write(
+        root_span_id,
+        span_id,
+        CachedSpan(span_id=span_id, output={"response": "world"}),
+    )
+
+    spans = cache.get_by_root_span_id(root_span_id)
+    assert spans is not None
+    assert len(spans) == 1
+    assert spans[0].span_id == span_id
+    assert spans[0].input == {"text": "hello"}
+    assert spans[0].output == {"response": "world"}
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_merge_metadata():
+    """Test that metadata objects are merged."""
+    cache = SpanCache()
+    cache.start()
+
+    root_span_id = "root-123"
+    span_id = "span-1"
+
+    cache.queue_write(
+        root_span_id,
+        span_id,
+        CachedSpan(span_id=span_id, metadata={"key1": "value1"}),
+    )
+
+    cache.queue_write(
+        root_span_id,
+        span_id,
+        CachedSpan(span_id=span_id, metadata={"key2": "value2"}),
+    )
+
+    spans = cache.get_by_root_span_id(root_span_id)
+    assert spans is not None
+    assert spans[0].metadata == {"key1": "value1", "key2": "value2"}
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_has():
+    """Test the has() method."""
+    cache = SpanCache()
+    cache.start()
+
+    cache.queue_write("root-123", "span-1", CachedSpan(span_id="span-1"))
+    assert cache.has("root-123") is True
+    assert cache.has("nonexistent") is False
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_clear():
+    """Test clearing spans for a specific rootSpanId."""
+    cache = SpanCache()
+    cache.start()
+
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    cache.queue_write("root-2", "span-2", CachedSpan(span_id="span-2"))
+
+    cache.clear("root-1")
+
+    assert cache.has("root-1") is False
+    assert cache.has("root-2") is True
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_clear_all():
+    """Test clearing all cached spans."""
+    cache = SpanCache()
+    cache.start()
+
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    cache.queue_write("root-2", "span-2", CachedSpan(span_id="span-2"))
+
+    cache.clear_all()
+
+    assert cache.size == 0
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_size():
+    """Test the size property."""
+    cache = SpanCache()
+    cache.start()
+
+    assert cache.size == 0
+
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    assert cache.size == 1
+
+    cache.queue_write("root-1", "span-2", CachedSpan(span_id="span-2"))  # Same root
+    assert cache.size == 1
+
+    cache.queue_write("root-2", "span-3", CachedSpan(span_id="span-3"))  # Different root
+    assert cache.size == 2
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_dispose():
+    """Test that dispose cleans up and allows reuse."""
+    cache = SpanCache()
+    cache.start()
+
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    assert cache.size == 1
+
+    # Stop first to decrement refcount, then dispose
+    cache.stop()
+    cache.dispose()
+
+    assert cache.size == 0
+    assert cache.has("root-1") is False
+
+    # Should be able to write again after dispose (if we start again)
+    cache.start()
+    cache.queue_write("root-2", "span-2", CachedSpan(span_id="span-2"))
+    assert cache.size == 1
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_disable():
+    """Test that disable() prevents writes."""
+    cache = SpanCache()
+    cache.start()
+
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    assert cache.size == 1
+
+    cache.disable()
+
+    # Writes after disable should be no-ops
+    cache.queue_write("root-2", "span-2", CachedSpan(span_id="span-2"))
+    assert cache.size == 1  # Still 1, not 2
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_disabled_getter():
+    """Test the disabled property."""
+    # Cache is disabled by default until start() is called
+    cache = SpanCache()
+    assert cache.disabled is True
+
+    cache.start()
+    assert cache.disabled is False
+
+    cache.disable()
+    assert cache.disabled is True
+
+    cache.dispose()
+
+
+def test_span_cache_disabled_from_constructor():
+    """Test that cache can be disabled via constructor."""
+    cache = SpanCache(disabled=True)
+    assert cache.disabled is True
+
+    # Writes should be no-ops
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    assert cache.size == 0
+    assert cache.get_by_root_span_id("root-1") is None
+
+    cache.dispose()
+
+
+def test_span_cache_start_stop_lifecycle():
+    """Test that stop() allows start() to work again."""
+    cache = SpanCache()
+
+    # Initially disabled by default
+    assert cache.disabled is True
+
+    # Start for first "eval"
+    cache.start()
+    assert cache.disabled is False
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    assert cache.size == 1
+
+    # Stop after first "eval"
+    cache.stop()
+    cache.dispose()
+    assert cache.disabled is True
+
+    # Start for second "eval" - should work!
+    cache.start()
+    assert cache.disabled is False
+    cache.queue_write("root-2", "span-2", CachedSpan(span_id="span-2"))
+    assert cache.size == 1
+
+    cache.stop()
+    cache.dispose()
+
+
+def test_span_cache_disable_prevents_start():
+    """Test that disable() prevents start() from working."""
+    cache = SpanCache()
+
+    # Simulate disable being called
+    cache.disable()
+    assert cache.disabled is True
+
+    # start() should be a no-op after disable()
+    cache.start()
+    assert cache.disabled is True
+
+    # Writes should still be no-ops
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    assert cache.size == 0
+
+    cache.dispose()
+
+
+def test_span_cache_parallel_eval_refcount():
+    """Test reference counting for parallel evals."""
+    cache = SpanCache()
+
+    # Simulate two evals starting
+    cache.start()  # Eval 1
+    assert cache.disabled is False
+
+    cache.start()  # Eval 2
+    assert cache.disabled is False
+
+    # Write data from both evals
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    cache.queue_write("root-2", "span-2", CachedSpan(span_id="span-2"))
+    assert cache.size == 2
+
+    # Eval 1 finishes first
+    cache.dispose()  # Should NOT dispose (refcount = 2)
+    cache.stop()  # Decrements to 1
+
+    # Cache should still be enabled and data intact
+    assert cache.disabled is False
+    assert cache.size == 2
+    assert cache.get_by_root_span_id("root-1") is not None
+    assert cache.get_by_root_span_id("root-2") is not None
+
+    # Eval 2 finishes
+    cache.dispose()  # Should NOT dispose yet (refcount = 1)
+    cache.stop()  # Decrements to 0, disables cache
+
+    # Now cache should be disabled
+    assert cache.disabled is True
+
+    # Final dispose should now work
+    cache.dispose()  # NOW it disposes (refcount = 0)
+    assert cache.size == 0
+
+
+def test_span_cache_refcount_underflow():
+    """Test that refcount handles underflow gracefully."""
+    cache = SpanCache()
+
+    # Call stop without start
+    cache.stop()
+
+    # Should work normally after
+    cache.start()
+    cache.queue_write("root-1", "span-1", CachedSpan(span_id="span-1"))
+    assert cache.size == 1
+
+    cache.stop()
+    cache.dispose()