|
2 | 2 |
|
3 | 3 | import pytest |
4 | 4 |
|
| 5 | +from agent_control.telemetry.trace_context import clear_trace_context_provider, set_trace_context_provider |
5 | 6 | from agent_control.tracing import ( |
6 | 7 | _generate_span_id, |
7 | 8 | _generate_trace_id, |
|
17 | 18 | ) |
18 | 19 |
|
19 | 20 |
|
| 21 | +def teardown_function() -> None: |
| 22 | + clear_trace_context_provider() |
| 23 | + |
| 24 | + |
20 | 25 | class TestIdGeneration: |
21 | 26 | """Tests for trace and span ID generation.""" |
22 | 27 |
|
@@ -132,6 +137,30 @@ def test_get_current_ids_without_context(self): |
132 | 137 | assert trace_id is None or isinstance(trace_id, str) |
133 | 138 | assert span_id is None or isinstance(span_id, str) |
134 | 139 |
|
| 140 | + def test_get_current_trace_id_uses_provider(self): |
| 141 | + """Test that get_current_trace_id uses external provider before OTEL fallback.""" |
| 142 | + expected_trace = "a" * 32 |
| 143 | + set_trace_context_provider( |
| 144 | + lambda: { |
| 145 | + "trace_id": expected_trace, |
| 146 | + "span_id": "b" * 16, |
| 147 | + } |
| 148 | + ) |
| 149 | + |
| 150 | + assert get_current_trace_id() == expected_trace |
| 151 | + |
| 152 | + def test_get_current_span_id_uses_provider(self): |
| 153 | + """Test that get_current_span_id uses external provider before OTEL fallback.""" |
| 154 | + expected_span = "b" * 16 |
| 155 | + set_trace_context_provider( |
| 156 | + lambda: { |
| 157 | + "trace_id": "a" * 32, |
| 158 | + "span_id": expected_span, |
| 159 | + } |
| 160 | + ) |
| 161 | + |
| 162 | + assert get_current_span_id() == expected_span |
| 163 | + |
135 | 164 |
|
136 | 165 | class TestWithTraceContextManager: |
137 | 166 | """Tests for the with_trace context manager.""" |
@@ -237,6 +266,23 @@ def test_get_trace_and_span_ids_uses_context(self): |
237 | 266 | assert trace_id == expected_trace |
238 | 267 | assert span_id == expected_span |
239 | 268 |
|
| 269 | + def test_get_trace_and_span_ids_uses_provider_before_otel(self): |
| 270 | + """Test that an external provider is checked before OTEL fallback.""" |
| 271 | + expected_trace = "c" * 32 |
| 272 | + expected_span = "d" * 16 |
| 273 | + |
| 274 | + set_trace_context_provider( |
| 275 | + lambda: { |
| 276 | + "trace_id": expected_trace, |
| 277 | + "span_id": expected_span, |
| 278 | + } |
| 279 | + ) |
| 280 | + |
| 281 | + trace_id, span_id = get_trace_and_span_ids() |
| 282 | + |
| 283 | + assert trace_id == expected_trace |
| 284 | + assert span_id == expected_span |
| 285 | + |
240 | 286 |
|
241 | 287 | class TestOtelAvailability: |
242 | 288 | """Tests for OpenTelemetry availability detection.""" |
|
0 commit comments