mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-08 15:48:29 +00:00
34a03f882c
Add crew context tracking using OpenTelemetry baggage for thread-safe propagation. Context is set during kickoff and cleaned up in finally block. Added thread safety tests with mocked agent execution.
89 lines
2.3 KiB
Python
89 lines
2.3 KiB
Python
import uuid
|
|
|
|
import pytest
|
|
from opentelemetry import baggage
|
|
from opentelemetry.context import attach, detach
|
|
|
|
from crewai.utilities.crew.crew_context import get_crew_context
|
|
from crewai.utilities.crew.models import CrewContext
|
|
|
|
|
|
def test_crew_context_creation():
|
|
crew_id = str(uuid.uuid4())
|
|
context = CrewContext(id=crew_id, key="test-crew")
|
|
assert context.id == crew_id
|
|
assert context.key == "test-crew"
|
|
|
|
|
|
def test_get_crew_context_with_baggage():
|
|
crew_id = str(uuid.uuid4())
|
|
assert get_crew_context() is None
|
|
|
|
crew_ctx = CrewContext(id=crew_id, key="test-key")
|
|
ctx = baggage.set_baggage("crew_context", crew_ctx)
|
|
token = attach(ctx)
|
|
|
|
try:
|
|
context = get_crew_context()
|
|
assert context is not None
|
|
assert context.id == crew_id
|
|
assert context.key == "test-key"
|
|
finally:
|
|
detach(token)
|
|
|
|
assert get_crew_context() is None
|
|
|
|
|
|
def test_get_crew_context_empty():
|
|
assert get_crew_context() is None
|
|
|
|
|
|
def test_baggage_nested_contexts():
|
|
crew_id1 = str(uuid.uuid4())
|
|
crew_id2 = str(uuid.uuid4())
|
|
|
|
crew_ctx1 = CrewContext(id=crew_id1, key="outer")
|
|
ctx1 = baggage.set_baggage("crew_context", crew_ctx1)
|
|
token1 = attach(ctx1)
|
|
|
|
try:
|
|
outer_context = get_crew_context()
|
|
assert outer_context.id == crew_id1
|
|
assert outer_context.key == "outer"
|
|
|
|
crew_ctx2 = CrewContext(id=crew_id2, key="inner")
|
|
ctx2 = baggage.set_baggage("crew_context", crew_ctx2)
|
|
token2 = attach(ctx2)
|
|
|
|
try:
|
|
inner_context = get_crew_context()
|
|
assert inner_context.id == crew_id2
|
|
assert inner_context.key == "inner"
|
|
finally:
|
|
detach(token2)
|
|
|
|
restored_context = get_crew_context()
|
|
assert restored_context.id == crew_id1
|
|
assert restored_context.key == "outer"
|
|
finally:
|
|
detach(token1)
|
|
|
|
assert get_crew_context() is None
|
|
|
|
|
|
def test_baggage_exception_handling():
|
|
crew_id = str(uuid.uuid4())
|
|
|
|
crew_ctx = CrewContext(id=crew_id, key="test")
|
|
ctx = baggage.set_baggage("crew_context", crew_ctx)
|
|
token = attach(ctx)
|
|
|
|
with pytest.raises(ValueError):
|
|
try:
|
|
assert get_crew_context() is not None
|
|
raise ValueError("Test exception")
|
|
finally:
|
|
detach(token)
|
|
|
|
assert get_crew_context() is None
|