Skip to content

Commit fd121a7

Browse files
tcdenttcdent
committed
Test coverage, type checking.
1 parent eec9738 commit fd121a7

3 files changed

Lines changed: 227 additions & 4 deletions

File tree

‎agentops/helpers/dashboard.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,16 @@
22
Helpers for interacting with the AgentOps dashboard.
33
"""
44
from typing import Union
5+
from uuid import UUID
56
from termcolor import colored
6-
from opentelemetry.sdk.trace import Span
7+
from opentelemetry.sdk.trace import Span, ReadableSpan
78
from agentops.logging import logger
89
from agentops.sdk.converters import trace_id_to_uuid
910

1011

1112
APP_URL = "https://app.agentops.ai"
1213

13-
def get_trace_url(span: Span) -> str:
14+
def get_trace_url(span: Union[Span, ReadableSpan]) -> str:
1415
"""
1516
Generate a trace URL for a direct link to the session on the AgentOps dashboard.
1617
@@ -20,7 +21,7 @@ def get_trace_url(span: Span) -> str:
2021
Returns:
2122
The session URL.
2223
"""
23-
trace_id: Union[str, int] = span.context.trace_id
24+
trace_id: Union[int, UUID] = span.context.trace_id
2425

2526
# Convert trace_id to hex string if it's not already
2627
if isinstance(trace_id, int):
@@ -29,7 +30,7 @@ def get_trace_url(span: Span) -> str:
2930
return f"{APP_URL}/sessions?trace_id={trace_id}"
3031

3132

32-
def log_trace_url(span: Span) -> None:
33+
def log_trace_url(span: Union[Span, ReadableSpan]) -> None:
3334
"""
3435
Log the trace URL for the AgentOps dashboard.
3536

