amq-adapter-python/tests/adapter/test_trace_info_adapter.py

"""Unit tests for the tracing module."""

import unittest
from unittest.mock import MagicMock, patch

from opentelemetry.sdk.resources import Resource
from opentelemetry.trace import Tracer

from amqp.service.tracing import (
    MetricsConfig,
    TracingConfig,
    create_resource,
    create_span,
    initialize_telemetry,
    setup_metrics_provider,
    setup_otlp_exporter,
    setup_trace_provider,
)


class TestTracingConfig(unittest.TestCase):
    """Test cases for TracingConfig."""

    def test_default_values(self):
        """Test that default values are set correctly."""
        config = TracingConfig(service_name="test-service", service_version="1.0.0")
        self.assertEqual(config.service_name, "test-service")
        self.assertEqual(config.service_version, "1.0.0")
        self.assertEqual(config.otlp_endpoint, "http://localhost:4317")
        self.assertTrue(config.insecure)
        self.assertFalse(config.debug)
        self.assertIsNone(config.additional_attributes)

    def test_custom_values(self):
        """Test that custom values are set correctly."""
        config = TracingConfig(
            service_name="test-service",
            service_version="1.0.0",
            otlp_endpoint="http://custom:4317",
            insecure=False,
            debug=True,
            additional_attributes={"custom": "attribute"},
        )
        self.assertEqual(config.otlp_endpoint, "http://custom:4317")
        self.assertFalse(config.insecure)
        self.assertTrue(config.debug)
        self.assertEqual(config.additional_attributes, {"custom": "attribute"})


class TestMetricsConfig(unittest.TestCase):
    """Test cases for MetricsConfig."""

    def test_default_values(self):
        """Test that default values are set correctly."""
        config = MetricsConfig()
        self.assertEqual(config.prometheus_port, 9464)
        self.assertEqual(config.prometheus_host, "localhost")
        self.assertEqual(config.export_interval_millis, 30000)
        self.assertIsNone(config.additional_attributes)

    def test_custom_values(self):
        """Test that custom values are set correctly."""
        config = MetricsConfig(
            prometheus_port=8000,
            prometheus_host="0.0.0.0",
            export_interval_millis=15000,
            additional_attributes={"custom": "attribute"},
        )
        self.assertEqual(config.prometheus_port, 8000)
        self.assertEqual(config.prometheus_host, "0.0.0.0")
        self.assertEqual(config.export_interval_millis, 15000)
        self.assertEqual(config.additional_attributes, {"custom": "attribute"})


class TestCreateResource(unittest.TestCase):
    """Test cases for create_resource function."""

    def test_create_resource_basic(self):
        """Test creating a resource with basic configuration."""
        config = TracingConfig(service_name="test-service", service_version="1.0.0")
        resource = create_resource(config)
        self.assertIsInstance(resource, Resource)
        self.assertEqual(resource.attributes.get("service.name"), "test-service")
        self.assertEqual(resource.attributes.get("service.version"), "1.0.0")

    def test_create_resource_with_attributes(self):
        """Test creating a resource with additional attributes."""
        config = TracingConfig(
            service_name="test-service",
            service_version="1.0.0",
            additional_attributes={"custom": "attribute"},
        )
        resource = create_resource(config)
        self.assertEqual(resource.attributes.get("custom"), "attribute")


class TestSetupExporters(unittest.TestCase):
    """Test cases for exporter setup functions."""

    @patch("amqp.service.tracing.OTLPSpanExporter")
    def test_setup_otlp_exporter(self, mock_otlp):
        """Test OTLP exporter setup."""
        config = TracingConfig(
            service_name="test-service",
            service_version="1.0.0",
            otlp_endpoint="http://jaeger:4317",
            insecure=True,
        )
        setup_otlp_exporter(config)
        mock_otlp.assert_called_once_with(endpoint="http://jaeger:4317", insecure=True)

    @patch("amqp.service.tracing.OTLPSpanExporter")
    def test_setup_otlp_exporter_secure(self, mock_otlp):
        """Test OTLP exporter setup with secure connection."""
        config = TracingConfig(
            service_name="test-service",
            service_version="1.0.0",
            otlp_endpoint="https://jaeger:4317",
            insecure=False,
        )
        setup_otlp_exporter(config)
        mock_otlp.assert_called_once_with(endpoint="https://jaeger:4317", insecure=False)


