mseep-agentops 0.4.18__py3-none-any.whl → 0.4.23__py3-none-any.whl

tests/unit/test_context_manager.py

@@ -1,777 +0,0 @@
-import unittest
-from unittest.mock import Mock, patch
-import threading
-import time
-import asyncio
-
-from agentops import start_trace
-from agentops.sdk.core import TraceContext
-from opentelemetry.trace import StatusCode
-
-
-class TestContextManager(unittest.TestCase):
-    """Test the context manager functionality of TraceContext"""
-
-    def test_trace_context_has_context_manager_methods(self):
-        """Test that TraceContext has __enter__ and __exit__ methods"""
-        # TraceContext should have context manager protocol methods
-        assert hasattr(TraceContext, "__enter__")
-        assert hasattr(TraceContext, "__exit__")
-
-    @patch("agentops.sdk.core.tracer")
-    def test_trace_context_enter_returns_self(self, mock_tracer):
-        """Test that __enter__ returns the TraceContext instance"""
-        mock_span = Mock()
-        mock_token = Mock()
-        trace_context = TraceContext(span=mock_span, token=mock_token)
-
-        # __enter__ should return self
-        result = trace_context.__enter__()
-        assert result is trace_context
-
-    @patch("agentops.sdk.core.tracer")
-    def test_trace_context_exit_calls_end_trace(self, mock_tracer):
-        """Test that __exit__ calls end_trace with appropriate state"""
-        mock_span = Mock()
-        mock_token = Mock()
-        trace_context = TraceContext(span=mock_span, token=mock_token)
-
-        # Test normal exit (no exception)
-        trace_context.__exit__(None, None, None)
-        mock_tracer.end_trace.assert_called_once_with(trace_context, StatusCode.OK)
-
-    @patch("agentops.sdk.core.tracer")
-    def test_trace_context_exit_with_exception_sets_error_state(self, mock_tracer):
-        """Test that __exit__ sets ERROR state when exception occurs"""
-        mock_span = Mock()
-        mock_token = Mock()
-        trace_context = TraceContext(span=mock_span, token=mock_token)
-
-        # Test exit with exception
-        mock_tracer.reset_mock()
-        exc_type = ValueError
-        exc_val = ValueError("test error")
-        exc_tb = None
-
-        trace_context.__exit__(exc_type, exc_val, exc_tb)
-        mock_tracer.end_trace.assert_called_once_with(trace_context, StatusCode.ERROR)
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_context_manager_usage_pattern(self, mock_agentops_tracer, mock_core_tracer):
-        """Test using start_trace as a context manager"""
-        # Create a mock TraceContext
-        mock_span = Mock()
-        mock_token = Mock()
-        mock_trace_context = TraceContext(span=mock_span, token=mock_token)
-
-        # Mock the tracer's start_trace method to return our TraceContext
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.return_value = mock_trace_context
-
-        # Use as context manager
-        with start_trace("test_trace") as trace:
-            assert trace is mock_trace_context
-            assert trace.span is mock_span
-            assert trace.token is mock_token
-
-        # Verify start_trace was called
-        mock_agentops_tracer.start_trace.assert_called_once_with(trace_name="test_trace", tags=None)
-        # Verify end_trace was called
-        mock_core_tracer.end_trace.assert_called_once_with(mock_trace_context, StatusCode.OK)
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_context_manager_with_exception(self, mock_agentops_tracer, mock_core_tracer):
-        """Test context manager handles exceptions properly"""
-        # Create a mock TraceContext
-        mock_span = Mock()
-        mock_token = Mock()
-        mock_trace_context = TraceContext(span=mock_span, token=mock_token)
-
-        # Mock the tracer's start_trace method
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.return_value = mock_trace_context
-
-        # Test exception handling
-        with self.assertRaises(ValueError):
-            with start_trace("test_trace"):
-                raise ValueError("Test exception")
-
-        # Verify end_trace was called with ERROR state
-        mock_core_tracer.end_trace.assert_called_once_with(mock_trace_context, StatusCode.ERROR)
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    @patch("agentops.init")
-    def test_start_trace_auto_initializes_if_needed(self, mock_init, mock_agentops_tracer, mock_core_tracer):
-        """Test that start_trace attempts to initialize SDK if not initialized"""
-        # First call: SDK not initialized
-        mock_agentops_tracer.initialized = False
-
-        # After init() is called, set initialized to True
-        def set_initialized():
-            mock_agentops_tracer.initialized = True
-
-        mock_init.side_effect = set_initialized
-
-        # Create a mock TraceContext for when start_trace is called after init
-        mock_span = Mock()
-        mock_token = Mock()
-        mock_trace_context = TraceContext(span=mock_span, token=mock_token)
-        mock_agentops_tracer.start_trace.return_value = mock_trace_context
-
-        # Call start_trace
-        result = start_trace("test_trace")
-
-        # Verify init was called
-        mock_init.assert_called_once()
-        # Verify start_trace was called on tracer
-        mock_agentops_tracer.start_trace.assert_called_once_with(trace_name="test_trace", tags=None)
-        assert result is mock_trace_context
-
-    def test_no_wrapper_classes_needed(self):
-        """Test that we don't need wrapper classes - TraceContext is the context manager"""
-        # TraceContext itself implements the context manager protocol
-        # No need for TraceContextManager wrapper
-        mock_span = Mock()
-        mock_token = Mock()
-        trace_context = TraceContext(span=mock_span, token=mock_token)
-
-        # Can use directly as context manager
-        assert hasattr(trace_context, "__enter__")
-        assert hasattr(trace_context, "__exit__")
-        assert callable(trace_context.__enter__)
-        assert callable(trace_context.__exit__)
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_parallel_traces_independence(self, mock_agentops_tracer, mock_core_tracer):
-        """Test that multiple traces can run in parallel independently"""
-        # Create mock TraceContexts
-        mock_trace1 = TraceContext(span=Mock(), token=Mock())
-        mock_trace2 = TraceContext(span=Mock(), token=Mock())
-
-        # Mock the tracer to return different traces
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.side_effect = [mock_trace1, mock_trace2]
-
-        # Start two traces
-        trace1 = start_trace("trace1")
-        trace2 = start_trace("trace2")
-
-        # They should be different instances
-        assert trace1 is not trace2
-        assert trace1 is mock_trace1
-        assert trace2 is mock_trace2
-
-        # End them independently using context manager protocol
-        trace1.__exit__(None, None, None)
-        trace2.__exit__(None, None, None)
-
-        # Verify both were ended
-        assert mock_core_tracer.end_trace.call_count == 2
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_nested_context_managers_create_parallel_traces(self, mock_agentops_tracer, mock_core_tracer):
-        """Test that nested context managers create parallel traces, not parent-child"""
-        # Create mock TraceContexts
-        mock_outer = TraceContext(span=Mock(), token=Mock())
-        mock_inner = TraceContext(span=Mock(), token=Mock())
-
-        # Mock the tracer
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.side_effect = [mock_outer, mock_inner]
-
-        # Use nested context managers
-        with start_trace("outer_trace") as outer:
-            assert outer is mock_outer
-            with start_trace("inner_trace") as inner:
-                assert inner is mock_inner
-                assert inner is not outer
-                # Both traces are active
-                assert mock_agentops_tracer.start_trace.call_count == 2
-
-        # Verify both were ended
-        assert mock_core_tracer.end_trace.call_count == 2
-        # Inner trace ended first, then outer
-        calls = mock_core_tracer.end_trace.call_args_list
-        assert calls[0][0][0] is mock_inner
-        assert calls[1][0][0] is mock_outer
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_exception_in_nested_traces(self, mock_agentops_tracer, mock_core_tracer):
-        """Test exception handling in nested traces"""
-        # Create mock TraceContexts
-        mock_outer = TraceContext(span=Mock(), token=Mock())
-        mock_inner = TraceContext(span=Mock(), token=Mock())
-
-        # Mock the tracer
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.side_effect = [mock_outer, mock_inner]
-
-        # Test exception in inner trace
-        with self.assertRaises(ValueError):
-            with start_trace("outer_trace"):
-                with start_trace("inner_trace"):
-                    raise ValueError("Inner exception")
-
-        # Both traces should be ended with ERROR state
-        assert mock_core_tracer.end_trace.call_count == 2
-        calls = mock_core_tracer.end_trace.call_args_list
-        # Inner trace ended with ERROR
-        assert calls[0][0][0] is mock_inner
-        assert calls[0][0][1] == StatusCode.ERROR
-        # Outer trace also ended with ERROR (exception propagated)
-        assert calls[1][0][0] is mock_outer
-        assert calls[1][0][1] == StatusCode.ERROR
-
-    @patch("agentops.sdk.core.tracer")
-    def test_trace_context_attributes_access(self, mock_tracer):
-        """Test accessing span and token attributes of TraceContext"""
-        mock_span = Mock()
-        mock_token = Mock()
-        trace_context = TraceContext(span=mock_span, token=mock_token)
-
-        # Direct attribute access
-        assert trace_context.span is mock_span
-        assert trace_context.token is mock_token
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_multiple_exceptions_in_sequence(self, mock_agentops_tracer, mock_core_tracer):
-        """Test handling multiple exceptions in sequence"""
-        # Mock the tracer
-        mock_agentops_tracer.initialized = True
-
-        # Create different mock traces for each attempt
-        mock_traces = [TraceContext(span=Mock(), token=Mock()) for _ in range(3)]
-        mock_agentops_tracer.start_trace.side_effect = mock_traces
-
-        # Multiple traces with exceptions
-        for i in range(3):
-            with self.assertRaises(RuntimeError):
-                with start_trace(f"trace_{i}"):
-                    raise RuntimeError(f"Error {i}")
-
-        # All should be ended with ERROR state
-        assert mock_core_tracer.end_trace.call_count == 3
-        for i, call in enumerate(mock_core_tracer.end_trace.call_args_list):
-            assert call[0][0] is mock_traces[i]
-            assert call[0][1] == StatusCode.ERROR
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_trace_with_tags_dict(self, mock_agentops_tracer, mock_core_tracer):
-        """Test starting trace with tags as dictionary"""
-        # Create a mock TraceContext
-        mock_trace = TraceContext(span=Mock(), token=Mock())
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.return_value = mock_trace
-
-        tags = {"environment": "test", "version": "1.0"}
-        with start_trace("tagged_trace", tags=tags) as trace:
-            assert trace is mock_trace
-
-        # Verify tags were passed
-        mock_agentops_tracer.start_trace.assert_called_once_with(trace_name="tagged_trace", tags=tags)
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_trace_with_tags_list(self, mock_agentops_tracer, mock_core_tracer):
-        """Test starting trace with tags as list"""
-        # Create a mock TraceContext
-        mock_trace = TraceContext(span=Mock(), token=Mock())
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.return_value = mock_trace
-
-        tags = ["test", "v1.0", "experimental"]
-        with start_trace("tagged_trace", tags=tags) as trace:
-            assert trace is mock_trace
-
-        # Verify tags were passed
-        mock_agentops_tracer.start_trace.assert_called_once_with(trace_name="tagged_trace", tags=tags)
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_trace_context_manager_thread_safety(self, mock_agentops_tracer, mock_core_tracer):
-        """Test that context managers work correctly in multi-threaded environment"""
-        # Mock the tracer
-        mock_agentops_tracer.initialized = True
-
-        # Create unique traces for each thread
-        thread_traces = {}
-        trace_lock = threading.Lock()
-
-        def create_trace(trace_name=None, tags=None, **kwargs):
-            trace = TraceContext(span=Mock(), token=Mock())
-            with trace_lock:
-                thread_traces[threading.current_thread().ident] = trace
-            return trace
-
-        mock_agentops_tracer.start_trace.side_effect = create_trace
-
-        results = []
-        errors = []
-
-        def worker(thread_id):
-            try:
-                with start_trace(f"thread_{thread_id}_trace") as trace:
-                    # Each thread should get its own trace
-                    results.append((thread_id, trace))
-                    time.sleep(0.01)  # Simulate some work
-            except Exception as e:
-                errors.append((thread_id, str(e)))
-
-        # Start multiple threads
-        threads = []
-        for i in range(5):
-            t = threading.Thread(target=worker, args=(i,))
-            threads.append(t)
-            t.start()
-
-        # Wait for all threads
-        for t in threads:
-            t.join()
-
-        # Check results
-        assert len(errors) == 0, f"Errors in threads: {errors}"
-        assert len(results) == 5
-
-        # Each thread should have gotten a unique trace
-        traces = [r[1] for r in results]
-        assert len(set(id(t) for t in traces)) == 5  # All unique
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_context_manager_with_early_return(self, mock_agentops_tracer, mock_core_tracer):
-        """Test that context manager properly cleans up with early return"""
-        # Create a mock TraceContext
-        mock_trace = TraceContext(span=Mock(), token=Mock())
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.return_value = mock_trace
-
-        def function_with_early_return():
-            with start_trace("early_return_trace"):
-                if True:  # Early return condition
-                    return "early"
-                return "normal"
-
-        result = function_with_early_return()
-        assert result == "early"
-
-        # Verify trace was still ended properly
-        mock_core_tracer.end_trace.assert_called_once_with(mock_trace, StatusCode.OK)
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_context_manager_with_finally_block(self, mock_agentops_tracer, mock_core_tracer):
-        """Test context manager with try-finally block"""
-        # Create a mock TraceContext
-        mock_trace = TraceContext(span=Mock(), token=Mock())
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.return_value = mock_trace
-
-        finally_executed = False
-
-        try:
-            with start_trace("finally_trace"):
-                try:
-                    raise ValueError("Test")
-                finally:
-                    finally_executed = True
-        except ValueError:
-            pass
-
-        assert finally_executed
-        # Trace should be ended with ERROR due to exception
-        mock_core_tracer.end_trace.assert_called_once_with(mock_trace, StatusCode.ERROR)
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_backwards_compatibility_existing_patterns(self, mock_agentops_tracer, mock_core_tracer):
-        """Test that existing usage patterns continue to work"""
-        # Create mock traces
-        mock_traces = [TraceContext(span=Mock(), token=Mock()) for _ in range(3)]
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.side_effect = mock_traces
-
-        # Pattern 1: Basic context manager
-        with start_trace("basic") as trace:
-            assert trace is mock_traces[0]
-
-        # Pattern 2: Manual start/end using context manager protocol
-        trace = start_trace("manual")
-        assert trace is mock_traces[1]
-        trace.__exit__(None, None, None)  # Use context manager exit instead of end_trace
-
-        # Pattern 3: With tags
-        with start_trace("tagged", tags=["production", "v2"]) as trace:
-            assert trace is mock_traces[2]
-
-        # All patterns should work
-        assert mock_agentops_tracer.start_trace.call_count == 3
-        assert mock_core_tracer.end_trace.call_count == 3
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_edge_case_none_trace_context(self, mock_agentops_tracer, mock_core_tracer):
-        """Test handling when start_trace returns None"""
-        # Mock SDK not initialized and init fails
-        mock_agentops_tracer.initialized = False
-
-        # When start_trace is called on uninitialized tracer, it returns None
-        with patch("agentops.init") as mock_init:
-            mock_init.side_effect = Exception("Init failed")
-
-            result = start_trace("test_trace")
-            assert result is None
-
-        # Verify start_trace was not called on tracer (since init failed)
-        mock_agentops_tracer.start_trace.assert_not_called()
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_edge_case_tracing_core_not_initialized(self, mock_agentops_tracer, mock_core_tracer):
-        """Test behavior when global tracer is not initialized"""
-        mock_agentops_tracer.initialized = False
-
-        # Mock init to succeed but tracer still not initialized
-        with patch("agentops.init") as mock_init:
-            mock_init.return_value = None  # init succeeds but doesn't set initialized
-
-            result = start_trace("test")
-            assert result is None
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_edge_case_exception_in_exit_method(self, mock_agentops_tracer, mock_core_tracer):
-        """Test handling when exception occurs in __exit__ method"""
-        # Create a mock TraceContext
-        mock_trace = TraceContext(span=Mock(), token=Mock())
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.return_value = mock_trace
-
-        # Make end_trace raise an exception
-        mock_core_tracer.end_trace.side_effect = RuntimeError("End trace failed")
-
-        # The exception in __exit__ should be suppressed
-        with start_trace("exception_in_exit"):
-            pass  # Should not raise
-
-        # Verify end_trace was attempted
-        mock_core_tracer.end_trace.assert_called_once()
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_performance_many_sequential_traces(self, mock_agentops_tracer, mock_core_tracer):
-        """Test performance with many sequential traces"""
-        # Mock the tracer
-        mock_agentops_tracer.initialized = True
-
-        # Create traces on demand
-        def create_trace(trace_name=None, tags=None, **kwargs):
-            return TraceContext(span=Mock(), token=Mock())
-
-        mock_agentops_tracer.start_trace.side_effect = create_trace
-
-        # Create many traces sequentially
-        start_time = time.time()
-        for i in range(100):
-            with start_trace(f"trace_{i}") as trace:
-                assert trace is not None
-                assert trace.span is not None
-
-        elapsed = time.time() - start_time
-
-        # Should complete reasonably quickly (< 1 second for 100 traces)
-        assert elapsed < 1.0, f"Too slow: {elapsed:.2f}s for 100 traces"
-
-        # Verify all traces were started and ended
-        assert mock_agentops_tracer.start_trace.call_count == 100
-        assert mock_core_tracer.end_trace.call_count == 100
-
-    @patch("agentops.sdk.core.tracer")
-    def test_trace_context_state_management(self, mock_tracer):
-        """Test that TraceContext properly manages its internal state"""
-        mock_span = Mock()
-        mock_token = Mock()
-        trace_context = TraceContext(span=mock_span, token=mock_token)
-
-        # Initial state
-        assert trace_context.span is mock_span
-        assert trace_context.token is mock_token
-
-        # Enter context
-        result = trace_context.__enter__()
-        assert result is trace_context
-
-        # Exit context normally
-        trace_context.__exit__(None, None, None)
-        mock_tracer.end_trace.assert_called_once_with(trace_context, StatusCode.OK)
-
-        # State should remain accessible after exit
-        assert trace_context.span is mock_span
-        assert trace_context.token is mock_token
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_context_manager_with_async_context(self, mock_agentops_tracer, mock_core_tracer):
-        """Test context manager works in async context"""
-        # Create a mock TraceContext
-        mock_trace = TraceContext(span=Mock(), token=Mock())
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.return_value = mock_trace
-
-        async def async_function():
-            with start_trace("async_context") as trace:
-                assert trace is mock_trace
-                await asyncio.sleep(0.01)
-                return "done"
-
-        # Run async function
-        result = asyncio.run(async_function())
-        assert result == "done"
-
-        # Verify trace was properly managed
-        mock_agentops_tracer.start_trace.assert_called_once_with(trace_name="async_context", tags=None)
-        mock_core_tracer.end_trace.assert_called_once_with(mock_trace, StatusCode.OK)
-
-
-class TestContextManagerBackwardCompatibility(unittest.TestCase):
-    """Test backward compatibility for context manager usage"""
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_existing_code_patterns_still_work(self, mock_agentops_tracer, mock_core_tracer):
-        """Test that code using the old patterns still works"""
-        # Create mock traces - need more than 3 for this test
-        mock_traces = [TraceContext(span=Mock(), token=Mock()) for _ in range(5)]
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.side_effect = mock_traces
-
-        # Old pattern 1: Simple context manager
-        with start_trace("basic") as trace:
-            # Should work without changes
-            assert trace.span is not None
-
-        # Old pattern 2: Context manager with exception handling
-        try:
-            with start_trace("with_error") as trace:
-                raise ValueError("test")
-        except ValueError:
-            pass
-
-        # Old pattern 3: Nested traces
-        with start_trace("outer") as outer:
-            with start_trace("inner") as inner:
-                assert outer is not inner
-
-        # All should work - 4 calls total (basic, with_error, outer, inner)
-        assert mock_agentops_tracer.start_trace.call_count == 4
-        assert mock_core_tracer.end_trace.call_count == 4
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_api_compatibility(self, mock_agentops_tracer, mock_core_tracer):
-        """Test that the API remains compatible"""
-        # Create mock TraceContexts for each call
-        mock_traces = [TraceContext(span=Mock(), token=Mock()) for _ in range(3)]
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.side_effect = mock_traces
-
-        # Test function signatures
-        # start_trace(trace_name, tags=None)
-        trace1 = start_trace("test1")
-        assert trace1 is mock_traces[0]
-
-        trace2 = start_trace("test2", tags=["tag1", "tag2"])
-        assert trace2 is mock_traces[1]
-
-        trace3 = start_trace("test3", tags={"key": "value"})
-        assert trace3 is mock_traces[2]
-
-        # Use context manager protocol to end traces
-        trace1.__exit__(None, None, None)
-        trace2.__exit__(ValueError, ValueError("test"), None)
-        trace3.__exit__(None, None, None)
-
-        # All calls should work
-        assert mock_agentops_tracer.start_trace.call_count == 3
-        assert mock_core_tracer.end_trace.call_count == 3
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_return_type_compatibility(self, mock_agentops_tracer, mock_core_tracer):
-        """Test that return types are compatible with existing code"""
-        mock_span = Mock()
-        mock_token = Mock()
-        mock_trace = TraceContext(span=mock_span, token=mock_token)
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.return_value = mock_trace
-
-        # start_trace returns TraceContext (or None)
-        trace = start_trace("test")
-        assert isinstance(trace, TraceContext)
-        assert hasattr(trace, "span")
-        assert hasattr(trace, "token")
-        assert hasattr(trace, "__enter__")
-        assert hasattr(trace, "__exit__")
-
-        # Can be used as context manager
-        with trace:
-            pass
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_context_manager_with_keyboard_interrupt(self, mock_agentops_tracer, mock_core_tracer):
-        """Test context manager handles KeyboardInterrupt properly"""
-        # Create a mock TraceContext
-        mock_trace = TraceContext(span=Mock(), token=Mock())
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.return_value = mock_trace
-
-        # Test KeyboardInterrupt handling
-        with self.assertRaises(KeyboardInterrupt):
-            with start_trace("keyboard_interrupt_trace"):
-                raise KeyboardInterrupt()
-
-        # Verify end_trace was called with ERROR state
-        mock_core_tracer.end_trace.assert_called_once_with(mock_trace, StatusCode.ERROR)
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_context_manager_with_system_exit(self, mock_agentops_tracer, mock_core_tracer):
-        """Test context manager handles SystemExit properly"""
-        # Create a mock TraceContext
-        mock_trace = TraceContext(span=Mock(), token=Mock())
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.return_value = mock_trace
-
-        # Test SystemExit handling
-        with self.assertRaises(SystemExit):
-            with start_trace("system_exit_trace"):
-                raise SystemExit(1)
-
-        # Verify end_trace was called with ERROR state
-        mock_core_tracer.end_trace.assert_called_once_with(mock_trace, StatusCode.ERROR)
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_context_manager_in_generator_function(self, mock_agentops_tracer, mock_core_tracer):
-        """Test context manager works correctly in generator functions"""
-        # Create mock traces
-        mock_traces = [TraceContext(span=Mock(), token=Mock()) for _ in range(3)]
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.side_effect = mock_traces
-
-        def trace_generator():
-            with start_trace("generator_trace"):
-                yield 1
-                yield 2
-                yield 3
-
-        # Consume the generator
-        results = list(trace_generator())
-        assert results == [1, 2, 3]
-
-        # Verify trace was properly managed
-        mock_agentops_tracer.start_trace.assert_called_once()
-        mock_core_tracer.end_trace.assert_called_once()
-
-    @patch("agentops.sdk.core.tracer")
-    def test_context_manager_exit_return_value(self, mock_tracer):
-        """Test that __exit__ returns None (doesn't suppress exceptions)"""
-        mock_span = Mock()
-        mock_token = Mock()
-        trace_context = TraceContext(span=mock_span, token=mock_token)
-
-        # __exit__ should return None (or falsy) to not suppress exceptions
-        result = trace_context.__exit__(None, None, None)
-        assert result is None or not result
-
-        # Also with exception
-        result = trace_context.__exit__(ValueError, ValueError("test"), None)
-        assert result is None or not result
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_context_manager_with_very_large_data(self, mock_agentops_tracer, mock_core_tracer):
-        """Test context manager with very large trace names and tags"""
-        # Create a mock TraceContext
-        mock_trace = TraceContext(span=Mock(), token=Mock())
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.return_value = mock_trace
-
-        # Very large trace name and tags
-        large_trace_name = "x" * 10000
-        large_tags = {f"key_{i}": f"value_{i}" * 100 for i in range(100)}
-
-        with start_trace(large_trace_name, tags=large_tags) as trace:
-            assert trace is mock_trace
-
-        # Should handle large data without issues
-        mock_agentops_tracer.start_trace.assert_called_once()
-        args, kwargs = mock_agentops_tracer.start_trace.call_args
-        assert kwargs["trace_name"] == large_trace_name
-        assert kwargs["tags"] == large_tags
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_context_manager_with_asyncio_tasks(self, mock_agentops_tracer, mock_core_tracer):
-        """Test context manager with multiple asyncio tasks"""
-        # Mock the tracer
-        mock_agentops_tracer.initialized = True
-
-        # Create traces for each task
-        trace_count = 0
-
-        def create_trace(trace_name=None, tags=None, **kwargs):
-            nonlocal trace_count
-            trace_count += 1
-            return TraceContext(span=Mock(name=f"span_{trace_count}"), token=Mock())
-
-        mock_agentops_tracer.start_trace.side_effect = create_trace
-
-        async def task_with_trace(task_id):
-            with start_trace(f"async_task_{task_id}"):
-                await asyncio.sleep(0.01)
-                return task_id
-
-        async def run_concurrent_tasks():
-            tasks = [task_with_trace(i) for i in range(5)]
-            results = await asyncio.gather(*tasks)
-            return results
-
-        # Run async tasks
-        results = asyncio.run(run_concurrent_tasks())
-        assert results == [0, 1, 2, 3, 4]
-
-        # All traces should be started and ended
-        assert mock_agentops_tracer.start_trace.call_count == 5
-        assert mock_core_tracer.end_trace.call_count == 5
-
-    @patch("agentops.sdk.core.tracer")
-    @patch("agentops.tracer")
-    def test_context_manager_resource_cleanup_on_exit_failure(self, mock_agentops_tracer, mock_core_tracer):
-        """Test that resources are cleaned up even if __exit__ fails"""
-        # Create a mock TraceContext
-        mock_trace = TraceContext(span=Mock(), token=Mock())
-        mock_agentops_tracer.initialized = True
-        mock_agentops_tracer.start_trace.return_value = mock_trace
-
-        # Make end_trace fail
-        mock_core_tracer.end_trace.side_effect = Exception("Cleanup failed")
-
-        # Should not raise exception from __exit__
-        with start_trace("cleanup_test") as trace:
-            assert trace is mock_trace
-
-        # end_trace was attempted despite failure
-        mock_core_tracer.end_trace.assert_called_once()
-
-
-if __name__ == "__main__":
-    unittest.main()