‎tests/unit/helpers/test_dashboard.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
"""
2+
Unit tests for dashboard URL generation and logging.
3+
"""
4+
5+
import unittest
6+
from unittest.mock import patch, MagicMock
7+
8+
from agentops.helpers.dashboard import get_trace_url, log_trace_url
9+
10+
11+
class TestDashboardHelpers(unittest.TestCase):
12+
"""Tests for dashboard URL generation and logging functions."""
13+
14+
def test_get_trace_url_with_hex_trace_id(self):
15+
"""Test get_trace_url with a hexadecimal trace ID."""
16+
# Create a mock span with a hex string trace ID
17+
mock_span = MagicMock()
18+
mock_span.context.trace_id = "1234567890abcdef"
19+
20+
# Call get_trace_url
21+
url = get_trace_url(mock_span)
22+
23+
# Assert that the URL is correctly formed
24+
self.assertEqual(url, "https://app.agentops.ai/sessions?trace_id=1234567890abcdef")
25+
26+
def test_get_trace_url_with_int_trace_id(self):
27+
"""Test get_trace_url with an integer trace ID."""
28+
# Create a mock span with an int trace ID
29+
mock_span = MagicMock()
30+
mock_span.context.trace_id = 12345
31+
32+
# Call get_trace_url
33+
url = get_trace_url(mock_span)
34+
35+
# Assert that the URL follows the expected format with a UUID
36+
self.assertTrue(url.startswith("https://app.agentops.ai/sessions?trace_id="))
37+
# Verify the format matches UUID pattern (8-4-4-4-12 hex digits)
38+
uuid_part = url.split("trace_id=")[1]
39+
self.assertRegex(uuid_part, r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$")
40+
41+
@patch('agentops.helpers.dashboard.logger')
42+
def test_log_trace_url(self, mock_logger):
43+
"""Test log_trace_url includes the session URL in the log message."""
44+
# Create a mock span
45+
mock_span = MagicMock()
46+
mock_span.context.trace_id = "test-trace-id"
47+
48+
# Mock get_trace_url to return a known value
49+
expected_url = "https://app.agentops.ai/sessions?trace_id=test-trace-id"
50+
with patch('agentops.helpers.dashboard.get_trace_url', return_value=expected_url):
51+
# Call log_trace_url
52+
log_trace_url(mock_span)
53+
54+
# Assert that logger.info was called with a message containing the URL
55+
mock_logger.info.assert_called_once()
56+
log_message = mock_logger.info.call_args[0][0]
57+
self.assertIn(expected_url, log_message)

‎tests/unit/sdk/test_internal_span_processor.py‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
"""
2+
Unit tests for the InternalSpanProcessor.
3+
"""
4+
5+
import unittest
6+
from unittest.mock import patch, MagicMock, call
7+
8+
from opentelemetry.sdk.trace import Span, ReadableSpan
9+
10+
from agentops.sdk.processors import InternalSpanProcessor
11+
12+
13+
class TestInternalSpanProcessor(unittest.TestCase):
14+
"""Tests for InternalSpanProcessor."""
15+
16+
def setUp(self):
17+
self.processor = InternalSpanProcessor()
18+
19+
# Reset the root span ID before each test
20+
self.processor._root_span_id = None
21+
22+
@patch('agentops.sdk.processors.log_trace_url')
23+
def test_logs_url_for_first_span(self, mock_log_trace_url):
24+
"""Test that the first span triggers a log_trace_url call."""
25+
# Create a mock span
26+
mock_span = MagicMock(spec=Span)
27+
mock_context = MagicMock()
28+
mock_context.trace_flags.sampled = True
29+
mock_context.span_id = 12345
30+
mock_span.context = mock_context
31+
32+
# Call on_start
33+
self.processor.on_start(mock_span)
34+
35+
# Assert that log_trace_url was called once
36+
mock_log_trace_url.assert_called_once_with(mock_span)
37+
38+
@patch('agentops.sdk.processors.log_trace_url')
39+
def test_logs_url_only_for_root_span(self, mock_log_trace_url):
40+
"""Test that log_trace_url is only called for the root span."""
41+
# First, create and start the root span
42+
mock_root_span = MagicMock(spec=Span)
43+
mock_root_context = MagicMock()
44+
mock_root_context.trace_flags.sampled = True
45+
mock_root_context.span_id = 12345
46+
mock_root_span.context = mock_root_context
47+
48+
self.processor.on_start(mock_root_span)
49+
50+
# Reset the mock after root span creation
51+
mock_log_trace_url.reset_mock()
52+
53+
# Now create and start a non-root span
54+
mock_non_root_span = MagicMock(spec=Span)
55+
mock_non_root_context = MagicMock()
56+
mock_non_root_context.trace_flags.sampled = True
57+
mock_non_root_context.span_id = 67890 # Different from root span ID
58+
mock_non_root_span.context = mock_non_root_context
59+
60+
self.processor.on_start(mock_non_root_span)
61+
62+
# Assert that log_trace_url was not called for the non-root span
63+
mock_log_trace_url.assert_not_called()
64+
65+
# End the non-root span
66+
mock_non_root_readable = MagicMock(spec=ReadableSpan)
67+
mock_non_root_readable.context = mock_non_root_context
68+
69+
self.processor.on_end(mock_non_root_readable)
70+
71+
# Assert that log_trace_url was still not called
72+
mock_log_trace_url.assert_not_called()
73+
74+
# Now end the root span
75+
mock_root_readable = MagicMock(spec=ReadableSpan)
76+
mock_root_readable.context = mock_root_context
77+
78+
self.processor.on_end(mock_root_readable)
79+
80+
# Assert that log_trace_url was called for the root span end
81+
mock_log_trace_url.assert_called_once_with(mock_root_readable)
82+
83+
@patch('agentops.sdk.processors.log_trace_url')
84+
def test_logs_url_exactly_twice_for_root_span(self, mock_log_trace_url):
85+
"""Test that log_trace_url is called exactly twice for the root span (start and end)."""
86+
# Create a mock root span
87+
mock_root_span = MagicMock(spec=Span)
88+
mock_root_context = MagicMock()
89+
mock_root_context.trace_flags.sampled = True
90+
mock_root_context.span_id = 12345
91+
mock_root_span.context = mock_root_context
92+
93+
# Start the root span
94+
self.processor.on_start(mock_root_span)
95+
96+
# Create a mock readable span for the end event
97+
mock_root_readable = MagicMock(spec=ReadableSpan)
98+
mock_root_readable.context = mock_root_context
99+
100+
# End the root span
101+
self.processor.on_end(mock_root_readable)
102+
103+
# Assert that log_trace_url was called exactly twice
104+
self.assertEqual(mock_log_trace_url.call_count, 2)
105+
mock_log_trace_url.assert_has_calls([
106+
call(mock_root_span),
107+
call(mock_root_readable)
108+
])
109+
110+
@patch('agentops.sdk.processors.log_trace_url')
111+
def test_ignores_unsampled_spans(self, mock_log_trace_url):
112+
"""Test that unsampled spans are ignored."""
113+
# Create a mock unsampled span
114+
mock_span = MagicMock(spec=Span)
115+
mock_context = MagicMock()
116+
mock_context.trace_flags.sampled = False
117+
mock_span.context = mock_context
118+
119+
# Start and end the span
120+
self.processor.on_start(mock_span)
121+
self.processor.on_end(mock_span)
122+
123+
# Assert that log_trace_url was not called
124+
mock_log_trace_url.assert_not_called()
125+
126+
# Assert that root_span_id was not set
127+
self.assertIsNone(self.processor._root_span_id)
128+
129+
@patch('agentops.sdk.processors.log_trace_url')
130+
def test_shutdown_resets_root_span_id(self, mock_log_trace_url):
131+
"""Test that shutdown resets the root span ID."""
132+
# First set a root span
133+
mock_root_span = MagicMock(spec=Span)
134+
mock_root_context = MagicMock()
135+
mock_root_context.trace_flags.sampled = True
136+
mock_root_context.span_id = 12345
137+
mock_root_span.context = mock_root_context
138+
139+
self.processor.on_start(mock_root_span)
140+
141+
# Verify root span ID was set
142+
self.assertEqual(self.processor._root_span_id, 12345)
143+
144+
# Call shutdown
145+
self.processor.shutdown()
146+
147+
# Verify root span ID was reset
148+
self.assertIsNone(self.processor._root_span_id)
149+
150+
# Create another span after shutdown
151+
mock_span = MagicMock(spec=Span)
152+
mock_context = MagicMock()
153+
mock_context.trace_flags.sampled = True
154+
mock_context.span_id = 67890
155+
mock_span.context = mock_context
156+
157+
# Reset mocks
158+
mock_log_trace_url.reset_mock()
159+
160+
# Start the span, it should be treated as a new root span
161+
self.processor.on_start(mock_span)
162+
163+
# Verify new root span was identified
164+
self.assertEqual(self.processor._root_span_id, 67890)
165+
mock_log_trace_url.assert_called_once_with(mock_span)

0 commit comments

Comments
 (0)