mseep-agentops 0.4.18__py3-none-any.whl

tests/unit/test_context_manager.py

@@ -0,0 +1,777 @@
+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()