From a17b93a7f8f2023295fcab3ae3001ea4eb21179d Mon Sep 17 00:00:00 2001 From: Greyson LaLonde Date: Mon, 18 Aug 2025 11:55:30 -0400 Subject: [PATCH] Mock telemetry in pytest tests (#3340) * Add telemetry mocking for pytest tests - Mock telemetry by default for all tests except telemetry-specific tests - Add @pytest.mark.telemetry marker for real telemetry tests - Reduce test overhead and improve isolation * Fix telemetry test isolation - Properly isolate telemetry tests from mocking environment - Preserve API keys and other necessary environment variables - Ensure telemetry tests can run with real telemetry instances --- pyproject.toml | 5 + tests/conftest.py | 117 ++++++++++++++++++++++ tests/telemetry/test_telemetry.py | 10 +- tests/telemetry/test_telemetry_disable.py | 52 ++++++---- 4 files changed, 161 insertions(+), 23 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index e19ec38e6..2c7c9770a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -98,6 +98,11 @@ exclude = ["cli/templates"] [tool.bandit] exclude_dirs = ["src/crewai/cli/templates"] +[tool.pytest.ini_options] +markers = [ + "telemetry: mark test as a telemetry test (don't mock telemetry)", +] + # PyTorch index configuration, since torch 2.5.0 is not compatible with python 3.13 [[tool.uv.index]] name = "pytorch-nightly" diff --git a/tests/conftest.py b/tests/conftest.py index f67ad7222..6a03a163d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,7 @@ import os import tempfile from pathlib import Path +from unittest.mock import Mock, patch import pytest from dotenv import load_dotenv @@ -42,6 +43,122 @@ def setup_test_environment(): # Cleanup is handled automatically when tempfile context exits +def pytest_configure(config): + config.addinivalue_line( + "markers", "telemetry: mark test as a telemetry test (don't mock telemetry)" + ) + + +@pytest.fixture(autouse=True) +def auto_mock_telemetry(request): + if request.node.get_closest_marker("telemetry"): + telemetry_env = { + key: value + for key, value in os.environ.items() + if key not in ["CREWAI_DISABLE_TELEMETRY", "OTEL_SDK_DISABLED"] + } + with patch.dict(os.environ, telemetry_env, clear=True): + yield + return + + if "telemetry" in str(request.fspath): + telemetry_env = { + key: value + for key, value in os.environ.items() + if key not in ["CREWAI_DISABLE_TELEMETRY", "OTEL_SDK_DISABLED"] + } + with patch.dict(os.environ, telemetry_env, clear=True): + yield + return + + with patch.dict( + os.environ, {"CREWAI_DISABLE_TELEMETRY": "true", "OTEL_SDK_DISABLED": "true"} + ): + with patch("crewai.telemetry.Telemetry") as mock_telemetry_class: + mock_instance = create_mock_telemetry_instance() + mock_telemetry_class.return_value = mock_instance + + with ( + patch( + "crewai.utilities.events.event_listener.Telemetry", + mock_telemetry_class, + ), + patch("crewai.tools.tool_usage.Telemetry", mock_telemetry_class), + patch("crewai.cli.command.Telemetry", mock_telemetry_class), + patch("crewai.cli.create_flow.Telemetry", mock_telemetry_class), + ): + yield mock_instance + + +def create_mock_telemetry_instance(): + mock_instance = Mock() + + mock_instance.ready = False + mock_instance.trace_set = False + mock_instance._initialized = True + + mock_instance._is_telemetry_disabled.return_value = True + mock_instance._should_execute_telemetry.return_value = False + + telemetry_methods = [ + "set_tracer", + "crew_creation", + "task_started", + "task_ended", + "tool_usage", + "tool_repeated_usage", + "tool_usage_error", + "crew_execution_span", + "end_crew", + "flow_creation_span", + "flow_execution_span", + "individual_test_result_span", + "test_execution_span", + "deploy_signup_error_span", + "start_deployment_span", + "create_crew_deployment_span", + "get_crew_logs_span", + "remove_crew_span", + "flow_plotting_span", + "_add_attribute", + "_safe_telemetry_operation", + ] + + for method in telemetry_methods: + setattr(mock_instance, method, Mock(return_value=None)) + + mock_instance.task_started.return_value = None + + return mock_instance + + +@pytest.fixture +def mock_opentelemetry_components(): + with ( + patch("opentelemetry.trace.get_tracer") as mock_get_tracer, + patch("opentelemetry.trace.set_tracer_provider") as mock_set_provider, + patch("opentelemetry.baggage.set_baggage") as mock_set_baggage, + patch("opentelemetry.baggage.get_baggage") as mock_get_baggage, + patch("opentelemetry.context.attach") as mock_attach, + patch("opentelemetry.context.detach") as mock_detach, + ): + mock_tracer = Mock() + mock_span = Mock() + mock_tracer.start_span.return_value = mock_span + mock_get_tracer.return_value = mock_tracer + + yield { + "get_tracer": mock_get_tracer, + "set_tracer_provider": mock_set_provider, + "tracer": mock_tracer, + "span": mock_span, + "set_baggage": mock_set_baggage, + "get_baggage": mock_get_baggage, + "attach": mock_attach, + "detach": mock_detach, + } + + @pytest.fixture(scope="module") def vcr_config(request) -> dict: return { diff --git a/tests/telemetry/test_telemetry.py b/tests/telemetry/test_telemetry.py index 277578327..593f1711e 100644 --- a/tests/telemetry/test_telemetry.py +++ b/tests/telemetry/test_telemetry.py @@ -1,4 +1,5 @@ import os +import threading from unittest.mock import patch import pytest @@ -11,12 +12,16 @@ from opentelemetry import trace @pytest.fixture(autouse=True) def cleanup_telemetry(): - """Automatically clean up Telemetry singleton between tests.""" Telemetry._instance = None + if hasattr(Telemetry, "_lock"): + Telemetry._lock = threading.Lock() yield Telemetry._instance = None + if hasattr(Telemetry, "_lock"): + Telemetry._lock = threading.Lock() +@pytest.mark.telemetry @pytest.mark.parametrize( "env_var,value,expected_ready", [ @@ -36,6 +41,7 @@ def test_telemetry_environment_variables(env_var, value, expected_ready): assert telemetry.ready is expected_ready +@pytest.mark.telemetry def test_telemetry_enabled_by_default(): """Test that telemetry is enabled by default.""" with patch.dict(os.environ, {}, clear=True): @@ -44,6 +50,7 @@ def test_telemetry_enabled_by_default(): assert telemetry.ready is True +@pytest.mark.telemetry @patch("crewai.telemetry.telemetry.logger.error") @patch( "opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter.export", @@ -76,6 +83,7 @@ def test_telemetry_fails_due_connect_timeout(export_mock, logger_mock): logger_mock.assert_called_once_with(error) +@pytest.mark.telemetry def test_telemetry_singleton_pattern(): """Test that Telemetry uses the singleton pattern correctly.""" Telemetry._instance = None diff --git a/tests/telemetry/test_telemetry_disable.py b/tests/telemetry/test_telemetry_disable.py index 96738ad5f..2168bc8c2 100644 --- a/tests/telemetry/test_telemetry_disable.py +++ b/tests/telemetry/test_telemetry_disable.py @@ -14,14 +14,18 @@ def cleanup_telemetry(): Telemetry._instance = None -@pytest.mark.parametrize("env_var,value,expected_ready", [ - ("OTEL_SDK_DISABLED", "true", False), - ("OTEL_SDK_DISABLED", "TRUE", False), - ("CREWAI_DISABLE_TELEMETRY", "true", False), - ("CREWAI_DISABLE_TELEMETRY", "TRUE", False), - ("OTEL_SDK_DISABLED", "false", True), - ("CREWAI_DISABLE_TELEMETRY", "false", True), -]) +@pytest.mark.telemetry +@pytest.mark.parametrize( + "env_var,value,expected_ready", + [ + ("OTEL_SDK_DISABLED", "true", False), + ("OTEL_SDK_DISABLED", "TRUE", False), + ("CREWAI_DISABLE_TELEMETRY", "true", False), + ("CREWAI_DISABLE_TELEMETRY", "TRUE", False), + ("OTEL_SDK_DISABLED", "false", True), + ("CREWAI_DISABLE_TELEMETRY", "false", True), + ], +) def test_telemetry_environment_variables(env_var, value, expected_ready): """Test telemetry state with different environment variable configurations.""" with patch.dict(os.environ, {env_var: value}): @@ -30,6 +34,7 @@ def test_telemetry_environment_variables(env_var, value, expected_ready): assert telemetry.ready is expected_ready +@pytest.mark.telemetry def test_telemetry_enabled_by_default(): """Test that telemetry is enabled by default.""" with patch.dict(os.environ, {}, clear=True): @@ -38,57 +43,60 @@ def test_telemetry_enabled_by_default(): assert telemetry.ready is True +@pytest.mark.telemetry def test_telemetry_disable_after_singleton_creation(): """Test that telemetry operations are disabled when env var is set after singleton creation.""" with patch.dict(os.environ, {}, clear=True): with patch("crewai.telemetry.telemetry.TracerProvider"): telemetry = Telemetry() assert telemetry.ready is True - + mock_operation = MagicMock() telemetry._safe_telemetry_operation(mock_operation) mock_operation.assert_called_once() - + mock_operation.reset_mock() - - os.environ['CREWAI_DISABLE_TELEMETRY'] = 'true' - + + os.environ["CREWAI_DISABLE_TELEMETRY"] = "true" + telemetry._safe_telemetry_operation(mock_operation) mock_operation.assert_not_called() +@pytest.mark.telemetry def test_telemetry_disable_with_multiple_instances(): """Test that multiple telemetry instances respect dynamically changed env vars.""" with patch.dict(os.environ, {}, clear=True): with patch("crewai.telemetry.telemetry.TracerProvider"): telemetry1 = Telemetry() assert telemetry1.ready is True - - os.environ['CREWAI_DISABLE_TELEMETRY'] = 'true' - + + os.environ["CREWAI_DISABLE_TELEMETRY"] = "true" + telemetry2 = Telemetry() assert telemetry2 is telemetry1 assert telemetry2.ready is True - + mock_operation = MagicMock() telemetry2._safe_telemetry_operation(mock_operation) mock_operation.assert_not_called() +@pytest.mark.telemetry def test_telemetry_otel_sdk_disabled_after_creation(): """Test that OTEL_SDK_DISABLED also works when set after singleton creation.""" with patch.dict(os.environ, {}, clear=True): with patch("crewai.telemetry.telemetry.TracerProvider"): telemetry = Telemetry() assert telemetry.ready is True - + mock_operation = MagicMock() telemetry._safe_telemetry_operation(mock_operation) mock_operation.assert_called_once() - + mock_operation.reset_mock() - - os.environ['OTEL_SDK_DISABLED'] = 'true' - + + os.environ["OTEL_SDK_DISABLED"] = "true" + telemetry._safe_telemetry_operation(mock_operation) mock_operation.assert_not_called()