class TestSetupTraceProvider(unittest.TestCase):
    """Test cases for trace provider setup."""

    @patch("amqp.service.tracing.TracerProvider")
    @patch("amqp.service.tracing.BatchSpanProcessor")
    @patch("amqp.service.tracing.setup_otlp_exporter")
    def test_setup_trace_provider_basic(self, mock_otlp_exporter, mock_processor, mock_provider):
        """Test basic trace provider setup."""
        config = TracingConfig(service_name="test-service", service_version="1.0.0")
        resource = create_resource(config)

        mock_provider_instance = MagicMock()
        mock_provider.return_value = mock_provider_instance

        setup_trace_provider(config, resource)

        mock_provider.assert_called_once_with(resource=resource)
        mock_otlp_exporter.assert_called_once()
        self.assertEqual(mock_provider_instance.add_span_processor.call_count, 1)

    @patch("amqp.service.tracing.TracerProvider")
    @patch("amqp.service.tracing.BatchSpanProcessor")
    @patch("amqp.service.tracing.setup_otlp_exporter")
    @patch("amqp.service.tracing.ConsoleSpanExporter")
    def test_setup_trace_provider_with_debug(
        self, mock_console_exporter, mock_otlp_exporter, mock_processor, mock_provider
    ):
        """Test trace provider setup with debug mode."""
        config = TracingConfig(service_name="test-service", service_version="1.0.0", debug=True)
        resource = create_resource(config)

        mock_provider_instance = MagicMock()
        mock_provider.return_value = mock_provider_instance

        setup_trace_provider(config, resource)

        self.assertEqual(mock_provider_instance.add_span_processor.call_count, 2)
        mock_console_exporter.assert_called_once()


class TestSetupMetricsProvider(unittest.TestCase):
    """Test cases for metrics provider setup."""

    @patch("amqp.service.tracing.start_http_server")
    @patch("amqp.service.tracing.PrometheusMetricReader")
    @patch("amqp.service.tracing.PeriodicExportingMetricReader")
    @patch("amqp.service.tracing.MeterProvider")
    def test_setup_metrics_provider_success(
        self, mock_meter_provider, mock_periodic_reader, mock_prometheus_reader, mock_start_server
    ):
        """Test successful metrics provider setup."""
        config = MetricsConfig(prometheus_port=9464, prometheus_host="localhost")
        resource = Resource({})

        setup_metrics_provider(config, resource)

        mock_start_server.assert_called_once_with(port=9464, addr="localhost")
        mock_prometheus_reader.assert_called_once()
        mock_periodic_reader.assert_called_once()
        mock_meter_provider.assert_called_once()

    @patch("amqp.service.tracing.start_http_server")
    def test_setup_metrics_provider_failure(self, mock_start_server):
        """Test metrics provider setup with failure."""
        mock_start_server.side_effect = Exception("Failed to start server")

        config = MetricsConfig()
        resource = Resource({})

        provider = setup_metrics_provider(config, resource)
        self.assertIsNotNone(provider)  # Should return no-op provider


class TestInitializeTelemetry(unittest.TestCase):
    """Test cases for initialize_telemetry function."""

    @patch("amqp.service.tracing.setup_trace_provider")
    @patch("amqp.service.tracing.setup_metrics_provider")
    @patch("amqp.service.tracing.trace")
    @patch("amqp.service.tracing.metrics")
    def test_initialize_telemetry_success(
        self, mock_metrics, mock_trace, mock_setup_metrics, mock_setup_trace
    ):
        """Test successful telemetry initialization."""
        config = TracingConfig(service_name="test-service", service_version="1.0.0")
        metrics_config = MetricsConfig()

        mock_tracer = MagicMock(spec=Tracer)
        mock_trace.get_tracer.return_value = mock_tracer

        tracer = initialize_telemetry(config, metrics_config)

        self.assertEqual(tracer, mock_tracer)
        mock_setup_trace.assert_called_once()
        mock_setup_metrics.assert_called_once()

    @patch("amqp.service.tracing.setup_trace_provider")
    def test_initialize_telemetry_failure(self, mock_setup_trace):
        """Test telemetry initialization with failure."""
        mock_setup_trace.side_effect = Exception("Setup failed")

        tracer = initialize_telemetry()

        self.assertIsNotNone(tracer)  # Should return no-op tracer


class TestCreateSpan(unittest.TestCase):
    """Test cases for create_span function."""

    def test_create_span_basic(self):
        """Test basic span creation."""
        mock_tracer = MagicMock(spec=Tracer)
        create_span(mock_tracer, "test-span")
        mock_tracer.start_as_current_span.assert_called_once_with("test-span", attributes={})

    def test_create_span_with_attributes(self):
        """Test span creation with attributes."""
        mock_tracer = MagicMock(spec=Tracer)
        attributes = {"custom": "attribute"}
        create_span(mock_tracer, "test-span", attributes)
        mock_tracer.start_as_current_span.assert_called_once_with(
            "test-span", attributes=attributes
        )


if __name__ == "__main__":
    unittest.main()