diff --git a/mellea/telemetry/metrics.py b/mellea/telemetry/metrics.py index b9d70fb8c..597fa86dd 100644 --- a/mellea/telemetry/metrics.py +++ b/mellea/telemetry/metrics.py @@ -106,69 +106,63 @@ # Provide dummy types for type hints metrics = None # type: ignore -# Configuration from environment variables -_METRICS_ENABLED = _OTEL_AVAILABLE and os.getenv( - "MELLEA_METRICS_ENABLED", "false" -).lower() in ("true", "1", "yes") -_METRICS_CONSOLE = os.getenv("MELLEA_METRICS_CONSOLE", "false").lower() in ( - "true", - "1", - "yes", -) -_METRICS_OTLP = os.getenv("MELLEA_METRICS_OTLP", "false").lower() in ( - "true", - "1", - "yes", -) -# Metrics-specific endpoint takes precedence over general OTLP endpoint -_OTLP_METRICS_ENDPOINT = os.getenv("OTEL_EXPORTER_OTLP_METRICS_ENDPOINT") or os.getenv( - "OTEL_EXPORTER_OTLP_ENDPOINT" -) -_METRICS_PROMETHEUS = os.getenv("MELLEA_METRICS_PROMETHEUS", "false").lower() in ( - "true", - "1", - "yes", -) -_SERVICE_NAME = os.getenv("OTEL_SERVICE_NAME", "mellea") - -# Parse export interval (default 60000 milliseconds = 60 seconds) -try: - _EXPORT_INTERVAL_MILLIS = int(os.getenv("OTEL_METRIC_EXPORT_INTERVAL", "60000")) - if _EXPORT_INTERVAL_MILLIS <= 0: + +def _env_true(name: str) -> bool: + """Return True if `name` is set to a truthy value (1/true/yes).""" + return os.getenv(name, "false").lower() in ("true", "1", "yes") + + +def _parse_export_interval() -> int: + """Read OTEL_METRIC_EXPORT_INTERVAL with warn-and-default fallback. + + Returns: + Export interval in milliseconds (default 60000 = 60 seconds). + """ + raw = os.getenv("OTEL_METRIC_EXPORT_INTERVAL", "60000") + try: + value = int(raw) + if value <= 0: + warnings.warn( + f"Invalid OTEL_METRIC_EXPORT_INTERVAL value: {value}. " + "Must be positive. Using default of 60000 milliseconds.", + UserWarning, + stacklevel=2, + ) + return 60000 + return value + except ValueError: warnings.warn( - f"Invalid OTEL_METRIC_EXPORT_INTERVAL value: {_EXPORT_INTERVAL_MILLIS}. " - "Must be positive. Using default of 60000 milliseconds.", + f"Invalid OTEL_METRIC_EXPORT_INTERVAL value: {raw}. " + "Must be an integer. Using default of 60000 milliseconds.", UserWarning, stacklevel=2, ) - _EXPORT_INTERVAL_MILLIS = 60000 -except ValueError: - warnings.warn( - f"Invalid OTEL_METRIC_EXPORT_INTERVAL value: {os.getenv('OTEL_METRIC_EXPORT_INTERVAL')}. " - "Must be an integer. Using default of 60000 milliseconds.", - UserWarning, - stacklevel=2, - ) - _EXPORT_INTERVAL_MILLIS = 60000 + return 60000 def _setup_meter_provider() -> Any: """Set up the MeterProvider with configured exporters. + Reads exporter, endpoint, and service-name env vars at call time so that + environment changes made after module import are respected without + requiring a module reload. + Returns: MeterProvider instance or None if OpenTelemetry is not available """ if not _OTEL_AVAILABLE: return None - resource = Resource.create({"service.name": _SERVICE_NAME}) # type: ignore + service_name = os.getenv("OTEL_SERVICE_NAME", "mellea") + export_interval_millis = _parse_export_interval() + resource = Resource.create({"service.name": service_name}) # type: ignore readers = [] # Add Prometheus metric reader if enabled. # This registers metrics with the prometheus_client default registry. # The application is responsible for exposing the registry (e.g. via # prometheus_client.start_http_server() or a framework integration). - if _METRICS_PROMETHEUS: + if _env_true("MELLEA_METRICS_PROMETHEUS"): try: from opentelemetry.exporter.prometheus import PrometheusMetricReader @@ -191,15 +185,18 @@ def _setup_meter_provider() -> Any: ) # Add OTLP exporter if explicitly enabled - if _METRICS_OTLP: - if _OTLP_METRICS_ENDPOINT: + if _env_true("MELLEA_METRICS_OTLP"): + otlp_endpoint = os.getenv("OTEL_EXPORTER_OTLP_METRICS_ENDPOINT") or os.getenv( + "OTEL_EXPORTER_OTLP_ENDPOINT" + ) + if otlp_endpoint: try: otlp_exporter = OTLPMetricExporter( # type: ignore - endpoint=_OTLP_METRICS_ENDPOINT + endpoint=otlp_endpoint ) readers.append( PeriodicExportingMetricReader( # type: ignore - otlp_exporter, export_interval_millis=_EXPORT_INTERVAL_MILLIS + otlp_exporter, export_interval_millis=export_interval_millis ) ) except Exception as e: @@ -218,12 +215,12 @@ def _setup_meter_provider() -> Any: ) # Add console exporter for debugging if enabled - if _METRICS_CONSOLE: + if _env_true("MELLEA_METRICS_CONSOLE"): try: console_exporter = ConsoleMetricExporter() # type: ignore readers.append( PeriodicExportingMetricReader( # type: ignore - console_exporter, export_interval_millis=_EXPORT_INTERVAL_MILLIS + console_exporter, export_interval_millis=export_interval_millis ) ) except Exception as e: @@ -267,14 +264,84 @@ def _setup_meter_provider() -> Any: return provider -# Initialize meter provider if metrics are enabled -_meter_provider = None -_meter = None +_meter_provider: Any = None +_meter: Any = None +_metrics_enabled: bool = False +_plugins_registered: bool = False + + +def _register_metrics_plugins() -> None: + """Register metrics plugins on the global plugin registry. + + Idempotent via `_plugins_registered`; subsequent calls are no-ops. The + plugin registry is process-global, so plugins remain registered for the + lifetime of the process even if the meter provider is rebuilt. + """ + global _plugins_registered + if _plugins_registered: + return + + from mellea.plugins.registry import _HAS_PLUGIN_FRAMEWORK, register + + if not _HAS_PLUGIN_FRAMEWORK: + warnings.warn( + "Metrics are enabled but the plugin framework is not installed. " + "Token usage and latency metrics will not be recorded automatically. " + "Install with: pip install mellea[telemetry]", + UserWarning, + stacklevel=2, + ) + return + + from mellea.telemetry.metrics_plugins import _METRICS_PLUGIN_CLASSES + + for plugin_cls in _METRICS_PLUGIN_CLASSES: + try: + register(plugin_cls()) + except ValueError as e: + warnings.warn( + f"{plugin_cls.__name__} already registered: {e}", + UserWarning, + stacklevel=2, + ) + _plugins_registered = True + + +def is_metrics_enabled() -> bool: + """Check if metrics collection is enabled. + + Returns: + True if `MELLEA_METRICS_ENABLED` is truthy AND OpenTelemetry is installed. + """ + return _metrics_enabled + + +def _setup_metrics() -> None: + """Build the MeterProvider and register metrics plugins. + + No-op if metrics are disabled. Env vars are read at call time inside + `_setup_meter_provider()` so changes made after module import are + respected without requiring a module reload. + """ + global _meter_provider, _meter, _metrics_enabled + + _metrics_enabled = False + if not (_OTEL_AVAILABLE and _env_true("MELLEA_METRICS_ENABLED")): + return -if _OTEL_AVAILABLE and _METRICS_ENABLED: _meter_provider = _setup_meter_provider() - if _meter_provider is not None: - _meter = metrics.get_meter("mellea.metrics", version("mellea")) # type: ignore + if _meter_provider is None: + return + # Bind directly off the provider we just built — OTel's + # `set_meter_provider()` is one-shot per process, so `metrics.get_meter()` + # would return a meter attached to the original (now shutdown) provider + # after a reset. + _meter = _meter_provider.get_meter("mellea.metrics", version("mellea")) + _metrics_enabled = True + _register_metrics_plugins() + + +_setup_metrics() # No-op instrument classes for when metrics are disabled @@ -396,15 +463,6 @@ def create_up_down_counter(name: str, description: str = "", unit: str = "1") -> return _meter.create_up_down_counter(name, description=description, unit=unit) -def is_metrics_enabled() -> bool: - """Check if metrics collection is enabled. - - Returns: - True if metrics are enabled, False otherwise - """ - return _METRICS_ENABLED - - # Token usage counters following Gen-AI semantic conventions # These are lazily initialized on first use and kept internal _input_token_counter: Any = None @@ -458,7 +516,7 @@ def record_token_usage_metrics( ) """ # Early return if metrics are disabled (zero overhead) - if not _METRICS_ENABLED: + if _meter is None: return # Get the token counters (lazily initialized) @@ -528,7 +586,7 @@ def record_request_duration( streaming=True, ) """ - if not _METRICS_ENABLED: + if _meter is None: return duration_hist, _ = _get_latency_histograms() @@ -558,7 +616,7 @@ def record_ttfb(ttfb_s: float, model: str, provider: str) -> None: provider="ollama", ) """ - if not _METRICS_ENABLED: + if _meter is None: return _, ttfb_hist = _get_latency_histograms() @@ -685,7 +743,7 @@ def record_error( exception_class="RateLimitError", ) """ - if not _METRICS_ENABLED: + if _meter is None: return counter = _get_error_counter() @@ -744,7 +802,7 @@ def record_cost(cost: float, model: str, provider: str) -> None: provider="openai", ) """ - if not _METRICS_ENABLED: + if _meter is None: return counter = _get_cost_counter() @@ -803,7 +861,7 @@ def record_sampling_attempt(strategy: str) -> None: Args: strategy: Sampling strategy class name (e.g. `"RejectionSamplingStrategy"`). """ - if not _METRICS_ENABLED: + if _meter is None: return _get_sampling_attempts_counter().add(1, {"strategy": strategy}) @@ -818,7 +876,7 @@ def record_sampling_outcome(strategy: str, success: bool) -> None: strategy: Sampling strategy class name (e.g. `"RejectionSamplingStrategy"`). success: `True` if at least one attempt passed all requirements. """ - if not _METRICS_ENABLED: + if _meter is None: return if success: @@ -865,7 +923,7 @@ def record_requirement_check(requirement: str) -> None: Args: requirement: Requirement class name (e.g. `"LLMaJRequirement"`). """ - if not _METRICS_ENABLED: + if _meter is None: return _get_requirement_checks_counter().add(1, {"requirement": requirement}) @@ -880,7 +938,7 @@ def record_requirement_failure(requirement: str, reason: str) -> None: requirement: Requirement class name (e.g. `"LLMaJRequirement"`). reason: Human-readable failure reason from `ValidationResult.reason`. """ - if not _METRICS_ENABLED: + if _meter is None: return _get_requirement_failures_counter().add( @@ -913,39 +971,13 @@ def record_tool_call(tool: str, status: str) -> None: tool: Name of the tool that was invoked. status: `"success"` if the tool executed without error, `"failure"` otherwise. """ - if not _METRICS_ENABLED: + if _meter is None: return counter = _get_tool_calls_counter() counter.add(1, {"tool": tool, "status": status}) -# Auto-register metrics plugins when metrics are enabled -if _OTEL_AVAILABLE and _METRICS_ENABLED: - from mellea.plugins.registry import _HAS_PLUGIN_FRAMEWORK, register - from mellea.telemetry.metrics_plugins import _METRICS_PLUGIN_CLASSES - - if not _HAS_PLUGIN_FRAMEWORK: - warnings.warn( - "Metrics are enabled but the plugin framework is not installed. " - "Token usage and latency metrics will not be recorded automatically. " - "Install with: pip install mellea[telemetry]", - UserWarning, - stacklevel=2, - ) - else: - for _plugin_cls in _METRICS_PLUGIN_CLASSES: - try: - register(_plugin_cls()) - except ValueError as e: - # Already registered (expected during module reloads in tests) - warnings.warn( - f"{_plugin_cls.__name__} already registered: {e}", - UserWarning, - stacklevel=2, - ) - - __all__ = [ "classify_error", "create_counter", diff --git a/mellea/telemetry/pricing.py b/mellea/telemetry/pricing.py index c0a24d248..800acd327 100644 --- a/mellea/telemetry/pricing.py +++ b/mellea/telemetry/pricing.py @@ -65,7 +65,7 @@ def _resolve_pricing_enabled() -> bool: return _LITELLM_AVAILABLE -_PRICING_ENABLED = _resolve_pricing_enabled() +_PRICING_ENABLED: bool = False _warned_models: set[str] = set() @@ -92,9 +92,23 @@ def _register_custom_pricing(path: str | Path) -> None: logger.warning("Failed to register custom pricing from %r: %s", str(path), exc) -_custom_path = os.getenv("MELLEA_PRICING_FILE") -if _PRICING_ENABLED and _custom_path: - _register_custom_pricing(_custom_path) +def _setup_pricing() -> None: + """Read env vars and register custom pricing if configured. + + Reads `MELLEA_PRICING_ENABLED` and `MELLEA_PRICING_FILE` at call time so + that environment changes made after module import are respected without + requiring a module reload. + """ + global _PRICING_ENABLED + _PRICING_ENABLED = _resolve_pricing_enabled() + if not _PRICING_ENABLED: + return + custom_path = os.getenv("MELLEA_PRICING_FILE") + if custom_path: + _register_custom_pricing(custom_path) + + +_setup_pricing() def compute_cost( diff --git a/mellea/telemetry/tracing.py b/mellea/telemetry/tracing.py index 015aef03e..37b92df3a 100644 --- a/mellea/telemetry/tracing.py +++ b/mellea/telemetry/tracing.py @@ -57,6 +57,7 @@ def _env_true(name: str) -> bool: _tracer_provider: Any = None _application_tracer: Any = None _backend_tracer: Any = None +_tracing_enabled: bool = False _plugins_registered: bool = False # Plugin registry is process-global; register once. @@ -150,20 +151,20 @@ def _register_tracing_plugins() -> None: def is_tracing_enabled() -> bool: - """Check if tracing is enabled via `MELLEA_TRACES_ENABLED`. + """Check if tracing is enabled. Returns: - True if `MELLEA_TRACES_ENABLED` is set to a truthy value AND - OpenTelemetry is installed. + True if `MELLEA_TRACES_ENABLED` is truthy AND OpenTelemetry is installed. """ - return _OTEL_AVAILABLE and _env_true("MELLEA_TRACES_ENABLED") + return _tracing_enabled def _setup_tracing() -> None: """Initialise the tracer provider, tracers, and register plugins.""" - global _tracer_provider, _application_tracer, _backend_tracer + global _tracer_provider, _application_tracer, _backend_tracer, _tracing_enabled - if not is_tracing_enabled(): + _tracing_enabled = False + if not (_OTEL_AVAILABLE and _env_true("MELLEA_TRACES_ENABLED")): return _tracer_provider = _setup_tracer_provider() @@ -175,11 +176,11 @@ def _setup_tracing() -> None: "mellea.application", mellea_version ) _backend_tracer = _tracer_provider.get_tracer("mellea.backend", mellea_version) + _tracing_enabled = True _register_tracing_plugins() -if is_tracing_enabled(): - _setup_tracing() +_setup_tracing() def is_content_tracing_enabled() -> bool: diff --git a/test/telemetry/__init__.py b/test/telemetry/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test/telemetry/conftest.py b/test/telemetry/conftest.py new file mode 100644 index 000000000..1e1c6678d --- /dev/null +++ b/test/telemetry/conftest.py @@ -0,0 +1,67 @@ +"""Shared test helpers for the `mellea.telemetry` submodules.""" + + +def reset_metrics_state() -> None: + """Reset metrics module state and re-run setup so env-var changes take effect.""" + from mellea.telemetry import metrics + + if metrics._meter_provider is not None: + metrics._meter_provider.shutdown() + metrics._meter_provider = None + metrics._meter = None + # Cached lazy instruments — bound to the old MeterProvider, must clear so + # the next record_* call re-creates them against the new provider. + metrics._input_token_counter = None + metrics._output_token_counter = None + metrics._duration_histogram = None + metrics._ttfb_histogram = None + metrics._error_counter = None + metrics._cost_counter = None + metrics._sampling_attempts_counter = None + metrics._sampling_successes_counter = None + metrics._sampling_failures_counter = None + metrics._requirement_checks_counter = None + metrics._requirement_failures_counter = None + metrics._tool_calls_counter = None + # _plugins_registered intentionally NOT reset — registry is process-global. + metrics._setup_metrics() + + +def reset_tracing_state() -> None: + """Reset tracing module state and re-run setup so env-var changes take effect.""" + from mellea.telemetry import tracing + + if tracing._tracer_provider is not None: + tracing._tracer_provider.shutdown() + tracing._tracer_provider = None + tracing._application_tracer = None + tracing._backend_tracer = None + tracing._in_flight_spans.clear() + tracing._setup_tracing() + + +def reset_logging_state() -> None: + """Reset logging module state and the MelleaLogger singleton. + + The logging module's lazy init re-runs on the next `get_otlp_log_handler()` + call. Tests that want to capture the OTLP warning or intercept the exporter + construction must do so inside the block where they invoke + `get_otlp_log_handler()`. + """ + import logging + + from mellea.core.utils import MelleaLogger + from mellea.telemetry import logging as mlogging + + logging.getLogger("mellea").handlers.clear() + MelleaLogger.logger = None + mlogging._logger_provider = None + mlogging._logger_provider_initialised = False + + +def reset_pricing_state() -> None: + """Reset pricing module state and re-run setup.""" + from mellea.telemetry import pricing + + pricing._warned_models.clear() + pricing._setup_pricing() diff --git a/test/telemetry/test_logging.py b/test/telemetry/test_logging.py index c96c62167..8e9b9ae58 100644 --- a/test/telemetry/test_logging.py +++ b/test/telemetry/test_logging.py @@ -1,17 +1,16 @@ """Unit tests for OpenTelemetry logging instrumentation.""" -import importlib import logging -import os -from unittest.mock import MagicMock, call, patch +from unittest.mock import call, patch import pytest +from test.telemetry.conftest import reset_logging_state + # Check if OpenTelemetry is available try: - from opentelemetry._logs import set_logger_provider from opentelemetry.exporter.otlp.proto.grpc._log_exporter import OTLPLogExporter - from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler + from opentelemetry.sdk._logs import LoggingHandler OTEL_AVAILABLE = True except ImportError: @@ -22,24 +21,6 @@ ) -def _reset_logging_modules(): - """Helper to reset logging state and reload modules.""" - import mellea.core.utils - import mellea.telemetry.logging - from mellea.core.utils import MelleaLogger - - # Clear any existing handlers from previous tests - fancy_logger = logging.getLogger("mellea") - fancy_logger.handlers.clear() - - # Reset MelleaLogger singleton - MelleaLogger.logger = None - - # Force reload of logging module and core.utils to pick up env vars - importlib.reload(mellea.telemetry.logging) - importlib.reload(mellea.core.utils) - - @pytest.fixture def clean_logging_env(monkeypatch): """Clean logging environment variables before each test.""" @@ -47,10 +28,9 @@ def clean_logging_env(monkeypatch): monkeypatch.delenv("OTEL_EXPORTER_OTLP_ENDPOINT", raising=False) monkeypatch.delenv("OTEL_EXPORTER_OTLP_LOGS_ENDPOINT", raising=False) monkeypatch.delenv("OTEL_SERVICE_NAME", raising=False) - - _reset_logging_modules() + reset_logging_state() yield - _reset_logging_modules() + reset_logging_state() @pytest.fixture @@ -58,10 +38,9 @@ def enable_otlp_logging(monkeypatch): """Enable OTLP logging with endpoint for tests.""" monkeypatch.setenv("MELLEA_LOGS_OTLP", "true") monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317") - - _reset_logging_modules() + reset_logging_state() yield - _reset_logging_modules() + reset_logging_state() # Configuration Tests @@ -88,8 +67,7 @@ def test_otlp_logging_enabled_without_endpoint_warns(monkeypatch, clean_logging_ """Test that enabling OTLP without endpoint produces warning on first handler request.""" monkeypatch.setenv("MELLEA_LOGS_OTLP", "true") # No endpoint set - - _reset_logging_modules() + reset_logging_state() from mellea.telemetry.logging import get_otlp_log_handler @@ -105,10 +83,7 @@ def test_otlp_logging_with_various_truthy_values(monkeypatch, clean_logging_env) for value in ["true", "True", "TRUE", "1", "yes", "Yes", "YES"]: monkeypatch.setenv("MELLEA_LOGS_OTLP", value) - - import mellea.telemetry.logging - - importlib.reload(mellea.telemetry.logging) + reset_logging_state() from mellea.telemetry.logging import get_otlp_log_handler @@ -121,8 +96,7 @@ def test_logs_specific_endpoint_takes_precedence(monkeypatch, clean_logging_env) monkeypatch.setenv("MELLEA_LOGS_OTLP", "true") monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317") monkeypatch.setenv("OTEL_EXPORTER_OTLP_LOGS_ENDPOINT", "http://localhost:4318/logs") - - _reset_logging_modules() + reset_logging_state() import mellea.telemetry.logging diff --git a/test/telemetry/test_metrics.py b/test/telemetry/test_metrics.py index 0bd1136fa..7ce6e5d7c 100644 --- a/test/telemetry/test_metrics.py +++ b/test/telemetry/test_metrics.py @@ -4,6 +4,8 @@ import pytest +from test.telemetry.conftest import reset_metrics_state + # Check if OpenTelemetry is available try: from opentelemetry import metrics @@ -30,22 +32,16 @@ def clean_metrics_env(monkeypatch): monkeypatch.delenv("MELLEA_METRICS_PROMETHEUS", raising=False) monkeypatch.delenv("OTEL_METRIC_EXPORT_INTERVAL", raising=False) monkeypatch.delenv("OTEL_SERVICE_NAME", raising=False) - # Force reload of metrics module to pick up env vars - import importlib - - import mellea.telemetry.metrics - - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() yield - # Reset after test - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() @pytest.fixture def enable_metrics(monkeypatch): """Enable metrics for tests.""" monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") - # Clear other env vars to prevent user-set values from leaking into reload + # Clear other env vars to prevent user-set values from leaking in monkeypatch.delenv("MELLEA_METRICS_CONSOLE", raising=False) monkeypatch.delenv("MELLEA_METRICS_OTLP", raising=False) monkeypatch.delenv("MELLEA_METRICS_PROMETHEUS", raising=False) @@ -53,16 +49,10 @@ def enable_metrics(monkeypatch): monkeypatch.delenv("OTEL_EXPORTER_OTLP_METRICS_ENDPOINT", raising=False) monkeypatch.delenv("OTEL_METRIC_EXPORT_INTERVAL", raising=False) monkeypatch.delenv("OTEL_SERVICE_NAME", raising=False) - # Force reload of metrics module to pick up env vars - import importlib - - import mellea.telemetry.metrics - - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() yield - # Reset after test monkeypatch.setenv("MELLEA_METRICS_ENABLED", "false") - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() @pytest.fixture @@ -98,13 +88,9 @@ def test_metrics_enabled_with_env_var(enable_metrics): def test_metrics_enabled_with_various_truthy_values(monkeypatch): """Test that various truthy values enable metrics.""" - import importlib - - import mellea.telemetry.metrics - for value in ["true", "True", "TRUE", "1", "yes", "Yes", "YES"]: monkeypatch.setenv("MELLEA_METRICS_ENABLED", value) - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() from mellea.telemetry.metrics import is_metrics_enabled assert is_metrics_enabled(), f"Failed for value: {value}" @@ -112,13 +98,9 @@ def test_metrics_enabled_with_various_truthy_values(monkeypatch): def test_metrics_disabled_with_falsy_values(monkeypatch): """Test that falsy values keep metrics disabled.""" - import importlib - - import mellea.telemetry.metrics - for value in ["false", "False", "FALSE", "0", "no", "No", "NO", ""]: monkeypatch.setenv("MELLEA_METRICS_ENABLED", value) - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() from mellea.telemetry.metrics import is_metrics_enabled assert not is_metrics_enabled(), f"Failed for value: {value}" @@ -153,6 +135,52 @@ def test_meter_reused_across_instruments(enable_metrics): assert _meter is not None +def test_meter_remains_functional_after_repeated_resets(monkeypatch): + """Recordings via `_meter` must keep landing in the active reader after reset cycles.""" + monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") + monkeypatch.delenv("MELLEA_METRICS_CONSOLE", raising=False) + monkeypatch.delenv("MELLEA_METRICS_OTLP", raising=False) + monkeypatch.delenv("MELLEA_METRICS_PROMETHEUS", raising=False) + + from mellea.telemetry import metrics as metrics_module + + readers: list[InMemoryMetricReader] = [] + + def fake_setup_meter_provider() -> MeterProvider: + reader = InMemoryMetricReader() + readers.append(reader) + provider = MeterProvider(metric_readers=[reader]) + # Must call set_meter_provider — its 2nd+ no-op is the bug trigger. + from opentelemetry import metrics as otel_metrics + + otel_metrics.set_meter_provider(provider) + return provider + + monkeypatch.setattr( + metrics_module, "_setup_meter_provider", fake_setup_meter_provider + ) + + for cycle in range(3): + reset_metrics_state() + metrics_module.record_token_usage_metrics( + input_tokens=10, output_tokens=5, model="m", provider="p" + ) + metrics_module._meter_provider.force_flush() + data = readers[-1].get_metrics_data() + recorded = [ + m.name + for rm in data.resource_metrics + for sm in rm.scope_metrics + for m in sm.metrics + ] + assert "mellea.llm.tokens.input" in recorded, ( + f"cycle {cycle}: token recording vanished — _meter is bound to " + f"a stale/shutdown MeterProvider" + ) + + reset_metrics_state() + + # Instrument Creation Tests @@ -253,26 +281,18 @@ def test_noop_updown_counter_methods_dont_raise(clean_metrics_env): # Import Safety Tests -def test_graceful_handling_without_opentelemetry(): +def test_graceful_handling_without_opentelemetry(monkeypatch): """Test that metrics module handles missing OpenTelemetry gracefully.""" - with patch.dict("sys.modules", {"opentelemetry": None}): - # Force reimport - import importlib + import mellea.telemetry.metrics as metrics_module - import mellea.telemetry.metrics + monkeypatch.setattr(metrics_module, "_OTEL_AVAILABLE", False) + reset_metrics_state() - importlib.reload(mellea.telemetry.metrics) + from mellea.telemetry.metrics import create_counter, is_metrics_enabled - # Should not raise, metrics should be disabled - from mellea.telemetry.metrics import ( - create_counter, - create_histogram, - is_metrics_enabled, - ) - - assert not is_metrics_enabled() - counter = create_counter("test.counter") - assert counter is not None # Should be no-op + assert not is_metrics_enabled() + counter = create_counter("test.counter") + assert counter is not None # Should be no-op # Functional Tests with Real Instruments @@ -356,23 +376,21 @@ def test_histogram_with_attributes(enable_metrics): def test_custom_service_name(monkeypatch, enable_metrics): """Test that custom service name is used.""" monkeypatch.setenv("OTEL_SERVICE_NAME", "my-custom-service") + reset_metrics_state() - import importlib - - import mellea.telemetry.metrics - - importlib.reload(mellea.telemetry.metrics) + import mellea.telemetry.metrics as m - from mellea.telemetry.metrics import _SERVICE_NAME - - assert _SERVICE_NAME == "my-custom-service" + assert ( + m._meter_provider._sdk_config.resource.attributes["service.name"] + == "my-custom-service" + ) def test_default_service_name(enable_metrics): """Test that default service name is 'mellea'.""" - from mellea.telemetry.metrics import _SERVICE_NAME + import mellea.telemetry.metrics as m - assert _SERVICE_NAME == "mellea" + assert m._meter_provider._sdk_config.resource.attributes["service.name"] == "mellea" # Console Exporter Tests @@ -380,25 +398,34 @@ def test_default_service_name(enable_metrics): def test_console_exporter_enabled(monkeypatch, shutdown_meter_provider): """Test that console exporter can be enabled.""" + from opentelemetry.sdk.metrics.export import ( + ConsoleMetricExporter, + PeriodicExportingMetricReader, + ) + monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") monkeypatch.setenv("MELLEA_METRICS_CONSOLE", "true") - import importlib - - import mellea.telemetry.metrics - - importlib.reload(mellea.telemetry.metrics) + with patch( + "mellea.telemetry.metrics.ConsoleMetricExporter", wraps=ConsoleMetricExporter + ) as mock_exporter: + reset_metrics_state() + mock_exporter.assert_called_once() - from mellea.telemetry.metrics import _METRICS_CONSOLE + import mellea.telemetry.metrics as m - assert _METRICS_CONSOLE is True + assert any( + isinstance(r, PeriodicExportingMetricReader) + for r in m._meter_provider._sdk_config.metric_readers + ) def test_console_exporter_disabled_by_default(enable_metrics): """Test that console exporter is disabled by default.""" - from mellea.telemetry.metrics import _METRICS_CONSOLE + import mellea.telemetry.metrics as m - assert _METRICS_CONSOLE is False + # No readers attached because no exporter env vars are set in `enable_metrics`. + assert not m._meter_provider._sdk_config.metric_readers # OTLP Exporter Tests @@ -406,69 +433,72 @@ def test_console_exporter_disabled_by_default(enable_metrics): def test_otlp_explicit_enablement(monkeypatch, shutdown_meter_provider): """Test that OTLP exporter requires explicit enablement via MELLEA_METRICS_OTLP.""" + from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import ( + OTLPMetricExporter, + ) + monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") monkeypatch.setenv("MELLEA_METRICS_OTLP", "true") monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317") - import importlib - - import mellea.telemetry.metrics - - importlib.reload(mellea.telemetry.metrics) - - from mellea.telemetry.metrics import _METRICS_OTLP - - assert _METRICS_OTLP is True + with patch( + "mellea.telemetry.metrics.OTLPMetricExporter", wraps=OTLPMetricExporter + ) as mock_exporter: + reset_metrics_state() + mock_exporter.assert_called_once_with(endpoint="http://localhost:4317") def test_metrics_specific_endpoint_precedence(monkeypatch): """Test that OTEL_EXPORTER_OTLP_METRICS_ENDPOINT takes precedence over general endpoint.""" + from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import ( + OTLPMetricExporter, + ) + monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") + monkeypatch.setenv("MELLEA_METRICS_OTLP", "true") monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317") monkeypatch.setenv("OTEL_EXPORTER_OTLP_METRICS_ENDPOINT", "http://localhost:4318") - import importlib - - import mellea.telemetry.metrics - - importlib.reload(mellea.telemetry.metrics) - - from mellea.telemetry.metrics import _OTLP_METRICS_ENDPOINT - - assert _OTLP_METRICS_ENDPOINT == "http://localhost:4318" + with patch( + "mellea.telemetry.metrics.OTLPMetricExporter", wraps=OTLPMetricExporter + ) as mock_exporter: + reset_metrics_state() + mock_exporter.assert_called_once_with(endpoint="http://localhost:4318") def test_custom_export_interval(monkeypatch): """Test that custom export interval is configured correctly.""" + from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader + monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") + monkeypatch.setenv("MELLEA_METRICS_CONSOLE", "true") monkeypatch.setenv("OTEL_METRIC_EXPORT_INTERVAL", "30000") - import importlib - - import mellea.telemetry.metrics - - importlib.reload(mellea.telemetry.metrics) - - from mellea.telemetry.metrics import _EXPORT_INTERVAL_MILLIS - - assert _EXPORT_INTERVAL_MILLIS == 30000 + with patch( + "mellea.telemetry.metrics.PeriodicExportingMetricReader", + wraps=PeriodicExportingMetricReader, + ) as mock_reader: + reset_metrics_state() + assert mock_reader.call_args.kwargs["export_interval_millis"] == 30000 def test_invalid_export_interval_warning(monkeypatch): - """Test that invalid export interval produces warning and uses default.""" + """Test that invalid export interval produces warning and falls back to 60000.""" + from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader + monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") + monkeypatch.setenv("MELLEA_METRICS_CONSOLE", "true") monkeypatch.setenv("OTEL_METRIC_EXPORT_INTERVAL", "invalid") - import importlib - - import mellea.telemetry.metrics - - with pytest.warns(UserWarning, match="Invalid OTEL_METRIC_EXPORT_INTERVAL"): - importlib.reload(mellea.telemetry.metrics) - - from mellea.telemetry.metrics import _EXPORT_INTERVAL_MILLIS - - assert _EXPORT_INTERVAL_MILLIS == 60000 + with ( + pytest.warns(UserWarning, match="Invalid OTEL_METRIC_EXPORT_INTERVAL"), + patch( + "mellea.telemetry.metrics.PeriodicExportingMetricReader", + wraps=PeriodicExportingMetricReader, + ) as mock_reader, + ): + reset_metrics_state() + assert mock_reader.call_args.kwargs["export_interval_millis"] == 60000 def test_otlp_enabled_without_endpoint_warning(monkeypatch): @@ -479,15 +509,11 @@ def test_otlp_enabled_without_endpoint_warning(monkeypatch): monkeypatch.delenv("OTEL_EXPORTER_OTLP_METRICS_ENDPOINT", raising=False) monkeypatch.delenv("OTEL_EXPORTER_OTLP_ENDPOINT", raising=False) - import importlib - - import mellea.telemetry.metrics - with pytest.warns( UserWarning, match="OTLP metrics exporter is enabled.*but no endpoint is configured", ): - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() # Prometheus Exporter Tests @@ -495,29 +521,26 @@ def test_otlp_enabled_without_endpoint_warning(monkeypatch): def test_prometheus_exporter_enabled(monkeypatch): """Test that Prometheus metric reader initializes when enabled.""" - import importlib - - import mellea.telemetry.metrics + from opentelemetry.exporter.prometheus import PrometheusMetricReader monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") monkeypatch.setenv("MELLEA_METRICS_PROMETHEUS", "true") + reset_metrics_state() - importlib.reload(mellea.telemetry.metrics) - - from mellea.telemetry.metrics import _METRICS_PROMETHEUS + import mellea.telemetry.metrics as m - assert _METRICS_PROMETHEUS is True + assert any( + isinstance(r, PrometheusMetricReader) + for r in m._meter_provider._sdk_config.metric_readers + ) def test_prometheus_exporter_import_error_warning(monkeypatch): """Test that missing Prometheus package produces helpful warning.""" - monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") - monkeypatch.setenv("MELLEA_METRICS_PROMETHEUS", "true") - - import importlib import sys - import mellea.telemetry.metrics + monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") + monkeypatch.setenv("MELLEA_METRICS_PROMETHEUS", "true") # Temporarily remove the module to simulate ImportError original_modules = sys.modules.copy() @@ -528,7 +551,7 @@ def test_prometheus_exporter_import_error_warning(monkeypatch): UserWarning, match="Prometheus exporter is enabled.*but opentelemetry-exporter-prometheus is not installed", ): - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() finally: # Restore modules sys.modules.clear() @@ -537,48 +560,53 @@ def test_prometheus_exporter_import_error_warning(monkeypatch): def test_prometheus_and_otlp_exporters_together(monkeypatch, shutdown_meter_provider): """Test that Prometheus and OTLP exporters can run simultaneously.""" + from opentelemetry.exporter.prometheus import PrometheusMetricReader + from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader + monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") monkeypatch.setenv("MELLEA_METRICS_PROMETHEUS", "true") monkeypatch.setenv("MELLEA_METRICS_OTLP", "true") monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317") + reset_metrics_state() - import importlib - - import mellea.telemetry.metrics - - importlib.reload(mellea.telemetry.metrics) + import mellea.telemetry.metrics as m - from mellea.telemetry.metrics import _METRICS_OTLP, _METRICS_PROMETHEUS - - assert _METRICS_PROMETHEUS is True - assert _METRICS_OTLP is True + reader_types = [type(r) for r in m._meter_provider._sdk_config.metric_readers] + assert PrometheusMetricReader in reader_types + # OTLP exporter is wrapped in a PeriodicExportingMetricReader + assert PeriodicExportingMetricReader in reader_types def test_prometheus_exporter_disabled_by_default(enable_metrics): """Test that Prometheus exporter is disabled by default.""" - from mellea.telemetry.metrics import _METRICS_PROMETHEUS + from opentelemetry.exporter.prometheus import PrometheusMetricReader + + import mellea.telemetry.metrics as m - assert _METRICS_PROMETHEUS is False + assert not any( + isinstance(r, PrometheusMetricReader) + for r in m._meter_provider._sdk_config.metric_readers + ) def test_prometheus_exporter_with_console_exporter( monkeypatch, shutdown_meter_provider ): """Test that Prometheus works alongside console exporter.""" + from opentelemetry.exporter.prometheus import PrometheusMetricReader + from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader + monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") monkeypatch.setenv("MELLEA_METRICS_PROMETHEUS", "true") monkeypatch.setenv("MELLEA_METRICS_CONSOLE", "true") + reset_metrics_state() - import importlib - - import mellea.telemetry.metrics - - importlib.reload(mellea.telemetry.metrics) - - from mellea.telemetry.metrics import _METRICS_CONSOLE, _METRICS_PROMETHEUS + import mellea.telemetry.metrics as m - assert _METRICS_PROMETHEUS is True - assert _METRICS_CONSOLE is True + reader_types = [type(r) for r in m._meter_provider._sdk_config.metric_readers] + assert PrometheusMetricReader in reader_types + # Console exporter is wrapped in a PeriodicExportingMetricReader + assert PeriodicExportingMetricReader in reader_types # Metric Instrument Tests diff --git a/test/telemetry/test_metrics_backend.py b/test/telemetry/test_metrics_backend.py index 2119501cf..b69158012 100644 --- a/test/telemetry/test_metrics_backend.py +++ b/test/telemetry/test_metrics_backend.py @@ -18,6 +18,7 @@ from mellea.stdlib.context import SimpleContext from test.conftest import hf_skip from test.predicates import require_api_key, require_gpu +from test.telemetry.conftest import reset_metrics_state # Check if OpenTelemetry is available try: @@ -47,16 +48,10 @@ def enable_metrics(monkeypatch): enable_background_collection() discard_background_tasks() monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") - # Force reload of metrics module to pick up env vars - import importlib - - import mellea.telemetry.metrics - - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() yield - # Reset after test monkeypatch.setenv("MELLEA_METRICS_ENABLED", "false") - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() disable_background_collection() diff --git a/test/telemetry/test_metrics_cost.py b/test/telemetry/test_metrics_cost.py index 8ed757b55..064cfbbdd 100644 --- a/test/telemetry/test_metrics_cost.py +++ b/test/telemetry/test_metrics_cost.py @@ -6,6 +6,8 @@ import pytest +from test.telemetry.conftest import reset_metrics_state + # Check if OpenTelemetry is available try: from opentelemetry.sdk.metrics import MeterProvider @@ -26,14 +28,9 @@ def clean_metrics_env(monkeypatch): """Enable metrics and reset module state for integration tests.""" monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") monkeypatch.delenv("MELLEA_METRICS_CONSOLE", raising=False) - - import importlib - - import mellea.telemetry.metrics - - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() yield - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() def _setup_in_memory_provider(metrics_module): diff --git a/test/telemetry/test_metrics_errors.py b/test/telemetry/test_metrics_errors.py index f1e2ebcb4..e266af830 100644 --- a/test/telemetry/test_metrics_errors.py +++ b/test/telemetry/test_metrics_errors.py @@ -6,6 +6,8 @@ import pytest +from test.telemetry.conftest import reset_metrics_state + # Check if OpenTelemetry is available try: from opentelemetry.sdk.metrics import MeterProvider @@ -26,14 +28,9 @@ def clean_metrics_env(monkeypatch): """Enable metrics and reset module state for integration tests.""" monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") monkeypatch.delenv("MELLEA_METRICS_CONSOLE", raising=False) - - import importlib - - import mellea.telemetry.metrics - - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() yield - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() def _setup_in_memory_provider(metrics_module): diff --git a/test/telemetry/test_metrics_latency.py b/test/telemetry/test_metrics_latency.py index fe5b691f6..4aca2dbad 100644 --- a/test/telemetry/test_metrics_latency.py +++ b/test/telemetry/test_metrics_latency.py @@ -6,6 +6,8 @@ import pytest +from test.telemetry.conftest import reset_metrics_state + # Check if OpenTelemetry is available try: from opentelemetry.sdk.metrics import MeterProvider @@ -26,14 +28,9 @@ def clean_metrics_env(monkeypatch): """Enable metrics and reset module state for integration tests.""" monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") monkeypatch.delenv("MELLEA_METRICS_CONSOLE", raising=False) - - import importlib - - import mellea.telemetry.metrics - - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() yield - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() def _setup_in_memory_provider(metrics_module): diff --git a/test/telemetry/test_metrics_operational.py b/test/telemetry/test_metrics_operational.py index 0d2187133..3c68a7a16 100644 --- a/test/telemetry/test_metrics_operational.py +++ b/test/telemetry/test_metrics_operational.py @@ -6,6 +6,8 @@ import pytest +from test.telemetry.conftest import reset_metrics_state + try: from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.metrics.export import InMemoryMetricReader @@ -24,14 +26,9 @@ def clean_metrics_env(monkeypatch): """Enable metrics and reset all module state for each test.""" monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") monkeypatch.delenv("MELLEA_METRICS_CONSOLE", raising=False) - - import importlib - - import mellea.telemetry.metrics - - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() yield - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() def _setup_in_memory_provider(metrics_module): @@ -312,12 +309,10 @@ def test_record_tool_call_multiple_tools(clean_metrics_env): def test_record_sampling_attempt_noop_when_disabled(monkeypatch): """record_sampling_attempt is a no-op when metrics are disabled.""" - import importlib - import mellea.telemetry.metrics as m - importlib.reload(m) - monkeypatch.setattr(m, "_METRICS_ENABLED", False) + reset_metrics_state() + monkeypatch.setattr(m, "_meter", None) # Should not raise and should not create any counter m.record_sampling_attempt("RejectionSamplingStrategy") @@ -326,12 +321,10 @@ def test_record_sampling_attempt_noop_when_disabled(monkeypatch): def test_record_sampling_outcome_noop_when_disabled(monkeypatch): """record_sampling_outcome is a no-op when metrics are disabled.""" - import importlib - import mellea.telemetry.metrics as m - importlib.reload(m) - monkeypatch.setattr(m, "_METRICS_ENABLED", False) + reset_metrics_state() + monkeypatch.setattr(m, "_meter", None) m.record_sampling_outcome("RejectionSamplingStrategy", success=True) assert m._sampling_successes_counter is None @@ -339,12 +332,10 @@ def test_record_sampling_outcome_noop_when_disabled(monkeypatch): def test_record_requirement_check_noop_when_disabled(monkeypatch): """record_requirement_check is a no-op when metrics are disabled.""" - import importlib - import mellea.telemetry.metrics as m - importlib.reload(m) - monkeypatch.setattr(m, "_METRICS_ENABLED", False) + reset_metrics_state() + monkeypatch.setattr(m, "_meter", None) m.record_requirement_check("LLMaJRequirement") assert m._requirement_checks_counter is None @@ -352,12 +343,10 @@ def test_record_requirement_check_noop_when_disabled(monkeypatch): def test_record_requirement_failure_noop_when_disabled(monkeypatch): """record_requirement_failure is a no-op when metrics are disabled.""" - import importlib - import mellea.telemetry.metrics as m - importlib.reload(m) - monkeypatch.setattr(m, "_METRICS_ENABLED", False) + reset_metrics_state() + monkeypatch.setattr(m, "_meter", None) m.record_requirement_failure("LLMaJRequirement", "reason") assert m._requirement_failures_counter is None @@ -365,12 +354,10 @@ def test_record_requirement_failure_noop_when_disabled(monkeypatch): def test_record_tool_call_noop_when_disabled(monkeypatch): """record_tool_call is a no-op when metrics are disabled.""" - import importlib - import mellea.telemetry.metrics as m - importlib.reload(m) - monkeypatch.setattr(m, "_METRICS_ENABLED", False) + reset_metrics_state() + monkeypatch.setattr(m, "_meter", None) m.record_tool_call("search", "success") assert m._tool_calls_counter is None diff --git a/test/telemetry/test_metrics_token.py b/test/telemetry/test_metrics_token.py index c688ca278..04fc95e93 100644 --- a/test/telemetry/test_metrics_token.py +++ b/test/telemetry/test_metrics_token.py @@ -6,6 +6,8 @@ import pytest +from test.telemetry.conftest import reset_metrics_state + # Check if OpenTelemetry is available try: from opentelemetry.sdk.metrics import MeterProvider @@ -24,21 +26,11 @@ @pytest.fixture def clean_metrics_env(monkeypatch): """Clean metrics environment variables and enable metrics for tests.""" - # Enable metrics for integration tests monkeypatch.setenv("MELLEA_METRICS_ENABLED", "true") - - # Clean other metrics env vars monkeypatch.delenv("MELLEA_METRICS_CONSOLE", raising=False) - - # Force reload of metrics module to pick up env vars - import importlib - - import mellea.telemetry.metrics - - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() yield - # Reset after test - importlib.reload(mellea.telemetry.metrics) + reset_metrics_state() def _setup_in_memory_provider(metrics_module): diff --git a/test/telemetry/test_pricing.py b/test/telemetry/test_pricing.py index 25b9746d2..f9d47ecab 100644 --- a/test/telemetry/test_pricing.py +++ b/test/telemetry/test_pricing.py @@ -1,27 +1,21 @@ """Unit tests for the litellm-backed pricing module.""" -import importlib import json -import sys from unittest.mock import MagicMock, patch import pytest from mellea.telemetry import pricing +from test.telemetry.conftest import reset_pricing_state @pytest.fixture() def restore_pricing(monkeypatch): - """Restore litellm import state and reload pricing module after each test.""" - original_litellm = sys.modules.get("litellm", ...) + """Restore pricing module state after each test.""" yield monkeypatch.delenv("MELLEA_PRICING_ENABLED", raising=False) monkeypatch.delenv("MELLEA_PRICING_FILE", raising=False) - if original_litellm is ...: - sys.modules.pop("litellm", None) - else: - sys.modules["litellm"] = original_litellm - importlib.reload(pricing) + reset_pricing_state() @pytest.fixture() @@ -45,16 +39,16 @@ def mock_litellm_pricing(): def test_tri_state_false_disables_explicitly(monkeypatch, restore_pricing): """MELLEA_PRICING_ENABLED=false disables pricing regardless of litellm.""" monkeypatch.setenv("MELLEA_PRICING_ENABLED", "false") - monkeypatch.setitem(sys.modules, "litellm", MagicMock()) - importlib.reload(pricing) + monkeypatch.setattr(pricing, "_LITELLM_AVAILABLE", True) + reset_pricing_state() assert pricing._PRICING_ENABLED is False def test_tri_state_true_with_litellm_enables(monkeypatch, restore_pricing): """MELLEA_PRICING_ENABLED=true with litellm present enables pricing.""" monkeypatch.setenv("MELLEA_PRICING_ENABLED", "true") - monkeypatch.setitem(sys.modules, "litellm", MagicMock()) - importlib.reload(pricing) + monkeypatch.setattr(pricing, "_LITELLM_AVAILABLE", True) + reset_pricing_state() assert pricing._PRICING_ENABLED is True @@ -63,17 +57,17 @@ def test_tri_state_true_without_litellm_warns_and_disables( ): """MELLEA_PRICING_ENABLED=true without litellm emits a warning and disables.""" monkeypatch.setenv("MELLEA_PRICING_ENABLED", "true") - monkeypatch.setitem(sys.modules, "litellm", None) + monkeypatch.setattr(pricing, "_LITELLM_AVAILABLE", False) with pytest.warns(UserWarning, match="litellm is not installed"): - importlib.reload(pricing) + reset_pricing_state() assert pricing._PRICING_ENABLED is False def test_tri_state_unset_with_litellm_auto_enables(monkeypatch, restore_pricing): """Unset MELLEA_PRICING_ENABLED with litellm present auto-enables pricing.""" monkeypatch.delenv("MELLEA_PRICING_ENABLED", raising=False) - monkeypatch.setitem(sys.modules, "litellm", MagicMock()) - importlib.reload(pricing) + monkeypatch.setattr(pricing, "_LITELLM_AVAILABLE", True) + reset_pricing_state() assert pricing._PRICING_ENABLED is True @@ -82,8 +76,8 @@ def test_tri_state_unset_without_litellm_silent_disable( ): """Unset MELLEA_PRICING_ENABLED without litellm silently disables (no warning).""" monkeypatch.delenv("MELLEA_PRICING_ENABLED", raising=False) - monkeypatch.setitem(sys.modules, "litellm", None) - importlib.reload(pricing) + monkeypatch.setattr(pricing, "_LITELLM_AVAILABLE", False) + reset_pricing_state() assert pricing._PRICING_ENABLED is False assert not any("litellm is not installed" in str(w.message) for w in recwarn.list) diff --git a/test/telemetry/test_tracing.py b/test/telemetry/test_tracing.py index 3e8f4ab37..f21ff004b 100644 --- a/test/telemetry/test_tracing.py +++ b/test/telemetry/test_tracing.py @@ -15,24 +15,16 @@ tracing, ) from mellea.telemetry.tracing import get_backend_tracer - - -def _reset_tracing_state() -> None: - """Reset module state and re-run setup so env-var changes take effect.""" - tracing._tracer_provider = None - tracing._application_tracer = None - tracing._backend_tracer = None - tracing._in_flight_spans.clear() - tracing._setup_tracing() +from test.telemetry.conftest import reset_tracing_state @pytest.fixture def enable_tracing(monkeypatch): """Enable tracing for the duration of a test.""" monkeypatch.setenv("MELLEA_TRACES_ENABLED", "true") - _reset_tracing_state() + reset_tracing_state() yield - _reset_tracing_state() + reset_tracing_state() @pytest.fixture @@ -41,9 +33,9 @@ def disable_tracing(monkeypatch): monkeypatch.delenv("MELLEA_TRACES_ENABLED", raising=False) monkeypatch.delenv("MELLEA_TRACES_OTLP", raising=False) monkeypatch.delenv("MELLEA_TRACES_CONSOLE", raising=False) - _reset_tracing_state() + reset_tracing_state() yield - _reset_tracing_state() + reset_tracing_state() def test_telemetry_disabled_by_default(disable_tracing): @@ -73,7 +65,7 @@ def test_content_tracing(monkeypatch, env, expected): ) for k, v in env.items(): monkeypatch.setenv(k, v) - _reset_tracing_state() + reset_tracing_state() assert is_content_tracing_enabled() is expected @@ -84,10 +76,10 @@ def test_otlp_traces_endpoint_honored(monkeypatch): monkeypatch.setenv("MELLEA_TRACES_OTLP", "true") monkeypatch.setenv("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", "http://localhost:4317") monkeypatch.delenv("OTEL_EXPORTER_OTLP_ENDPOINT", raising=False) - _reset_tracing_state() + reset_tracing_state() assert get_backend_tracer() is not None - _reset_tracing_state() + reset_tracing_state() def test_otlp_warns_when_no_endpoint_configured(monkeypatch, recwarn): @@ -120,7 +112,7 @@ def test_otlp_falls_back_to_generic_endpoint(monkeypatch, recwarn): def test_trace_application_context_manager(): """Test that trace_application works as a context manager.""" - _reset_tracing_state() + reset_tracing_state() # Should not raise even when tracing is disabled with trace_application("test_span", test_attr="value") as span: diff --git a/test/telemetry/test_tracing_backend.py b/test/telemetry/test_tracing_backend.py index ca80ec0e1..4a564569f 100644 --- a/test/telemetry/test_tracing_backend.py +++ b/test/telemetry/test_tracing_backend.py @@ -12,6 +12,7 @@ ) from mellea.stdlib.components import Message from mellea.stdlib.context import SimpleContext +from test.telemetry.conftest import reset_tracing_state # Check if OpenTelemetry is available try: @@ -32,27 +33,17 @@ ] -def _reset_tracing_state() -> None: - import mellea.telemetry.tracing as tracing_mod - - tracing_mod._tracer_provider = None - tracing_mod._application_tracer = None - tracing_mod._backend_tracer = None - tracing_mod._in_flight_spans.clear() - tracing_mod._setup_tracing() - - @pytest.fixture(scope="module", autouse=True) def setup_telemetry(): """Enable tracing for all tests in this module.""" mp = pytest.MonkeyPatch() mp.setenv("MELLEA_TRACES_ENABLED", "true") - _reset_tracing_state() + reset_tracing_state() yield mp.undo() - _reset_tracing_state() + reset_tracing_state() @pytest.fixture diff --git a/test/telemetry/test_tracing_plugins.py b/test/telemetry/test_tracing_plugins.py index 63e84e9a4..381f15afa 100644 --- a/test/telemetry/test_tracing_plugins.py +++ b/test/telemetry/test_tracing_plugins.py @@ -21,17 +21,7 @@ ) from mellea.telemetry import tracing from mellea.telemetry.tracing_plugins import BackendTracingPlugin - - -def _reset_tracing_state() -> None: - """Reset module state and re-run setup so env-var changes take effect.""" - if tracing._tracer_provider is not None: - tracing._tracer_provider.shutdown() - tracing._tracer_provider = None - tracing._application_tracer = None - tracing._backend_tracer = None - tracing._in_flight_spans.clear() - tracing._setup_tracing() +from test.telemetry.conftest import reset_tracing_state @pytest.fixture @@ -42,17 +32,17 @@ def plugin(): @pytest.fixture def enabled_tracing(monkeypatch): monkeypatch.setenv("MELLEA_TRACES_ENABLED", "true") - _reset_tracing_state() + reset_tracing_state() yield - _reset_tracing_state() + reset_tracing_state() @pytest.fixture def disabled_tracing(monkeypatch): monkeypatch.delenv("MELLEA_TRACES_ENABLED", raising=False) - _reset_tracing_state() + reset_tracing_state() yield - _reset_tracing_state() + reset_tracing_state() def _attrs(span: MagicMock) -> dict: