|
| 1 | +# Copyright (c) Microsoft Corporation. |
| 2 | +# Licensed under the MIT License. |
| 3 | + |
| 4 | +"""Tests that observability code handles NonRecordingSpan gracefully. |
| 5 | +
|
| 6 | +When the TracerProvider has no valid exporter (e.g. token resolver returns |
| 7 | +None on the first turn), it produces NonRecordingSpan instances. These only |
| 8 | +expose get_span_context() — the .context attribute does not exist. The code |
| 9 | +must use get_span_context() everywhere to avoid AttributeError crashes. |
| 10 | +""" |
| 11 | + |
| 12 | +import os |
| 13 | +import unittest |
| 14 | +from unittest.mock import MagicMock, patch |
| 15 | + |
| 16 | +from microsoft_agents_a365.observability.core.exporters.enriched_span import ( |
| 17 | + EnrichedReadableSpan, |
| 18 | +) |
| 19 | +from microsoft_agents_a365.observability.core.opentelemetry_scope import ( |
| 20 | + OpenTelemetryScope, |
| 21 | +) |
| 22 | +from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags |
| 23 | + |
| 24 | + |
| 25 | +def _make_non_recording_span(trace_id: int = 0, span_id: int = 0) -> NonRecordingSpan: |
| 26 | + """Create a NonRecordingSpan with the given trace/span IDs.""" |
| 27 | + return NonRecordingSpan( |
| 28 | + SpanContext( |
| 29 | + trace_id=trace_id, |
| 30 | + span_id=span_id, |
| 31 | + is_remote=False, |
| 32 | + trace_flags=TraceFlags(0), |
| 33 | + ) |
| 34 | + ) |
| 35 | + |
| 36 | + |
| 37 | +class TestOpenTelemetryScopeNonRecordingSpan(unittest.TestCase): |
| 38 | + """Verify OpenTelemetryScope works when the tracer returns a NonRecordingSpan.""" |
| 39 | + |
| 40 | + @patch.dict(os.environ, {"ENABLE_OBSERVABILITY": "true"}) |
| 41 | + @patch.object(OpenTelemetryScope, "_get_tracer") |
| 42 | + def test_init_with_non_recording_span(self, mock_get_tracer: MagicMock) -> None: |
| 43 | + """__init__ must not crash when tracer.start_span returns NonRecordingSpan.""" |
| 44 | + mock_tracer = MagicMock() |
| 45 | + nr_span = _make_non_recording_span(trace_id=0xABCD, span_id=0x1234) |
| 46 | + mock_tracer.start_span.return_value = nr_span |
| 47 | + mock_get_tracer.return_value = mock_tracer |
| 48 | + |
| 49 | + # Should not raise AttributeError |
| 50 | + scope = OpenTelemetryScope( |
| 51 | + operation_name="invoke_agent", |
| 52 | + activity_name="test_activity", |
| 53 | + ) |
| 54 | + |
| 55 | + self.assertIs(scope._span, nr_span) |
| 56 | + |
| 57 | + @patch.dict(os.environ, {"ENABLE_OBSERVABILITY": "true"}) |
| 58 | + @patch.object(OpenTelemetryScope, "_get_tracer") |
| 59 | + def test_end_with_non_recording_span(self, mock_get_tracer: MagicMock) -> None: |
| 60 | + """_end() must not crash when the span is a NonRecordingSpan.""" |
| 61 | + mock_tracer = MagicMock() |
| 62 | + nr_span = _make_non_recording_span(trace_id=0xABCD, span_id=0x1234) |
| 63 | + mock_tracer.start_span.return_value = nr_span |
| 64 | + mock_get_tracer.return_value = mock_tracer |
| 65 | + |
| 66 | + scope = OpenTelemetryScope( |
| 67 | + operation_name="invoke_agent", |
| 68 | + activity_name="test_activity", |
| 69 | + ) |
| 70 | + |
| 71 | + # Should not raise AttributeError |
| 72 | + scope._end() |
| 73 | + |
| 74 | + @patch.dict(os.environ, {"ENABLE_OBSERVABILITY": "true"}) |
| 75 | + @patch.object(OpenTelemetryScope, "_get_tracer") |
| 76 | + def test_dispose_with_non_recording_span(self, mock_get_tracer: MagicMock) -> None: |
| 77 | + """Full dispose lifecycle must not crash with NonRecordingSpan.""" |
| 78 | + mock_tracer = MagicMock() |
| 79 | + nr_span = _make_non_recording_span(trace_id=0xABCD, span_id=0x1234) |
| 80 | + mock_tracer.start_span.return_value = nr_span |
| 81 | + mock_get_tracer.return_value = mock_tracer |
| 82 | + |
| 83 | + scope = OpenTelemetryScope( |
| 84 | + operation_name="invoke_agent", |
| 85 | + activity_name="test_activity", |
| 86 | + ) |
| 87 | + |
| 88 | + # Context manager exit should not raise |
| 89 | + scope.__exit__(None, None, None) |
| 90 | + |
| 91 | + |
| 92 | +class TestEnrichedReadableSpanNonRecordingSpan(unittest.TestCase): |
| 93 | + """Verify EnrichedReadableSpan delegates context via get_span_context().""" |
| 94 | + |
| 95 | + def test_context_property_uses_get_span_context(self) -> None: |
| 96 | + """EnrichedReadableSpan.context must call get_span_context(), not .context.""" |
| 97 | + mock_span = MagicMock() |
| 98 | + expected_ctx = SpanContext( |
| 99 | + trace_id=0xFF, |
| 100 | + span_id=0xAA, |
| 101 | + is_remote=False, |
| 102 | + trace_flags=TraceFlags(0), |
| 103 | + ) |
| 104 | + mock_span.get_span_context.return_value = expected_ctx |
| 105 | + # Ensure .context attribute is NOT used |
| 106 | + del mock_span.context |
| 107 | + |
| 108 | + enriched = EnrichedReadableSpan(mock_span, extra_attributes={"key": "val"}) |
| 109 | + |
| 110 | + result = enriched.context |
| 111 | + self.assertIs(result, expected_ctx) |
| 112 | + mock_span.get_span_context.assert_called_once() |
| 113 | + |
| 114 | + def test_to_json_with_non_recording_span_context(self) -> None: |
| 115 | + """to_json() must not crash when wrapped span uses get_span_context().""" |
| 116 | + mock_span = MagicMock() |
| 117 | + mock_span.name = "test" |
| 118 | + mock_span.get_span_context.return_value = SpanContext( |
| 119 | + trace_id=0x1234, |
| 120 | + span_id=0x5678, |
| 121 | + is_remote=False, |
| 122 | + trace_flags=TraceFlags(0), |
| 123 | + ) |
| 124 | + del mock_span.context |
| 125 | + mock_span.attributes = {"key": "value"} |
| 126 | + mock_span.kind = "INTERNAL" |
| 127 | + mock_span.parent = None |
| 128 | + mock_span.start_time = 1000000000 |
| 129 | + mock_span.end_time = 2000000000 |
| 130 | + mock_span.status = None |
| 131 | + mock_span.events = [] |
| 132 | + mock_span.links = [] |
| 133 | + mock_span.resource = None |
| 134 | + |
| 135 | + enriched = EnrichedReadableSpan(mock_span, extra_attributes={}) |
| 136 | + |
| 137 | + # Should not raise |
| 138 | + json_str = enriched.to_json() |
| 139 | + self.assertIn("test", json_str) |
| 140 | + self.assertIn("0x00000000000000000000000000001234", json_str) |
| 141 | + |
| 142 | + |
| 143 | +class TestAgent365ExporterMapSpanGetSpanContext(unittest.TestCase): |
| 144 | + """Verify _map_span uses get_span_context() instead of .context.""" |
| 145 | + |
| 146 | + @patch.dict(os.environ, {}, clear=True) |
| 147 | + def test_map_span_calls_get_span_context(self) -> None: |
| 148 | + """_map_span must use get_span_context(), not .context attribute.""" |
| 149 | + from microsoft_agents_a365.observability.core.exporters.agent365_exporter import ( |
| 150 | + _Agent365Exporter, |
| 151 | + ) |
| 152 | + |
| 153 | + exporter = _Agent365Exporter(token_resolver=lambda a, t: "token") |
| 154 | + |
| 155 | + span = MagicMock() |
| 156 | + span.name = "test_span" |
| 157 | + span.attributes = {"gen_ai.operation.name": "invoke_agent"} |
| 158 | + |
| 159 | + expected_ctx = SpanContext( |
| 160 | + trace_id=0xABCD, |
| 161 | + span_id=0x1234, |
| 162 | + is_remote=False, |
| 163 | + trace_flags=TraceFlags(0), |
| 164 | + ) |
| 165 | + span.get_span_context.return_value = expected_ctx |
| 166 | + # Remove .context so it would fail if accessed |
| 167 | + del span.context |
| 168 | + |
| 169 | + span.parent = None |
| 170 | + span.kind = MagicMock() |
| 171 | + span.start_time = 1000000000 |
| 172 | + span.end_time = 2000000000 |
| 173 | + span.status = MagicMock() |
| 174 | + span.status.status_code = MagicMock() |
| 175 | + span.status.description = "" |
| 176 | + span.events = [] |
| 177 | + span.links = [] |
| 178 | + span.instrumentation_scope = MagicMock() |
| 179 | + span.instrumentation_scope.name = "test" |
| 180 | + span.instrumentation_scope.version = "1.0" |
| 181 | + |
| 182 | + # Should not raise AttributeError |
| 183 | + mapped = exporter._map_span(span) |
| 184 | + |
| 185 | + span.get_span_context.assert_called_once() |
| 186 | + self.assertIn("traceId", mapped) |
| 187 | + self.assertIn("spanId", mapped) |
| 188 | + exporter.shutdown() |
| 189 | + |
| 190 | + |
| 191 | +if __name__ == "__main__": |
| 192 | + unittest.main() |
0 commit comments