feat: add opentelemetry metrics support for worker backend (#72)

Author: leoparente

worker/pyproject.toml

@@ -33,11 +33,14 @@ dependencies = [
     "pydantic~=2.9",
     "uvicorn~=0.32",
     "PyYAML~=6.0",
+    "opentelemetry-api~=1.32",
+    "opentelemetry-sdk~=1.32",
+    "opentelemetry-exporter-otlp~=1.32",
 ]
 
 [project.optional-dependencies]
 dev = ["black", "check-manifest", "ruff"]
-test = ["coverage", "pytest", "pytest-cov"]
+test = ["coverage", "pytest", "pytest-cov==6.0.0"]
 
 [project.urls]
 "Homepage" = "https://netboxlabs.com/"

worker/tests/policy/test_runner.py

@@ -8,7 +8,7 @@
 from apscheduler.triggers.date import DateTrigger
 
 from worker.backend import Backend
-from worker.models import Config, DiodeConfig, Policy, Status
+from worker.models import Config, DiodeConfig, Metadata, Policy, Status
 from worker.policy.runner import PolicyRunner
 
 
@@ -151,7 +151,8 @@ def test_run_ingestion_errors(
     mock_backend.run.assert_called_once_with(policy_runner.name, sample_policy)
     mock_diode_client.ingest.assert_called_once_with(mock_backend.run.return_value)
     assert (
-        "ERROR ingestion failed for test_policy : ['error1', 'error2']" in caplog.text
+        "Policy test_policy: Ingestion failed with errors: ['error1', 'error2']"
+        in caplog.text
     )
 
 
@@ -186,3 +187,108 @@ def test_stop_policy_runner(policy_runner):
         # Ensure scheduler shutdown is called and status is updated
         mock_shutdown.assert_called_once()
         assert policy_runner.status == Status.FINISHED
+
+
+def test_metrics_during_policy_lifecycle(
+    policy_runner, sample_policy, mock_diode_client, mock_backend
+):
+    """Test that metrics are properly updated during the policy lifecycle."""
+    # Create mock metrics
+    mock_active_policies = MagicMock()
+    mock_policy_executions = MagicMock()
+    mock_backend_execution_success = MagicMock()
+    mock_backend_execution_failure = MagicMock()
+    mock_backend_execution_latency = MagicMock()
+
+    # Map of metric names to mock objects
+    mock_metrics = {
+        "active_policies": mock_active_policies,
+        "policy_executions": mock_policy_executions,
+        "backend_execution_success": mock_backend_execution_success,
+        "backend_execution_failure": mock_backend_execution_failure,
+        "backend_execution_latency": mock_backend_execution_latency,
+    }
+
+    policy_runner.name = "test_policy"
+    policy_runner.metadata = Metadata(
+        name="my_backend",
+        app_name="test_app",
+        app_version="1.0",
+    )
+
+    # Setup mock for get_metric function
+    def mock_get_metric(name):
+        return mock_metrics.get(name)
+
+    with patch("worker.policy.runner.get_metric", side_effect=mock_get_metric):
+
+        mock_diode_client.ingest.return_value.errors = []
+
+        policy_runner.run(mock_diode_client, mock_backend, sample_policy)
+
+        mock_backend.run.assert_called_once_with(policy_runner.name, sample_policy)
+        mock_diode_client.ingest.assert_called_once_with(mock_backend.run.return_value)
+
+        mock_policy_executions.add.assert_called_once_with(1, {"policy": "test_policy"})
+        mock_backend_execution_success.add.assert_called_once_with(
+            1,
+            {
+                "policy": "test_policy",
+                "backend": "my_backend",
+                "app_name": "test_app",
+                "app_version": "1.0",
+            },
+        )
+
+        # Test stop - should decrement active_policies
+        with patch.object(policy_runner.scheduler, "shutdown") as mock_shutdown:
+            policy_runner.stop()
+            mock_shutdown.assert_called_once()
+            mock_active_policies.add.assert_called_with(-1, {"policy": "test_policy"})
+
+
+def test_metrics_during_failed_discovery(
+    policy_runner, sample_policy, mock_diode_client, mock_backend
+):
+    """Test that metrics are properly updated when discovery fails."""
+    mock_backend_execution_failure = MagicMock()
+    mock_backend_execution_latency = MagicMock()
+
+    mock_metrics = {
+        "backend_execution_failure": mock_backend_execution_failure,
+        "backend_execution_latency": mock_backend_execution_latency,
+    }
+
+    policy_runner.name = "test_policy"
+    policy_runner.metadata = Metadata(
+        name="my_backend",
+        app_name="test_app",
+        app_version="1.0",
+    )
+
+    def mock_get_metric(name):
+        return mock_metrics.get(name)
+
+    # Simulate backend throwing an exception
+    mock_backend.run.side_effect = Exception("Backend error")
+
+    with patch("worker.policy.runner.get_metric", side_effect=mock_get_metric):
+        mock_diode_client = MagicMock(name="MockDiodeClient")
+        policy_runner.run(mock_diode_client, sample_diode_config, sample_policy)
+        # Verify failure metric was called
+        mock_backend_execution_failure.add.assert_called_once_with(
+            1,
+            {
+                "policy": "test_policy",
+                "backend": "my_backend",
+                "app_name": "test_app",
+                "app_version": "1.0",
+            },
+        )
+
+        # Verify backend execution latency recorded with failure status
+        mock_backend_execution_latency.record.assert_called_once()
+        latency_args = mock_backend_execution_latency.record.call_args[0][0]
+        latency_kwargs = mock_backend_execution_latency.record.call_args[0][1]
+        assert latency_args > 0
+        assert latency_kwargs["backend"] == "my_backend"

worker/tests/test_metrics.py

@@ -0,0 +1,211 @@
+#!/usr/bin/env python
+# Copyright 2025 NetBox Labs Inc
+"""NetBox Labs - Metrics Unit Tests."""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from worker.metrics import get_metric, setup_metrics_export
+
+
+@pytest.fixture
+def mock_opentelemetry():
+    """Mock the OpenTelemetry SDK imports and components."""
+    with patch("worker.metrics.OTLPMetricExporter") as mock_exporter, patch(
+        "worker.metrics.PeriodicExportingMetricReader"
+    ) as mock_reader, patch("opentelemetry.sdk.metrics.Meter") as mock_meter, patch(
+        "worker.metrics.MeterProvider"
+    ) as mock_provider:
+        # Setup return values
+        mock_provider.return_value.get_meter.return_value = mock_meter
+
+        yield {
+            "exporter": mock_exporter,
+            "reader": mock_reader,
+            "meter": mock_meter,
+            "provider": mock_provider,
+        }
+
+
+@pytest.fixture
+def reset_metrics_cache():
+    """Reset the metrics cache before test."""
+    with patch("worker.metrics._metrics_cache", {}), \
+         patch("worker.metrics._metric_factories", {}), \
+         patch("worker.metrics._metrics_enabled", True):
+        yield
+
+
+def test_setup_metrics_export(mock_opentelemetry):
+    """Test that metrics export setup creates the correct OpenTelemetry components."""
+    endpoint = "http://localhost:4317"
+    export_period = 30
+
+    setup_metrics_export(endpoint, export_period)
+
+    # Verify exporter was created with correct endpoint
+    mock_opentelemetry["exporter"].assert_called_once()
+    args, kwargs = mock_opentelemetry["exporter"].call_args
+    assert kwargs["endpoint"] == endpoint
+
+    # Verify reader was created with correct exporter and export interval
+    mock_opentelemetry["reader"].assert_called_once()
+    args, kwargs = mock_opentelemetry["reader"].call_args
+    assert kwargs["export_interval_millis"] == export_period * 1000
+
+    # Verify meter provider was configured
+    mock_opentelemetry["provider"].assert_called_once()
+
+    # Verify meter was created
+    mock_opentelemetry["provider"].return_value.get_meter.assert_called_once_with(
+        "device-discovery", "0.0.0"
+    )
+
+
+def test_setup_metrics_export_no_endpoint(mock_opentelemetry, reset_metrics_cache):
+    """Test that metrics export setup is properly disabled when no endpoint is provided."""
+    with patch("worker.metrics.logger") as mock_logger:
+        # Call with None endpoint
+        setup_metrics_export(None, 30)
+
+        # Verify logger message
+        mock_logger.info.assert_called_once_with(
+            "No metrics endpoint provided, metrics collection is disabled"
+        )
+
+        # Verify no OpenTelemetry components were created
+        mock_opentelemetry["exporter"].assert_not_called()
+        mock_opentelemetry["reader"].assert_not_called()
+        mock_opentelemetry["provider"].assert_not_called()
+
+        # Verify get_metric returns None after setup with no endpoint
+        metric = get_metric("api_requests")
+        assert metric is None
+
+
+def test_get_metric_returns_counter(reset_metrics_cache):
+    """Test that get_metric returns a counter for counter-type metrics."""
+    mock_counter = MagicMock()
+    mock_meter = MagicMock()
+    mock_meter.create_counter.return_value = mock_counter
+
+    with patch("worker.metrics._meter", mock_meter):
+        # Test accessing a counter metric
+        metric = get_metric("api_requests")
+
+        # Verify counter was created with correct name and description
+        mock_meter.create_counter.assert_called_once()
+        args, kwargs = mock_meter.create_counter.call_args
+        assert kwargs["name"] == "api_requests"
+        assert "description" in kwargs
+
+        # Should return the mock counter
+        assert metric == mock_counter
+
+
+def test_get_metric_returns_histogram(reset_metrics_cache):
+    """Test that get_metric returns a histogram for latency-type metrics."""
+    mock_histogram = MagicMock()
+    mock_meter = MagicMock()
+    mock_meter.create_histogram.return_value = mock_histogram
+
+    with patch("worker.metrics._meter", mock_meter):
+        # Test accessing a histogram metric
+        metric = get_metric("api_response_latency")
+
+        # Verify histogram was created with correct name and description
+        mock_meter.create_histogram.assert_called_once()
+        args, kwargs = mock_meter.create_histogram.call_args
+        assert kwargs["name"] == "api_response_latency"
+        assert "description" in kwargs
+
+        # Should return the mock histogram
+        assert metric == mock_histogram
+
+
+def test_get_metric_returns_none_when_not_initialized():
+    """Test that get_metric returns None when metrics are not initialized."""
+    with patch("worker.metrics._meter", None):
+        metric = get_metric("api_requests")
+        assert metric is None
+
+
+def test_get_metric_creates_metric_only_once(reset_metrics_cache):
+    """Test that get_metric only creates a metric once and returns cached value."""
+    mock_counter = MagicMock()
+    mock_meter = MagicMock()
+    mock_meter.create_counter.return_value = mock_counter
+
+    with patch("worker.metrics._meter", mock_meter), patch(
+        "worker.metrics._metrics_cache", {}
+    ):
+
+        # First call should create the metric
+        metric1 = get_metric("api_requests")
+        assert metric1 == mock_counter
+        mock_meter.create_counter.assert_called_once()
+
+        # Reset the mock to check if it's called again
+        mock_meter.create_counter.reset_mock()
+
+        # Second call should return cached metric without creating it again
+        metric2 = get_metric("api_requests")
+        assert metric2 == mock_counter
+        mock_meter.create_counter.assert_not_called()
+
+
+def test_all_expected_metrics_exist(reset_metrics_cache):
+    """Test that all expected metrics can be retrieved."""
+    expected_metrics = [
+        "api_requests",
+        "api_response_latency",
+        "active_policies",
+        "policy_executions",
+        "backend_execution_success",
+        "backend_execution_failure",
+        "backend_execution_latency",
+    ]
+
+    mock_meter = MagicMock()
+    mock_meter.create_counter.return_value = MagicMock()
+    mock_meter.create_histogram.return_value = MagicMock()
+
+    with patch("worker.metrics._meter", mock_meter), patch(
+        "worker.metrics._metrics_cache", {}
+    ):
+
+        for metric_name in expected_metrics:
+            metric = get_metric(metric_name)
+            assert metric is not None, f"Expected metric {metric_name} to exist"
+
+
+def test_setup_metrics_export_meter_provider_error(mock_opentelemetry, reset_metrics_cache):
+    """Test handling of errors when setting the meter provider."""
+    endpoint = "http://localhost:4317"
+    export_period = 30
+
+    # Mock set_meter_provider to raise an exception
+    with patch("worker.metrics.otlp_metrics.set_meter_provider",
+              side_effect=Exception("Provider error")), \
+         patch("worker.metrics.logger") as mock_logger:
+
+        # Call function
+        setup_metrics_export(endpoint, export_period)
+
+        # Verify components were created but meter provider wasn't set
+        mock_opentelemetry["exporter"].assert_called_once()
+        mock_opentelemetry["reader"].assert_called_once()
+        mock_opentelemetry["provider"].assert_called_once()
+
+        # Verify warning was logged
+        mock_logger.warning.assert_called_once()
+        warning_message = mock_logger.warning.call_args[0][0]
+        assert "Could not set meter provider" in warning_message
+
+        # Verify meter was not created
+        mock_opentelemetry["provider"].return_value.get_meter.assert_not_called()
+
+        # Verify metrics are not enabled and get_metric returns None
+        metric = get_metric("api_requests")
+        assert metric is None

worker/tests/test_server.py

@@ -2,7 +2,7 @@
 # Copyright 2025 NetBox Labs Inc
 """NetBox Labs - Server Unit Tests."""
 
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 import yaml
@@ -139,11 +139,25 @@ def test_read_status(mock_version_semver):
         mock_version_semver: Mocked version_semver function.
 
     """
-    response = client.get("/api/v1/status")
-    mock_version_semver.assert_called_once()
-    assert response.status_code == 200
-    assert response.json()["version"] == "1.0.0"
-    assert "up_time_seconds" in response.json()
+    mock_api_requests = MagicMock()
+    mock_api_response_latency = MagicMock()
+
+    mock_metrics = {
+        "api_requests": mock_api_requests,
+        "api_response_latency": mock_api_response_latency,
+    }
+
+    def mock_get_metric(name):
+        return mock_metrics.get(name)
+
+    with patch("worker.server.get_metric", side_effect=mock_get_metric):
+        response = client.get("/api/v1/status")
+        mock_version_semver.assert_called_once()
+        assert response.status_code == 200
+        assert response.json()["version"] == "1.0.0"
+        assert "up_time_seconds" in response.json()
+        assert mock_api_requests.add.call_count == 1
+        assert mock_api_response_latency.record.call_count == 1
 
 
 def test_read_capabilities():