398 lines
16 KiB
Python
398 lines
16 KiB
Python
"""
|
|
Unit tests for URL logging functionality and InternalSpanProcessor.
|
|
"""
|
|
|
|
import unittest
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
from opentelemetry.sdk.trace import Span, ReadableSpan
|
|
|
|
from agentops.sdk.processors import InternalSpanProcessor
|
|
from agentops.sdk.core import TraceContext, tracer
|
|
|
|
|
|
class TestURLLogging(unittest.TestCase):
|
|
"""Tests for URL logging functionality in global tracer."""
|
|
|
|
def setUp(self):
|
|
self.tracing_core = tracer
|
|
# Mock the initialization to avoid actual setup
|
|
self.tracing_core._initialized = True
|
|
self.tracing_core._config = {"project_id": "test_project"}
|
|
|
|
@patch("agentops.sdk.core.log_trace_url")
|
|
@patch("agentops.sdk.core.tracer.make_span")
|
|
def test_start_trace_logs_url(self, mock_make_span, mock_log_trace_url):
|
|
"""Test that start_trace logs the trace URL."""
|
|
# Create a mock span
|
|
mock_span = MagicMock(spec=Span)
|
|
mock_context = MagicMock()
|
|
mock_token = MagicMock()
|
|
mock_span.get_span_context.return_value.span_id = 12345
|
|
mock_make_span.return_value = (mock_span, mock_context, mock_token)
|
|
|
|
# Call start_trace
|
|
trace_context = self.tracing_core.start_trace(trace_name="test_trace")
|
|
|
|
# Assert that log_trace_url was called with the span and title
|
|
mock_log_trace_url.assert_called_once_with(mock_span, title="test_trace")
|
|
self.assertIsInstance(trace_context, TraceContext)
|
|
self.assertEqual(trace_context.span, mock_span)
|
|
|
|
@patch("agentops.sdk.core.log_trace_url")
|
|
@patch("agentops.sdk.core.tracer.finalize_span")
|
|
def test_end_trace_logs_url(self, mock_finalize_span, mock_log_trace_url):
|
|
"""Test that end_trace logs the trace URL."""
|
|
# Create a mock trace context
|
|
mock_span = MagicMock(spec=Span)
|
|
mock_span.name = "test_trace"
|
|
mock_span.get_span_context.return_value.span_id = 12345
|
|
mock_token = MagicMock()
|
|
trace_context = TraceContext(mock_span, mock_token)
|
|
|
|
# Call end_trace
|
|
self.tracing_core.end_trace(trace_context, "Success")
|
|
|
|
# Assert that log_trace_url was called with the span and title
|
|
mock_log_trace_url.assert_called_once_with(mock_span, title="test_trace")
|
|
|
|
@patch("agentops.sdk.core.log_trace_url")
|
|
@patch("agentops.sdk.core.tracer.make_span")
|
|
def test_start_trace_url_logging_failure_does_not_break_trace(self, mock_make_span, mock_log_trace_url):
|
|
"""Test that URL logging failure doesn't break trace creation."""
|
|
# Create a mock span
|
|
mock_span = MagicMock(spec=Span)
|
|
mock_context = MagicMock()
|
|
mock_token = MagicMock()
|
|
mock_span.get_span_context.return_value.span_id = 12345
|
|
mock_make_span.return_value = (mock_span, mock_context, mock_token)
|
|
|
|
# Make log_trace_url raise an exception
|
|
mock_log_trace_url.side_effect = Exception("URL logging failed")
|
|
|
|
# Call start_trace - should not raise exception
|
|
trace_context = self.tracing_core.start_trace(trace_name="test_trace")
|
|
|
|
# Assert that trace was still created successfully
|
|
self.assertIsInstance(trace_context, TraceContext)
|
|
self.assertEqual(trace_context.span, mock_span)
|
|
mock_log_trace_url.assert_called_once_with(mock_span, title="test_trace")
|
|
|
|
@patch("agentops.sdk.core.log_trace_url")
|
|
@patch("agentops.sdk.core.tracer.finalize_span")
|
|
def test_end_trace_url_logging_failure_does_not_break_trace(self, mock_finalize_span, mock_log_trace_url):
|
|
"""Test that URL logging failure doesn't break trace ending."""
|
|
# Create a mock trace context
|
|
mock_span = MagicMock(spec=Span)
|
|
mock_span.name = "test_trace"
|
|
mock_span.get_span_context.return_value.span_id = 12345
|
|
mock_token = MagicMock()
|
|
trace_context = TraceContext(mock_span, mock_token)
|
|
|
|
# Make log_trace_url raise an exception
|
|
mock_log_trace_url.side_effect = Exception("URL logging failed")
|
|
|
|
# Call end_trace - should not raise exception
|
|
self.tracing_core.end_trace(trace_context, "Success")
|
|
|
|
# Assert that finalize_span was still called
|
|
mock_finalize_span.assert_called_once()
|
|
mock_log_trace_url.assert_called_once_with(mock_span, title="test_trace")
|
|
|
|
@patch("agentops.sdk.core.log_trace_url")
|
|
@patch("agentops.sdk.core.tracer.make_span")
|
|
def test_start_trace_with_tags_logs_url(self, mock_make_span, mock_log_trace_url):
|
|
"""Test that start_trace with tags logs the trace URL."""
|
|
# Create a mock span
|
|
mock_span = MagicMock(spec=Span)
|
|
mock_context = MagicMock()
|
|
mock_token = MagicMock()
|
|
mock_span.get_span_context.return_value.span_id = 12345
|
|
mock_make_span.return_value = (mock_span, mock_context, mock_token)
|
|
|
|
# Call start_trace with tags
|
|
trace_context = self.tracing_core.start_trace(trace_name="tagged_trace", tags=["test", "integration"])
|
|
|
|
# Assert that log_trace_url was called with the span and title
|
|
mock_log_trace_url.assert_called_once_with(mock_span, title="tagged_trace")
|
|
self.assertIsInstance(trace_context, TraceContext)
|
|
|
|
|
|
class TestSessionDecoratorURLLogging(unittest.TestCase):
|
|
"""Tests for URL logging functionality in session decorators."""
|
|
|
|
def setUp(self):
|
|
self.tracing_core = tracer
|
|
# Mock the initialization to avoid actual setup
|
|
self.tracing_core._initialized = True
|
|
self.tracing_core._config = {"project_id": "test_project"}
|
|
|
|
@patch("agentops.sdk.core.log_trace_url")
|
|
@patch("agentops.sdk.core.tracer.make_span")
|
|
@patch("agentops.sdk.core.tracer.finalize_span")
|
|
def test_session_decorator_logs_url_on_start_and_end(self, mock_finalize_span, mock_make_span, mock_log_trace_url):
|
|
"""Test that session decorator logs URLs on both start and end."""
|
|
from agentops.sdk.decorators import session
|
|
|
|
# Create a mock span
|
|
mock_span = MagicMock(spec=Span)
|
|
mock_span.name = "test_function"
|
|
mock_context = MagicMock()
|
|
mock_token = MagicMock()
|
|
mock_span.get_span_context.return_value.span_id = 12345
|
|
mock_make_span.return_value = (mock_span, mock_context, mock_token)
|
|
|
|
@session(name="test_session")
|
|
def test_function():
|
|
return "test_result"
|
|
|
|
# Call the decorated function
|
|
result = test_function()
|
|
|
|
# Assert that log_trace_url was called (start and end)
|
|
# Note: The actual number of calls may vary based on implementation details
|
|
self.assertGreaterEqual(mock_log_trace_url.call_count, 2)
|
|
# Verify that the calls include the expected session name
|
|
call_args_list = [
|
|
call_args[1]["title"] for call_args in mock_log_trace_url.call_args_list if "title" in call_args[1]
|
|
]
|
|
self.assertIn("test_session", call_args_list)
|
|
self.assertEqual(result, "test_result")
|
|
|
|
@patch("agentops.sdk.core.log_trace_url")
|
|
@patch("agentops.sdk.core.tracer.make_span")
|
|
@patch("agentops.sdk.core.tracer.finalize_span")
|
|
def test_session_decorator_with_default_name_logs_url(self, mock_finalize_span, mock_make_span, mock_log_trace_url):
|
|
"""Test that session decorator with default name logs URLs."""
|
|
from agentops.sdk.decorators import session
|
|
|
|
# Create a mock span
|
|
mock_span = MagicMock(spec=Span)
|
|
mock_span.name = "my_function"
|
|
mock_context = MagicMock()
|
|
mock_token = MagicMock()
|
|
mock_span.get_span_context.return_value.span_id = 12345
|
|
mock_make_span.return_value = (mock_span, mock_context, mock_token)
|
|
|
|
@session
|
|
def my_function():
|
|
return "result"
|
|
|
|
# Call the decorated function
|
|
result = my_function()
|
|
|
|
# Assert that log_trace_url was called with function name as title
|
|
self.assertGreaterEqual(mock_log_trace_url.call_count, 2)
|
|
# Verify that the calls include the expected function name
|
|
call_args_list = [
|
|
call_args[1]["title"] for call_args in mock_log_trace_url.call_args_list if "title" in call_args[1]
|
|
]
|
|
self.assertIn("my_function", call_args_list)
|
|
self.assertEqual(result, "result")
|
|
|
|
@patch("agentops.sdk.core.log_trace_url")
|
|
@patch("agentops.sdk.core.tracer.make_span")
|
|
@patch("agentops.sdk.core.tracer.finalize_span")
|
|
def test_session_decorator_handles_url_logging_failure(
|
|
self, mock_finalize_span, mock_make_span, mock_log_trace_url
|
|
):
|
|
"""Test that session decorator handles URL logging failures gracefully."""
|
|
from agentops.sdk.decorators import session
|
|
|
|
# Create a mock span
|
|
mock_span = MagicMock(spec=Span)
|
|
mock_span.name = "test_function"
|
|
mock_context = MagicMock()
|
|
mock_token = MagicMock()
|
|
mock_span.get_span_context.return_value.span_id = 12345
|
|
mock_make_span.return_value = (mock_span, mock_context, mock_token)
|
|
|
|
# Make log_trace_url raise an exception
|
|
mock_log_trace_url.side_effect = Exception("URL logging failed")
|
|
|
|
@session(name="failing_session")
|
|
def test_function():
|
|
return "test_result"
|
|
|
|
# Call the decorated function - should not raise exception
|
|
result = test_function()
|
|
|
|
# Assert that function still executed successfully
|
|
self.assertEqual(result, "test_result")
|
|
# Assert that log_trace_url was called (even though it failed)
|
|
self.assertGreaterEqual(mock_log_trace_url.call_count, 2)
|
|
|
|
|
|
class TestInternalSpanProcessor(unittest.TestCase):
|
|
"""Tests for InternalSpanProcessor functionality."""
|
|
|
|
def setUp(self):
|
|
self.processor = InternalSpanProcessor()
|
|
# Reset the root span ID before each test
|
|
self.processor._root_span_id = None
|
|
|
|
def test_tracks_root_span_on_start(self):
|
|
"""Test that the processor tracks the first span as root span."""
|
|
# Create a mock span
|
|
mock_span = MagicMock(spec=Span)
|
|
mock_context = MagicMock()
|
|
mock_context.trace_flags.sampled = True
|
|
mock_context.span_id = 12345
|
|
mock_span.context = mock_context
|
|
|
|
# Call on_start
|
|
self.processor.on_start(mock_span)
|
|
|
|
# Assert that root span ID was set
|
|
self.assertEqual(self.processor._root_span_id, 12345)
|
|
|
|
def test_ignores_unsampled_spans_on_start(self):
|
|
"""Test that unsampled spans are ignored on start."""
|
|
# Create a mock unsampled span
|
|
mock_span = MagicMock(spec=Span)
|
|
mock_context = MagicMock()
|
|
mock_context.trace_flags.sampled = False
|
|
mock_span.context = mock_context
|
|
|
|
# Call on_start
|
|
self.processor.on_start(mock_span)
|
|
|
|
# Assert that root span ID was not set
|
|
self.assertIsNone(self.processor._root_span_id)
|
|
|
|
def test_only_tracks_first_span_as_root(self):
|
|
"""Test that only the first span is tracked as root span."""
|
|
# First span
|
|
mock_span1 = MagicMock(spec=Span)
|
|
mock_context1 = MagicMock()
|
|
mock_context1.trace_flags.sampled = True
|
|
mock_context1.span_id = 12345
|
|
mock_span1.context = mock_context1
|
|
|
|
# Second span
|
|
mock_span2 = MagicMock(spec=Span)
|
|
mock_context2 = MagicMock()
|
|
mock_context2.trace_flags.sampled = True
|
|
mock_context2.span_id = 67890
|
|
mock_span2.context = mock_context2
|
|
|
|
# Start first span
|
|
self.processor.on_start(mock_span1)
|
|
self.assertEqual(self.processor._root_span_id, 12345)
|
|
|
|
# Start second span - should not change root span ID
|
|
self.processor.on_start(mock_span2)
|
|
self.assertEqual(self.processor._root_span_id, 12345)
|
|
|
|
@patch("agentops.sdk.processors.upload_logfile")
|
|
def test_uploads_logfile_on_root_span_end(self, mock_upload_logfile):
|
|
"""Test that logfile is uploaded when root span ends."""
|
|
# Set up root span
|
|
mock_span = MagicMock(spec=Span)
|
|
mock_context = MagicMock()
|
|
mock_context.trace_flags.sampled = True
|
|
mock_context.span_id = 12345
|
|
mock_context.trace_id = 98765
|
|
mock_span.context = mock_context
|
|
|
|
# Start the span to set it as root
|
|
self.processor.on_start(mock_span)
|
|
|
|
# Create readable span for end event
|
|
mock_readable_span = MagicMock(spec=ReadableSpan)
|
|
mock_readable_span.context = mock_context
|
|
|
|
# End the span
|
|
self.processor.on_end(mock_readable_span)
|
|
|
|
# Assert that upload_logfile was called with trace_id
|
|
mock_upload_logfile.assert_called_once_with(98765)
|
|
|
|
@patch("agentops.sdk.processors.upload_logfile")
|
|
def test_does_not_upload_logfile_for_non_root_span(self, mock_upload_logfile):
|
|
"""Test that logfile is not uploaded for non-root spans."""
|
|
# Set up root span
|
|
root_span = MagicMock(spec=Span)
|
|
root_context = MagicMock()
|
|
root_context.trace_flags.sampled = True
|
|
root_context.span_id = 12345
|
|
root_span.context = root_context
|
|
|
|
# Start root span
|
|
self.processor.on_start(root_span)
|
|
|
|
# Create non-root span
|
|
non_root_span = MagicMock(spec=ReadableSpan)
|
|
non_root_context = MagicMock()
|
|
non_root_context.trace_flags.sampled = True
|
|
non_root_context.span_id = 67890 # Different from root
|
|
non_root_span.context = non_root_context
|
|
|
|
# End non-root span
|
|
self.processor.on_end(non_root_span)
|
|
|
|
# Assert that upload_logfile was not called
|
|
mock_upload_logfile.assert_not_called()
|
|
|
|
@patch("agentops.sdk.processors.upload_logfile")
|
|
def test_handles_upload_logfile_error(self, mock_upload_logfile):
|
|
"""Test that processor handles upload_logfile errors gracefully."""
|
|
# Set up root span
|
|
mock_span = MagicMock(spec=Span)
|
|
mock_context = MagicMock()
|
|
mock_context.trace_flags.sampled = True
|
|
mock_context.span_id = 12345
|
|
mock_context.trace_id = 98765
|
|
mock_span.context = mock_context
|
|
|
|
# Start the span to set it as root
|
|
self.processor.on_start(mock_span)
|
|
|
|
# Make upload_logfile raise an exception
|
|
mock_upload_logfile.side_effect = Exception("Upload failed")
|
|
|
|
# Create readable span for end event
|
|
mock_readable_span = MagicMock(spec=ReadableSpan)
|
|
mock_readable_span.context = mock_context
|
|
|
|
# End the span - should not raise exception
|
|
self.processor.on_end(mock_readable_span)
|
|
|
|
# Assert that upload_logfile was called
|
|
mock_upload_logfile.assert_called_once_with(98765)
|
|
|
|
def test_ignores_unsampled_spans_on_end(self):
|
|
"""Test that unsampled spans are ignored on end."""
|
|
# Create a mock unsampled span
|
|
mock_span = MagicMock(spec=ReadableSpan)
|
|
mock_context = MagicMock()
|
|
mock_context.trace_flags.sampled = False
|
|
mock_span.context = mock_context
|
|
|
|
# Call on_end - should not raise exception
|
|
self.processor.on_end(mock_span)
|
|
|
|
def test_shutdown_resets_root_span_id(self):
|
|
"""Test that shutdown resets the root span ID."""
|
|
# Set up root span
|
|
mock_span = MagicMock(spec=Span)
|
|
mock_context = MagicMock()
|
|
mock_context.trace_flags.sampled = True
|
|
mock_context.span_id = 12345
|
|
mock_span.context = mock_context
|
|
|
|
# Start span to set root span ID
|
|
self.processor.on_start(mock_span)
|
|
self.assertEqual(self.processor._root_span_id, 12345)
|
|
|
|
# Call shutdown
|
|
self.processor.shutdown()
|
|
|
|
# Verify root span ID was reset
|
|
self.assertIsNone(self.processor._root_span_id)
|
|
|
|
def test_force_flush_returns_true(self):
|
|
"""Test that force_flush returns True."""
|
|
result = self.processor.force_flush()
|
|
self.assertTrue(result)
|