add tests

Author: nikhilc-microsoft

tests/observability/core/exporters/test_enriched_span.py

@@ -0,0 +1,65 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Tests for EnrichedReadableSpan."""
+
+import unittest
+from unittest.mock import Mock
+
+from microsoft_agents_a365.observability.core.exporters.enriched_span import EnrichedReadableSpan
+
+
+class TestEnrichedReadableSpan(unittest.TestCase):
+    """Test suite for EnrichedReadableSpan."""
+
+    def test_attributes_merges_original_and_extra(self):
+        """Test that attributes property merges original span attributes with extra attributes."""
+        # Create mock span with original attributes
+        mock_span = Mock()
+        mock_span.attributes = {"original_key": "original_value", "shared_key": "original"}
+
+        # Create enriched span with extra attributes
+        extra_attributes = {"extra_key": "extra_value", "shared_key": "overwritten"}
+        enriched_span = EnrichedReadableSpan(mock_span, extra_attributes)
+
+        # Verify merged attributes
+        attributes = enriched_span.attributes
+        self.assertEqual(attributes["original_key"], "original_value")
+        self.assertEqual(attributes["extra_key"], "extra_value")
+        self.assertEqual(attributes["shared_key"], "overwritten")  # Extra should overwrite original
+
+    def test_delegates_all_properties_to_wrapped_span(self):
+        """Test that all span properties are delegated to the wrapped span."""
+        # Create mock span with all properties
+        mock_span = Mock()
+        mock_span.name = "test-span"
+        mock_span.context = Mock(trace_id=123, span_id=456)
+        mock_span.parent = Mock(span_id=789)
+        mock_span.start_time = 1000000000
+        mock_span.end_time = 2000000000
+        mock_span.status = Mock(status_code="OK", description=None)
+        mock_span.kind = "INTERNAL"
+        mock_span.events = []
+        mock_span.links = []
+        mock_span.resource = Mock(attributes={"service.name": "test"})
+        mock_span.instrumentation_scope = Mock(name="test-scope")
+        mock_span.attributes = {}
+
+        enriched_span = EnrichedReadableSpan(mock_span, {})
+
+        # Verify all properties delegate correctly
+        self.assertEqual(enriched_span.name, "test-span")
+        self.assertEqual(enriched_span.context, mock_span.context)
+        self.assertEqual(enriched_span.parent, mock_span.parent)
+        self.assertEqual(enriched_span.start_time, 1000000000)
+        self.assertEqual(enriched_span.end_time, 2000000000)
+        self.assertEqual(enriched_span.status, mock_span.status)
+        self.assertEqual(enriched_span.kind, "INTERNAL")
+        self.assertEqual(enriched_span.events, [])
+        self.assertEqual(enriched_span.links, [])
+        self.assertEqual(enriched_span.resource, mock_span.resource)
+        self.assertEqual(enriched_span.instrumentation_scope, mock_span.instrumentation_scope)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/observability/core/exporters/test_enriching_span_processor.py

@@ -0,0 +1,142 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Tests for enriching_span_processor module."""
+
+import unittest
+from unittest.mock import Mock
+
+from microsoft_agents_a365.observability.core.exporters.enriching_span_processor import (
+    _EnrichingBatchSpanProcessor,
+    _span_enrichers,
+    register_span_enricher,
+    unregister_span_enricher,
+)
+
+
+class TestSpanEnricherRegistry(unittest.TestCase):
+    """Test suite for span enricher registration functions."""
+
+    def setUp(self):
+        """Clear enrichers before each test."""
+        _span_enrichers.clear()
+
+    def tearDown(self):
+        """Clear enrichers after each test."""
+        _span_enrichers.clear()
+
+    def test_register_and_unregister_enricher(self):
+        """Test that enrichers can be registered and unregistered."""
+
+        # Define a simple enricher
+        def my_enricher(span):
+            return span
+
+        # Register
+        register_span_enricher(my_enricher)
+        self.assertIn(my_enricher, _span_enrichers)
+        self.assertEqual(len(_span_enrichers), 1)
+
+        # Duplicate registration should not add again
+        register_span_enricher(my_enricher)
+        self.assertEqual(len(_span_enrichers), 1)
+
+        # Unregister
+        unregister_span_enricher(my_enricher)
+        self.assertNotIn(my_enricher, _span_enrichers)
+        self.assertEqual(len(_span_enrichers), 0)
+
+    def test_unregister_nonexistent_enricher_does_not_raise(self):
+        """Test that unregistering a non-existent enricher doesn't raise an error."""
+
+        def my_enricher(span):
+            return span
+
+        # Should not raise
+        unregister_span_enricher(my_enricher)
+        self.assertEqual(len(_span_enrichers), 0)
+
+
+class TestEnrichingBatchSpanProcessor(unittest.TestCase):
+    """Test suite for _EnrichingBatchSpanProcessor."""
+
+    def setUp(self):
+        """Clear enrichers before each test."""
+        _span_enrichers.clear()
+
+    def tearDown(self):
+        """Clear enrichers after each test."""
+        _span_enrichers.clear()
+
+    def test_on_end_applies_enrichers_to_span(self):
+        """Test that on_end applies all registered enrichers to the span."""
+        # Create processor with a mock exporter
+        mock_exporter = Mock()
+        processor = _EnrichingBatchSpanProcessor(mock_exporter)
+
+        # Register an enricher that tracks what it receives and returns
+        received_spans = []
+
+        def enricher(span):
+            received_spans.append(span)
+            # Return a mock enriched span
+            enriched = Mock(name="enriched_span")
+            enriched.context = span.context
+            return enriched
+
+        register_span_enricher(enricher)
+
+        # Create a mock span
+        original_span = Mock(name="original_span")
+        original_span.context = Mock()
+        original_span.context.trace_id = 123
+        original_span.context.span_id = 456
+
+        # Call on_end
+        processor.on_end(original_span)
+
+        # Verify enricher was called with the original span
+        self.assertEqual(len(received_spans), 1)
+        self.assertEqual(received_spans[0], original_span)
+
+        # Cleanup
+        processor.shutdown()
+
+    def test_on_end_continues_if_enricher_raises_exception(self):
+        """Test that on_end continues processing even if an enricher raises an exception."""
+        mock_exporter = Mock()
+        processor = _EnrichingBatchSpanProcessor(mock_exporter)
+
+        # Track which enrichers were called
+        called_enrichers = []
+
+        def failing_enricher(span):
+            called_enrichers.append("failing")
+            raise ValueError("Enricher failed!")
+
+        def succeeding_enricher(span):
+            called_enrichers.append("succeeding")
+            return span
+
+        register_span_enricher(failing_enricher)
+        register_span_enricher(succeeding_enricher)
+
+        # Create a mock span
+        original_span = Mock(name="original_span")
+        original_span.context = Mock()
+        original_span.context.trace_id = 123
+        original_span.context.span_id = 456
+
+        # Should not raise despite failing enricher
+        processor.on_end(original_span)
+
+        # Verify both enrichers were called (failing one didn't stop the chain)
+        self.assertIn("failing", called_enrichers)
+        self.assertIn("succeeding", called_enrichers)
+
+        # Cleanup
+        processor.shutdown()
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/observability/extensions/semantickernel/test_span_enricher.py

