677 lines
24 KiB
Python
677 lines
24 KiB
Python
import pytest
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
# Tests for the new session management functionality
|
|
# These tests call the actual public API but mock the underlying implementation
|
|
# to avoid making real API calls or initializing the full telemetry pipeline
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def mock_tracing_core():
|
|
"""Mock the global tracer to avoid actual initialization"""
|
|
# Patch both the main location and where it's imported in client
|
|
with (
|
|
patch("agentops.tracer") as mock_tracer,
|
|
patch("agentops.client.client.tracer", mock_tracer),
|
|
patch("agentops.sdk.decorators.factory.tracer", mock_tracer),
|
|
patch("agentops.legacy.tracer", mock_tracer),
|
|
):
|
|
# Create a mock instance
|
|
mock_tracer.initialized = True
|
|
|
|
yield mock_tracer
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def mock_api_client():
|
|
"""Mock the API client to avoid actual API calls"""
|
|
with patch("agentops.client.api.ApiClient") as mock_api:
|
|
# Configure the v3.fetch_auth_token method to return a valid response
|
|
mock_v3 = MagicMock()
|
|
mock_v3.fetch_auth_token.return_value = {"token": "mock-jwt-token", "project_id": "mock-project-id"}
|
|
mock_api.return_value.v3 = mock_v3
|
|
|
|
yield mock_api
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def mock_trace_context():
|
|
"""Mock the TraceContext creation"""
|
|
mock_span = MagicMock()
|
|
mock_token = MagicMock()
|
|
mock_trace_context_instance = MagicMock()
|
|
mock_trace_context_instance.span = mock_span
|
|
mock_trace_context_instance.token = mock_token
|
|
mock_trace_context_instance.is_init_trace = False
|
|
|
|
return mock_trace_context_instance
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def reset_client():
|
|
"""Reset the AgentOps client before each test"""
|
|
import agentops
|
|
|
|
# Create a fresh client instance for each test
|
|
agentops._client = agentops.Client()
|
|
# Reset all client state
|
|
agentops._client._initialized = False
|
|
agentops._client._init_trace_context = None
|
|
agentops._client._legacy_session_for_init_trace = None
|
|
yield
|
|
# Clean up after test
|
|
try:
|
|
if hasattr(agentops._client, "_initialized"):
|
|
agentops._client._initialized = False
|
|
if hasattr(agentops._client, "_init_trace_context"):
|
|
agentops._client._init_trace_context = None
|
|
if hasattr(agentops._client, "_legacy_session_for_init_trace"):
|
|
agentops._client._legacy_session_for_init_trace = None
|
|
except:
|
|
pass
|
|
|
|
|
|
def test_explicit_init_then_explicit_start_trace(mock_tracing_core, mock_api_client, mock_trace_context, reset_client):
|
|
"""Test explicitly initializing followed by explicitly starting a trace"""
|
|
import agentops
|
|
|
|
# Explicitly initialize with auto_start_session=False
|
|
agentops.init(api_key="test-api-key", auto_start_session=False)
|
|
|
|
# Verify that no auto-trace was started
|
|
mock_tracing_core.start_trace.assert_not_called()
|
|
|
|
# Mock the start_trace method to return our mock trace context
|
|
mock_tracing_core.start_trace.return_value = mock_trace_context
|
|
|
|
# Explicitly start a trace
|
|
trace_context = agentops.start_trace(trace_name="test_trace", tags=["test"])
|
|
|
|
# Verify the trace was created
|
|
mock_tracing_core.start_trace.assert_called_once_with(trace_name="test_trace", tags=["test"])
|
|
assert trace_context == mock_trace_context
|
|
|
|
|
|
def test_auto_start_session_true(mock_tracing_core, mock_api_client, mock_trace_context, reset_client):
|
|
"""Test initializing with auto_start_session=True"""
|
|
import agentops
|
|
from agentops.legacy import Session
|
|
|
|
# Mock the start_trace method to return our mock trace context
|
|
mock_tracing_core.start_trace.return_value = mock_trace_context
|
|
|
|
# Initialize with auto_start_session=True
|
|
result = agentops.init(api_key="test-api-key", auto_start_session=True)
|
|
|
|
# Verify a trace was auto-started
|
|
mock_tracing_core.start_trace.assert_called_once()
|
|
# init() should return a Session object when auto-starting a session
|
|
assert isinstance(result, Session)
|
|
# Check that the session's trace_context has the expected properties
|
|
assert result.trace_context is not None
|
|
|
|
|
|
def test_auto_start_session_default(mock_tracing_core, mock_api_client, mock_trace_context, reset_client):
|
|
"""Test initializing with default auto_start_session behavior"""
|
|
import agentops
|
|
from agentops.legacy import Session
|
|
|
|
# Mock the start_trace method to return our mock trace context
|
|
mock_tracing_core.start_trace.return_value = mock_trace_context
|
|
|
|
# Initialize without explicitly setting auto_start_session (defaults to True)
|
|
result = agentops.init(api_key="test-api-key")
|
|
|
|
# Verify that the client was initialized
|
|
assert agentops._client.initialized
|
|
# Since auto_start_session defaults to True, init() should return a Session object
|
|
assert isinstance(result, Session)
|
|
# Check that the session's trace_context has the expected properties
|
|
assert result.trace_context is not None
|
|
|
|
|
|
def test_start_trace_without_init():
|
|
"""Test starting a trace without initialization triggers auto-init"""
|
|
import agentops
|
|
|
|
# Reset client for test
|
|
agentops._client = agentops.Client()
|
|
|
|
# Mock global tracer to be uninitialized initially, then initialized after init
|
|
with (
|
|
patch("agentops.tracer") as mock_tracer,
|
|
patch("agentops.client.client.tracer", mock_tracer),
|
|
patch("agentops.sdk.decorators.factory.tracer", mock_tracer),
|
|
patch("agentops.legacy.tracer", mock_tracer),
|
|
):
|
|
mock_tracer.initialized = False
|
|
|
|
# Mock the init function to simulate successful initialization
|
|
with patch("agentops.init") as mock_init:
|
|
|
|
def side_effect():
|
|
# After init is called, mark global tracer as initialized
|
|
mock_tracer.initialized = True
|
|
|
|
mock_init.side_effect = side_effect
|
|
mock_tracer.start_trace.return_value = None
|
|
|
|
# Try to start a trace without initialization
|
|
result = agentops.start_trace(trace_name="test_trace")
|
|
|
|
# Verify that init was called automatically
|
|
mock_init.assert_called_once()
|
|
# Should return None since start_trace returned None
|
|
assert result is None
|
|
|
|
|
|
def test_end_trace(mock_tracing_core, mock_trace_context):
|
|
"""Test ending a trace"""
|
|
import agentops
|
|
|
|
# End the trace
|
|
agentops.end_trace(mock_trace_context, end_state="Success")
|
|
|
|
# Verify end_trace was called on global tracer
|
|
mock_tracing_core.end_trace.assert_called_once_with(trace_context=mock_trace_context, end_state="Success")
|
|
|
|
|
|
def test_session_decorator_creates_trace(mock_tracing_core, mock_api_client, mock_trace_context, reset_client):
|
|
"""Test that the @session decorator creates a trace-level span"""
|
|
import agentops
|
|
from agentops.sdk.decorators import session
|
|
|
|
# Initialize AgentOps
|
|
agentops.init(api_key="test-api-key", auto_start_session=False)
|
|
|
|
# Reset the call count to start fresh
|
|
mock_tracing_core.reset_mock()
|
|
# Mock the start_trace and end_trace methods
|
|
mock_tracing_core.start_trace.return_value = mock_trace_context
|
|
|
|
@session(name="test_session", tags=["decorator_test"])
|
|
def test_function():
|
|
return "test_result"
|
|
|
|
# Execute the decorated function
|
|
result = test_function()
|
|
|
|
# Verify the function executed successfully
|
|
assert result == "test_result"
|
|
|
|
# Verify that start_trace and end_trace were called
|
|
# Note: The decorator might call start_trace multiple times due to initialization
|
|
assert mock_tracing_core.start_trace.call_count >= 1
|
|
assert mock_tracing_core.end_trace.call_count >= 1
|
|
|
|
|
|
def test_session_decorator_with_exception(mock_tracing_core, mock_api_client, mock_trace_context, reset_client):
|
|
"""Test that the @session decorator handles exceptions properly"""
|
|
import agentops
|
|
from agentops.sdk.decorators import session
|
|
|
|
# Initialize AgentOps
|
|
agentops.init(api_key="test-api-key", auto_start_session=False)
|
|
|
|
# Reset the call count to start fresh
|
|
mock_tracing_core.reset_mock()
|
|
# Mock the start_trace method
|
|
mock_tracing_core.start_trace.return_value = mock_trace_context
|
|
|
|
@session(name="failing_session")
|
|
def failing_function():
|
|
raise ValueError("Test exception")
|
|
|
|
# Execute the decorated function and expect an exception
|
|
with pytest.raises(ValueError, match="Test exception"):
|
|
failing_function()
|
|
|
|
# Verify that start_trace was called
|
|
assert mock_tracing_core.start_trace.call_count >= 1
|
|
# Verify that end_trace was called
|
|
assert mock_tracing_core.end_trace.call_count >= 1
|
|
|
|
|
|
def test_legacy_start_session_compatibility(mock_tracing_core, mock_api_client, mock_trace_context, reset_client):
|
|
"""Test that legacy start_session still works and calls tracer.start_trace"""
|
|
import agentops
|
|
from agentops.legacy import Session
|
|
|
|
# Initialize AgentOps
|
|
agentops.init(api_key="test-api-key", auto_start_session=False)
|
|
|
|
# Reset the call count to start fresh
|
|
mock_tracing_core.reset_mock()
|
|
# Mock the start_trace method
|
|
mock_tracing_core.start_trace.return_value = mock_trace_context
|
|
|
|
# Start a legacy session
|
|
session = agentops.start_session(tags=["legacy_test"])
|
|
|
|
# Verify the session was created
|
|
assert isinstance(session, Session)
|
|
# Check that the session's trace_context has the expected properties
|
|
assert session.trace_context is not None
|
|
|
|
# Verify that tracer.start_trace was called
|
|
# Note: May be called multiple times due to initialization
|
|
assert mock_tracing_core.start_trace.call_count >= 1
|
|
|
|
|
|
def test_legacy_end_session_compatibility(mock_tracing_core, mock_api_client, mock_trace_context, reset_client):
|
|
"""Test that legacy end_session still works and calls tracer.end_trace"""
|
|
import agentops
|
|
from agentops.legacy import Session
|
|
|
|
# Initialize AgentOps
|
|
agentops.init(api_key="test-api-key", auto_start_session=False)
|
|
|
|
# Reset the call count to start fresh
|
|
mock_tracing_core.reset_mock()
|
|
|
|
# Create a legacy session object
|
|
session = Session(mock_trace_context)
|
|
|
|
# End the session
|
|
agentops.end_session(session)
|
|
|
|
# Verify that tracer.end_trace was called
|
|
mock_tracing_core.end_trace.assert_called_once_with(mock_trace_context, end_state="Success")
|
|
|
|
|
|
def test_no_double_init(mock_tracing_core, mock_api_client, reset_client):
|
|
"""Test that calling init multiple times doesn't reinitialize"""
|
|
import agentops
|
|
|
|
# Initialize once
|
|
agentops.init(api_key="test-api-key", auto_start_session=False)
|
|
|
|
# Track the call count
|
|
call_count = mock_api_client.call_count
|
|
|
|
# Call init again with the same API key
|
|
agentops.init(api_key="test-api-key", auto_start_session=False)
|
|
|
|
# Verify that API client wasn't constructed again
|
|
assert mock_api_client.call_count == call_count
|
|
|
|
|
|
def test_client_initialization_behavior(mock_tracing_core, mock_api_client, mock_trace_context, reset_client):
|
|
"""Test basic client initialization behavior"""
|
|
import agentops
|
|
|
|
# Mock the start_trace method
|
|
mock_tracing_core.start_trace.return_value = mock_trace_context
|
|
|
|
# Test that initialization works
|
|
agentops.init(api_key="test-api-key", auto_start_session=False)
|
|
|
|
# Verify that the client was initialized
|
|
assert agentops._client.initialized
|
|
|
|
# The API client might not be called if already mocked at a higher level
|
|
# Just verify that the initialization completed successfully
|
|
|
|
# Test that calling init again doesn't cause issues
|
|
agentops.init(api_key="test-api-key", auto_start_session=False)
|
|
|
|
# Should still be initialized
|
|
assert agentops._client.initialized
|
|
|
|
|
|
def test_multiple_concurrent_traces(mock_tracing_core, mock_api_client, reset_client):
|
|
"""Test that multiple traces can be started concurrently"""
|
|
import agentops
|
|
|
|
# Initialize AgentOps
|
|
agentops.init(api_key="test-api-key", auto_start_session=False)
|
|
|
|
# Create mock trace contexts for different traces
|
|
mock_trace_context_1 = MagicMock()
|
|
mock_trace_context_2 = MagicMock()
|
|
|
|
# Mock start_trace to return different contexts
|
|
mock_tracing_core.start_trace.side_effect = [
|
|
mock_trace_context_1,
|
|
mock_trace_context_2,
|
|
]
|
|
|
|
# Start multiple traces
|
|
trace1 = agentops.start_trace(trace_name="trace1", tags=["test1"])
|
|
trace2 = agentops.start_trace(trace_name="trace2", tags=["test2"])
|
|
|
|
# Verify both traces were created
|
|
assert trace1 == mock_trace_context_1
|
|
assert trace2 == mock_trace_context_2
|
|
|
|
# Verify start_trace was called twice
|
|
assert mock_tracing_core.start_trace.call_count == 2
|
|
|
|
|
|
def test_trace_context_properties(mock_trace_context):
|
|
"""Test that TraceContext properties work correctly"""
|
|
from agentops.legacy import Session
|
|
|
|
# Create a legacy session with the mock trace context
|
|
session = Session(mock_trace_context)
|
|
|
|
# Test that properties are accessible
|
|
assert session.span == mock_trace_context.span
|
|
assert session.token == mock_trace_context.token
|
|
assert session.trace_context == mock_trace_context
|
|
|
|
|
|
def test_session_decorator_async_function(mock_tracing_core, mock_api_client, mock_trace_context, reset_client):
|
|
"""Test that the @session decorator works with async functions"""
|
|
import agentops
|
|
import asyncio
|
|
from agentops.sdk.decorators import session
|
|
|
|
# Initialize AgentOps
|
|
agentops.init(api_key="test-api-key", auto_start_session=False)
|
|
|
|
# Reset the call count to start fresh
|
|
mock_tracing_core.reset_mock()
|
|
# Mock the start_trace method
|
|
mock_tracing_core.start_trace.return_value = mock_trace_context
|
|
|
|
@session(name="async_test_session")
|
|
async def async_test_function():
|
|
await asyncio.sleep(0.01) # Simulate async work
|
|
return "async_result"
|
|
|
|
# Execute the decorated async function
|
|
result = asyncio.run(async_test_function())
|
|
|
|
# Verify the function executed successfully
|
|
assert result == "async_result"
|
|
|
|
# Verify that start_trace and end_trace were called
|
|
assert mock_tracing_core.start_trace.call_count >= 1
|
|
assert mock_tracing_core.end_trace.call_count >= 1
|
|
|
|
|
|
def test_trace_context_creation():
|
|
"""Test that TraceContext can be created with proper attributes"""
|
|
from agentops.sdk.core import TraceContext
|
|
|
|
mock_span = MagicMock()
|
|
mock_token = MagicMock()
|
|
|
|
# Test creating a TraceContext
|
|
trace_context = TraceContext(span=mock_span, token=mock_token, is_init_trace=False)
|
|
|
|
assert trace_context.span == mock_span
|
|
assert trace_context.token == mock_token
|
|
assert trace_context.is_init_trace == False
|
|
|
|
|
|
def test_session_management_integration():
|
|
"""Test the integration between new and legacy session management"""
|
|
import agentops
|
|
|
|
# Reset client for test
|
|
agentops._client = agentops.Client()
|
|
|
|
# Test that we can use both new and legacy APIs together
|
|
with (
|
|
patch("agentops.tracer") as mock_tracer,
|
|
patch("agentops.client.client.tracer", mock_tracer),
|
|
patch("agentops.sdk.decorators.factory.tracer", mock_tracer),
|
|
patch("agentops.legacy.tracer", mock_tracer),
|
|
):
|
|
mock_tracer.initialized = True
|
|
|
|
# Mock API client
|
|
with patch("agentops.client.api.ApiClient") as mock_api:
|
|
mock_v3 = MagicMock()
|
|
mock_v3.fetch_auth_token.return_value = {"token": "mock-jwt-token", "project_id": "mock-project-id"}
|
|
mock_api.return_value.v3 = mock_v3
|
|
|
|
# Initialize AgentOps
|
|
agentops.init(api_key="test-api-key", auto_start_session=False)
|
|
|
|
# Reset call counts after initialization
|
|
mock_tracer.reset_mock()
|
|
|
|
# Create mock trace context
|
|
mock_trace_context = MagicMock()
|
|
mock_tracer.start_trace.return_value = mock_trace_context
|
|
|
|
# Test new API
|
|
trace_context = agentops.start_trace(trace_name="new_api_trace")
|
|
assert trace_context == mock_trace_context
|
|
|
|
# Test legacy API
|
|
session = agentops.start_session(tags=["legacy"])
|
|
# Check that the session's trace_context has the expected properties
|
|
assert session.trace_context is not None
|
|
|
|
# Test ending both
|
|
agentops.end_trace(trace_context)
|
|
agentops.end_session(session)
|
|
|
|
# Verify calls were made
|
|
assert mock_tracer.start_trace.call_count >= 2
|
|
assert mock_tracer.end_trace.call_count >= 2
|
|
|
|
|
|
# CrewAI Backwards Compatibility Tests
|
|
# These tests ensure CrewAI integration patterns continue to work
|
|
|
|
|
|
def test_session_auto_start(mock_tracing_core, mock_api_client, mock_trace_context, reset_client):
|
|
"""Test auto-start session functionality for CrewAI compatibility"""
|
|
import agentops
|
|
from agentops.legacy import Session
|
|
|
|
# Configure mocks for session initialization
|
|
mock_tracing_core.start_trace.return_value = mock_trace_context
|
|
|
|
# Pass a dummy API key for the test
|
|
session = agentops.init(api_key="test-api-key", auto_start_session=True)
|
|
|
|
assert isinstance(session, Session)
|
|
|
|
|
|
def test_crewai_backwards_compatibility(mock_tracing_core, mock_api_client, mock_trace_context, reset_client):
|
|
"""
|
|
CrewAI needs to access:
|
|
|
|
agentops.track_agent
|
|
agentops.track_tool
|
|
agentops.start_session
|
|
agentops.end_session
|
|
agentops.ActionEvent
|
|
agentops.ToolEvent
|
|
"""
|
|
import agentops
|
|
from agentops.legacy import Session
|
|
|
|
# Configure mocks
|
|
mock_tracing_core.start_trace.return_value = mock_trace_context
|
|
|
|
# Test initialization with API key
|
|
agentops.init(api_key="test-api-key")
|
|
|
|
# Test session management functions
|
|
session = agentops.start_session(tags=["test", "crewai"])
|
|
assert isinstance(session, Session)
|
|
|
|
# Test that passing a string to end_session doesn't raise an error
|
|
agentops.end_session("Success") # This pattern is used in CrewAI
|
|
|
|
# Test track_agent function exists and doesn't raise errors
|
|
try:
|
|
# Mock an agent object similar to what CrewAI would provide
|
|
class MockAgent:
|
|
def __init__(self):
|
|
self.role = "Test Agent"
|
|
self.goal = "Testing"
|
|
self.id = "test-agent-id"
|
|
|
|
agent = MockAgent()
|
|
agentops.track_agent(agent)
|
|
except Exception as e:
|
|
assert False, f"track_agent raised an exception: {e}"
|
|
|
|
# Test track_tool function exists and doesn't raise errors
|
|
try:
|
|
# Mock a tool object similar to what CrewAI would provide
|
|
class MockTool:
|
|
def __init__(self):
|
|
self.name = "Test Tool"
|
|
self.description = "A test tool"
|
|
|
|
tool = MockTool()
|
|
agentops.track_tool(tool, "Test Agent")
|
|
except Exception as e:
|
|
assert False, f"track_tool raised an exception: {e}"
|
|
|
|
# Test events that CrewAI might use
|
|
tool_event = agentops.ToolEvent(name="test_tool")
|
|
action_event = agentops.ActionEvent(action_type="test_action")
|
|
|
|
# Verify that record function works with these events
|
|
agentops.record(tool_event)
|
|
agentops.record(action_event)
|
|
|
|
|
|
def test_crewai_kwargs_pattern(mock_tracing_core, mock_api_client, mock_trace_context, reset_client):
|
|
"""
|
|
Test the CrewAI < 0.105.0 pattern where end_session is called with only kwargs.
|
|
|
|
In versions < 0.105.0, CrewAI directly calls:
|
|
agentops.end_session(
|
|
end_state="Success",
|
|
end_state_reason="Finished Execution",
|
|
is_auto_end=True
|
|
)
|
|
"""
|
|
import agentops
|
|
from agentops.legacy import Session
|
|
|
|
# Configure mocks
|
|
mock_tracing_core.start_trace.return_value = mock_trace_context
|
|
|
|
# Initialize with test API key
|
|
agentops.init(api_key="test-api-key")
|
|
|
|
# Create a session
|
|
session = agentops.start_session(tags=["test", "crewai-kwargs"])
|
|
assert isinstance(session, Session)
|
|
|
|
# Test the CrewAI < 0.105.0 pattern - calling end_session with only kwargs
|
|
agentops.end_session(end_state="Success", end_state_reason="Finished Execution", is_auto_end=True)
|
|
|
|
# After calling end_session, creating a new session should work correctly
|
|
# (this implicitly tests that the internal state is reset properly)
|
|
new_session = agentops.start_session(tags=["test", "post-end"])
|
|
assert isinstance(new_session, Session)
|
|
|
|
|
|
def test_crewai_kwargs_pattern_no_session(mock_tracing_core, mock_api_client, reset_client):
|
|
"""
|
|
Test the CrewAI < 0.105.0 pattern where end_session is called with only kwargs,
|
|
but no session has been created.
|
|
|
|
This should log a warning but not fail.
|
|
"""
|
|
import agentops
|
|
|
|
# Initialize with test API key
|
|
agentops.init(api_key="test-api-key")
|
|
|
|
# We don't need to explicitly clear the session state
|
|
# Just make sure we start with a clean state by calling init
|
|
|
|
# Test the CrewAI < 0.105.0 pattern - calling end_session with only kwargs
|
|
# when no session exists. This should not raise an error.
|
|
agentops.end_session(end_state="Success", end_state_reason="Finished Execution", is_auto_end=True)
|
|
|
|
|
|
def test_crewai_kwargs_force_flush(mock_tracing_core, mock_api_client, mock_trace_context, reset_client):
|
|
"""
|
|
Test that when using the CrewAI < 0.105.0 pattern (end_session with kwargs),
|
|
the spans are properly exported to the backend with force_flush.
|
|
|
|
This is a more comprehensive test that ensures spans are actually sent
|
|
to the backend when using the CrewAI integration pattern.
|
|
"""
|
|
import agentops
|
|
import time
|
|
|
|
# Configure mocks
|
|
mock_tracing_core.start_trace.return_value = mock_trace_context
|
|
|
|
# Initialize AgentOps with API key
|
|
agentops.init(api_key="test-api-key")
|
|
|
|
# Create a session
|
|
agentops.start_session(tags=["test", "crewai-integration"])
|
|
|
|
# Simulate some work
|
|
time.sleep(0.1)
|
|
|
|
# End session with kwargs (CrewAI < 0.105.0 pattern)
|
|
agentops.end_session(end_state="Success", end_state_reason="Test Finished", is_auto_end=True)
|
|
|
|
# Explicitly ensure the core isn't already shut down for the test
|
|
# Note: We're using mocks, so we check the mock's initialized state
|
|
assert mock_tracing_core.initialized, "Mocked tracer should still be initialized"
|
|
|
|
|
|
def test_crewai_task_instrumentation(mock_tracing_core, mock_api_client, mock_trace_context, reset_client):
|
|
"""
|
|
Test the CrewAI task instrumentation focusing on span attributes and tags.
|
|
This test verifies that task spans are properly created with correct attributes
|
|
and tags without requiring a session.
|
|
"""
|
|
import agentops
|
|
from opentelemetry.trace import SpanKind
|
|
from agentops.semconv import SpanAttributes, AgentOpsSpanKindValues
|
|
from opentelemetry import trace
|
|
from agentops.semconv.core import CoreAttributes
|
|
|
|
# Configure mocks
|
|
mock_tracing_core.start_trace.return_value = mock_trace_context
|
|
|
|
# Initialize AgentOps with API key and default tags
|
|
agentops.init(
|
|
api_key="test-api-key",
|
|
)
|
|
agentops.start_session(tags=["test", "crewai-integration"])
|
|
# Get the tracer
|
|
tracer = trace.get_tracer(__name__)
|
|
|
|
# Create a mock task instance
|
|
class MockTask:
|
|
def __init__(self):
|
|
self.description = "Test Task Description"
|
|
self.agent = "Test Agent"
|
|
self.tools = ["tool1", "tool2"]
|
|
|
|
task = MockTask()
|
|
|
|
# Start a span for the task
|
|
with tracer.start_as_current_span(
|
|
f"{task.description}.task",
|
|
kind=SpanKind.CLIENT,
|
|
attributes={
|
|
SpanAttributes.AGENTOPS_SPAN_KIND: AgentOpsSpanKindValues.TASK.value,
|
|
CoreAttributes.TAGS: ["crewai", "task-test"],
|
|
},
|
|
) as span:
|
|
# Verify span attributes
|
|
assert span.attributes[SpanAttributes.AGENTOPS_SPAN_KIND] == AgentOpsSpanKindValues.TASK.value
|
|
assert "crewai" in span.attributes[CoreAttributes.TAGS]
|
|
assert "task-test" in span.attributes[CoreAttributes.TAGS]
|
|
|
|
# Verify span name
|
|
assert span.name == f"{task.description}.task"
|
|
|
|
# Verify span kind
|
|
assert span.kind == SpanKind.CLIENT
|
|
|
|
agentops.end_session(end_state="Success", end_state_reason="Test Finished", is_auto_end=True)
|