mirror of
https://github.com/langgenius/dify.git
synced 2026-07-25 21:48:30 +08:00
52 lines
1.6 KiB
Python
52 lines
1.6 KiB
Python
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
import pytest
|
|
from opentelemetry import context as otel_context
|
|
|
|
from extensions.otel.context import propagate_context
|
|
|
|
|
|
def test_propagate_context_captures_context_when_wrapped() -> None:
|
|
context_key = otel_context.create_key("test-context")
|
|
captured_context = otel_context.set_value(context_key, "captured")
|
|
|
|
token = otel_context.attach(captured_context)
|
|
try:
|
|
wrapped = propagate_context(lambda: otel_context.get_value(context_key))
|
|
finally:
|
|
otel_context.detach(token)
|
|
|
|
with ThreadPoolExecutor(max_workers=1) as executor:
|
|
assert executor.submit(wrapped).result() == "captured"
|
|
|
|
|
|
def test_propagate_context_detaches_context_after_exception() -> None:
|
|
context_key = otel_context.create_key("test-context")
|
|
captured_context = otel_context.set_value(context_key, "captured")
|
|
|
|
def raise_error() -> None:
|
|
raise RuntimeError("retrieval failed")
|
|
|
|
token = otel_context.attach(captured_context)
|
|
try:
|
|
wrapped = propagate_context(raise_error)
|
|
finally:
|
|
otel_context.detach(token)
|
|
|
|
def invoke_and_read_context() -> str | None:
|
|
with pytest.raises(RuntimeError, match="retrieval failed"):
|
|
wrapped()
|
|
return otel_context.get_value(context_key)
|
|
|
|
with ThreadPoolExecutor(max_workers=1) as executor:
|
|
assert executor.submit(invoke_and_read_context).result() is None
|
|
|
|
|
|
def test_propagate_context_preserves_function_metadata() -> None:
|
|
def retrieve() -> None:
|
|
pass
|
|
|
|
wrapped = propagate_context(retrieve)
|
|
|
|
assert wrapped.__name__ == "retrieve"
|