chore: import upstream snapshot with attribution

tests/tracing/test_setup.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import atexit
+from typing import Any, cast
+
+import pytest
+
+from agents.tracing import (
+    processors as tracing_processors,
+    provider as tracing_provider,
+    setup as tracing_setup,
+)
+
+
+class _DummyProvider:
+    def __init__(self) -> None:
+        self.shutdown_calls = 0
+
+    def shutdown(self) -> None:
+        self.shutdown_calls += 1
+
+
+class _DefaultProviderWithTimeout(tracing_provider.DefaultTraceProvider):
+    def __init__(self) -> None:
+        super().__init__()
+        self.shutdown_timeout: float | None = None
+
+    def shutdown(self, timeout: float | None = None) -> None:
+        self.shutdown_timeout = timeout
+
+
+class _BootstrapProvider:
+    def __init__(self) -> None:
+        self.processors: list[Any] = []
+        self.shutdown_calls = 0
+
+    def register_processor(self, processor: Any) -> None:
+        self.processors.append(processor)
+
+    def shutdown(self) -> None:
+        self.shutdown_calls += 1
+
+
+def test_shutdown_global_trace_provider_calls_shutdown(monkeypatch: pytest.MonkeyPatch) -> None:
+    provider = _DummyProvider()
+    monkeypatch.setattr(tracing_setup, "GLOBAL_TRACE_PROVIDER", provider)
+
+    tracing_setup._shutdown_global_trace_provider()
+
+    assert provider.shutdown_calls == 1
+
+
+def test_shutdown_global_trace_provider_passes_timeout_to_default_provider(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    provider = _DefaultProviderWithTimeout()
+    monkeypatch.setattr(tracing_setup, "GLOBAL_TRACE_PROVIDER", provider)
+
+    tracing_setup._shutdown_global_trace_provider()
+
+    assert provider.shutdown_timeout == tracing_setup._DEFAULT_SHUTDOWN_TIMEOUT
+
+
+def test_set_trace_provider_registers_shutdown_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    registrations: list[Any] = []
+
+    def fake_register(callback: Any) -> Any:
+        registrations.append(callback)
+        return callback
+
+    first = _DummyProvider()
+    second = _DummyProvider()
+
+    monkeypatch.setattr(atexit, "register", fake_register)
+    monkeypatch.setattr(tracing_setup, "GLOBAL_TRACE_PROVIDER", None)
+    monkeypatch.setattr(tracing_setup, "_SHUTDOWN_HANDLER_REGISTERED", False)
+
+    tracing_setup.set_trace_provider(cast(Any, first))
+    tracing_setup.set_trace_provider(cast(Any, second))
+
+    assert cast(Any, tracing_setup.GLOBAL_TRACE_PROVIDER) is second
+    assert registrations == [tracing_setup._shutdown_global_trace_provider]
+
+
+def test_get_trace_provider_returns_existing_provider(monkeypatch: pytest.MonkeyPatch) -> None:
+    provider = _DummyProvider()
+
+    def fail_register(_: Any) -> None:
+        raise AssertionError("atexit.register should not be called for an existing provider.")
+
+    monkeypatch.setattr(atexit, "register", fail_register)
+    monkeypatch.setattr(tracing_setup, "GLOBAL_TRACE_PROVIDER", provider)
+
+    assert cast(Any, tracing_setup.get_trace_provider()) is provider
+
+
+def test_get_trace_provider_bootstraps_provider_in_process(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    registrations: list[Any] = []
+    default_processor = object()
+
+    def fake_register(callback: Any) -> Any:
+        registrations.append(callback)
+        return callback
+
+    monkeypatch.setattr(atexit, "register", fake_register)
+    monkeypatch.setattr(tracing_setup, "GLOBAL_TRACE_PROVIDER", None)
+    monkeypatch.setattr(tracing_setup, "_SHUTDOWN_HANDLER_REGISTERED", False)
+    monkeypatch.setattr(tracing_processors, "default_processor", lambda: default_processor)
+    monkeypatch.setattr(tracing_provider, "DefaultTraceProvider", _BootstrapProvider)
+
+    provider = tracing_setup.get_trace_provider()
+
+    assert isinstance(provider, _BootstrapProvider)
+    assert provider.processors == [default_processor]
+    assert tracing_setup.GLOBAL_TRACE_PROVIDER is provider
+    assert registrations == [tracing_setup._shutdown_global_trace_provider]