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()