Skip to content

Commit 784249a

Browse files
hassiebphassiebp
committed
push
1 parent da81030 commit 784249a

1 file changed

Lines changed: 243 additions & 0 deletions

File tree

‎tests/test_propagate_attributes.py‎

Lines changed: 243 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1150,3 +1150,246 @@ def test_attributes_persist_across_tracer_changes(
11501150
LangfuseOtelSpanAttributes.TRACE_USER_ID,
11511151
"persistent_user",
11521152
)
1153+
1154+
1155+
class TestPropagateAttributesBaggage(TestPropagateAttributesBase):
1156+
"""Tests for as_baggage=True parameter and OpenTelemetry baggage propagation."""
1157+
1158+
def test_baggage_is_set_when_as_baggage_true(self, langfuse_client):
1159+
"""Verify baggage entries are created with correct keys when as_baggage=True."""
1160+
from opentelemetry import baggage
1161+
from opentelemetry import context as otel_context
1162+
1163+
with langfuse_client.start_as_current_span(name="parent"):
1164+
with propagate_attributes(
1165+
user_id="user_123",
1166+
session_id="session_abc",
1167+
metadata={"env": "test", "version": "2.0"},
1168+
as_baggage=True,
1169+
):
1170+
# Get current context and inspect baggage
1171+
current_context = otel_context.get_current()
1172+
baggage_entries = baggage.get_all(context=current_context)
1173+
1174+
# Verify baggage entries exist with correct keys
1175+
assert "langfuse_user_id" in baggage_entries
1176+
assert baggage_entries["langfuse_user_id"] == "user_123"
1177+
1178+
assert "langfuse_session_id" in baggage_entries
1179+
assert baggage_entries["langfuse_session_id"] == "session_abc"
1180+
1181+
assert "langfuse_metadata_env" in baggage_entries
1182+
assert baggage_entries["langfuse_metadata_env"] == "test"
1183+
1184+
assert "langfuse_metadata_version" in baggage_entries
1185+
assert baggage_entries["langfuse_metadata_version"] == "2.0"
1186+
1187+
def test_spans_receive_attributes_from_baggage(
1188+
self, langfuse_client, memory_exporter
1189+
):
1190+
"""Verify child spans get attributes when parent uses as_baggage=True."""
1191+
with langfuse_client.start_as_current_span(name="parent"):
1192+
with propagate_attributes(
1193+
user_id="baggage_user",
1194+
session_id="baggage_session",
1195+
metadata={"source": "baggage"},
1196+
as_baggage=True,
1197+
):
1198+
# Create child span
1199+
child = langfuse_client.start_span(name="child-span")
1200+
child.end()
1201+
1202+
# Verify child span has all attributes
1203+
child_span = self.get_span_by_name(memory_exporter, "child-span")
1204+
self.verify_span_attribute(
1205+
child_span, LangfuseOtelSpanAttributes.TRACE_USER_ID, "baggage_user"
1206+
)
1207+
self.verify_span_attribute(
1208+
child_span,
1209+
LangfuseOtelSpanAttributes.TRACE_SESSION_ID,
1210+
"baggage_session",
1211+
)
1212+
self.verify_span_attribute(
1213+
child_span,
1214+
f"{LangfuseOtelSpanAttributes.TRACE_METADATA}.source",
1215+
"baggage",
1216+
)
1217+
1218+
def test_baggage_disabled_by_default(self, langfuse_client):
1219+
"""Verify as_baggage=False (default) doesn't create baggage entries."""
1220+
from opentelemetry import baggage
1221+
from opentelemetry import context as otel_context
1222+
1223+
with langfuse_client.start_as_current_span(name="parent"):
1224+
with propagate_attributes(
1225+
user_id="user_123",
1226+
session_id="session_abc",
1227+
):
1228+
# Get current context and inspect baggage
1229+
current_context = otel_context.get_current()
1230+
baggage_entries = baggage.get_all(context=current_context)
1231+
assert len(baggage_entries) == 0
1232+
1233+
def test_metadata_key_with_user_id_substring_doesnt_collide(
1234+
self, langfuse_client, memory_exporter
1235+
):
1236+
"""Verify metadata key containing 'user_id' substring doesn't map to TRACE_USER_ID."""
1237+
with langfuse_client.start_as_current_span(name="parent"):
1238+
with propagate_attributes(
1239+
metadata={"user_info": "some_data", "user_id_copy": "another"},
1240+
as_baggage=True,
1241+
):
1242+
child = langfuse_client.start_span(name="child-span")
1243+
child.end()
1244+
1245+
child_span = self.get_span_by_name(memory_exporter, "child-span")
1246+
1247+
# Should NOT have TRACE_USER_ID attribute
1248+
self.verify_missing_attribute(
1249+
child_span, LangfuseOtelSpanAttributes.TRACE_USER_ID
1250+
)
1251+
1252+
# Should have metadata attributes with correct keys
1253+
self.verify_span_attribute(
1254+
child_span,
1255+
f"{LangfuseOtelSpanAttributes.TRACE_METADATA}.user_info",
1256+
"some_data",
1257+
)
1258+
self.verify_span_attribute(
1259+
child_span,
1260+
f"{LangfuseOtelSpanAttributes.TRACE_METADATA}.user_id_copy",
1261+
"another",
1262+
)
1263+
1264+
def test_metadata_key_with_session_substring_doesnt_collide(
1265+
self, langfuse_client, memory_exporter
1266+
):
1267+
"""Verify metadata key containing 'session_id' substring doesn't map to TRACE_SESSION_ID."""
1268+
with langfuse_client.start_as_current_span(name="parent"):
1269+
with propagate_attributes(
1270+
metadata={"session_data": "value1", "session_id_backup": "value2"},
1271+
as_baggage=True,
1272+
):
1273+
child = langfuse_client.start_span(name="child-span")
1274+
child.end()
1275+
1276+
child_span = self.get_span_by_name(memory_exporter, "child-span")
1277+
1278+
# Should NOT have TRACE_SESSION_ID attribute
1279+
self.verify_missing_attribute(
1280+
child_span, LangfuseOtelSpanAttributes.TRACE_SESSION_ID
1281+
)
1282+
1283+
# Should have metadata attributes with correct keys
1284+
self.verify_span_attribute(
1285+
child_span,
1286+
f"{LangfuseOtelSpanAttributes.TRACE_METADATA}.session_data",
1287+
"value1",
1288+
)
1289+
self.verify_span_attribute(
1290+
child_span,
1291+
f"{LangfuseOtelSpanAttributes.TRACE_METADATA}.session_id_backup",
1292+
"value2",
1293+
)
1294+
1295+
def test_metadata_keys_extract_correctly_from_baggage(
1296+
self, langfuse_client, memory_exporter
1297+
):
1298+
"""Verify metadata keys are correctly formatted in baggage and extracted back."""
1299+
with langfuse_client.start_as_current_span(name="parent"):
1300+
with propagate_attributes(
1301+
metadata={
1302+
"env": "production",
1303+
"region": "us-west",
1304+
"experiment_id": "exp_123",
1305+
},
1306+
as_baggage=True,
1307+
):
1308+
child = langfuse_client.start_span(name="child-span")
1309+
child.end()
1310+
1311+
child_span = self.get_span_by_name(memory_exporter, "child-span")
1312+
1313+
# All metadata should be under TRACE_METADATA prefix
1314+
self.verify_span_attribute(
1315+
child_span,
1316+
f"{LangfuseOtelSpanAttributes.TRACE_METADATA}.env",
1317+
"production",
1318+
)
1319+
self.verify_span_attribute(
1320+
child_span,
1321+
f"{LangfuseOtelSpanAttributes.TRACE_METADATA}.region",
1322+
"us-west",
1323+
)
1324+
self.verify_span_attribute(
1325+
child_span,
1326+
f"{LangfuseOtelSpanAttributes.TRACE_METADATA}.experiment_id",
1327+
"exp_123",
1328+
)
1329+
1330+
def test_baggage_and_context_both_propagate(self, langfuse_client, memory_exporter):
1331+
"""Verify attributes propagate when both baggage and context mechanisms are active."""
1332+
with langfuse_client.start_as_current_span(name="parent"):
1333+
# Enable baggage
1334+
with propagate_attributes(
1335+
user_id="user_both",
1336+
session_id="session_both",
1337+
metadata={"source": "both"},
1338+
as_baggage=True,
1339+
):
1340+
# Create multiple levels of nesting
1341+
with langfuse_client.start_as_current_span(name="middle"):
1342+
child = langfuse_client.start_span(name="leaf")
1343+
child.end()
1344+
1345+
# Verify all spans have attributes
1346+
for span_name in ["parent", "middle", "leaf"]:
1347+
span_data = self.get_span_by_name(memory_exporter, span_name)
1348+
self.verify_span_attribute(
1349+
span_data, LangfuseOtelSpanAttributes.TRACE_USER_ID, "user_both"
1350+
)
1351+
self.verify_span_attribute(
1352+
span_data, LangfuseOtelSpanAttributes.TRACE_SESSION_ID, "session_both"
1353+
)
1354+
self.verify_span_attribute(
1355+
span_data,
1356+
f"{LangfuseOtelSpanAttributes.TRACE_METADATA}.source",
1357+
"both",
1358+
)
1359+
1360+
def test_baggage_survives_context_isolation(self, langfuse_client, memory_exporter):
1361+
"""Simulate cross-process scenario: baggage persists when context is detached/reattached."""
1362+
from opentelemetry import context as otel_context
1363+
1364+
# Step 1: Create context with baggage
1365+
with langfuse_client.start_as_current_span(name="original-process"):
1366+
with propagate_attributes(
1367+
user_id="cross_process_user",
1368+
session_id="cross_process_session",
1369+
as_baggage=True,
1370+
):
1371+
# Capture the context with baggage
1372+
context_with_baggage = otel_context.get_current()
1373+
1374+
# Step 2: Simulate "remote" process by creating span in saved context
1375+
# This mimics what happens when receiving an HTTP request with baggage headers
1376+
token = otel_context.attach(context_with_baggage)
1377+
try:
1378+
with langfuse_client.start_as_current_span(name="remote-process"):
1379+
child = langfuse_client.start_span(name="remote-child")
1380+
child.end()
1381+
finally:
1382+
otel_context.detach(token)
1383+
1384+
# Verify remote spans have the propagated attributes from baggage
1385+
remote_child = self.get_span_by_name(memory_exporter, "remote-child")
1386+
self.verify_span_attribute(
1387+
remote_child,
1388+
LangfuseOtelSpanAttributes.TRACE_USER_ID,
1389+
"cross_process_user",
1390+
)
1391+
self.verify_span_attribute(
1392+
remote_child,
1393+
LangfuseOtelSpanAttributes.TRACE_SESSION_ID,
1394+
"cross_process_session",
1395+
)

0 commit comments

Comments
 (0)