@@ -0,0 +1,58 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Tests for Semantic Kernel span enricher."""
+
+import unittest
+from unittest.mock import Mock
+
+from microsoft_agents_a365.observability.core.constants import (
+    GEN_AI_INPUT_MESSAGES_KEY,
+    GEN_AI_OUTPUT_MESSAGES_KEY,
+)
+from microsoft_agents_a365.observability.extensions.semantickernel.span_enricher import (
+    enrich_semantic_kernel_span,
+)
+
+
+class TestSemanticKernelSpanEnricher(unittest.TestCase):
+    """Test suite for enrich_semantic_kernel_span function."""
+
+    def test_invoke_agent_span_extracts_content_from_messages(self):
+        """Test that invoke_agent spans have content extracted from input/output messages."""
+        # Create a mock span with invoke_agent name and message attributes
+        mock_span = Mock()
+        mock_span.name = "invoke_agent test-agent"
+        mock_span.attributes = {
+            "gen_ai.agent.invocation_input": '[{"role": "user", "content": "Hello"}]',
+            "gen_ai.agent.invocation_output": '[{"role": "assistant", "content": "Hi there!"}]',
+        }
+
+        # Enrich the span
+        enriched = enrich_semantic_kernel_span(mock_span)
+
+        # Verify it returns an EnrichedReadableSpan with extracted content
+        self.assertNotEqual(enriched, mock_span)
+        attributes = enriched.attributes
+        # extract_content_as_string_list returns a JSON string
+        self.assertEqual(attributes[GEN_AI_INPUT_MESSAGES_KEY], '["Hello"]')
+        self.assertEqual(attributes[GEN_AI_OUTPUT_MESSAGES_KEY], '["Hi there!"]')
+
+    def test_non_matching_span_returns_original(self):
+        """Test that spans not matching invoke_agent or execute_tool are returned unchanged."""
+        # Create a mock span with a different operation name
+        mock_span = Mock()
+        mock_span.name = "some_other_operation"
+        mock_span.attributes = {
+            "some_key": "some_value",
+        }
+
+        # Enrich the span
+        result = enrich_semantic_kernel_span(mock_span)
+
+        # Verify it returns the original span unchanged
+        self.assertEqual(result, mock_span)
+
+
+if __name__ == "__main__":
+    unittest.main()