3838import pytest
3939
4040import pymongo
41+ import pymongo ._otel as _otel
4142from bson import SON , json_util
4243from bson .codec_options import DEFAULT_CODEC_OPTIONS
4344from bson .objectid import ObjectId
9394 PLACEHOLDER_MAP ,
9495 EventListenerUtil ,
9596 MatchEvaluatorUtil ,
97+ _shared_test_provider ,
9698 coerce_result ,
9799 parse_bulk_write_error_result ,
98100 parse_bulk_write_result ,
113115
114116_IS_SYNC = False
115117
118+ _HAS_OTEL_TEST_DEPS = False
119+ if _otel ._HAS_OPENTELEMETRY :
120+ try :
121+ from opentelemetry .sdk .trace .export import SimpleSpanProcessor
122+ from opentelemetry .sdk .trace .export .in_memory_span_exporter import InMemorySpanExporter
123+
124+ _HAS_OTEL_TEST_DEPS = True
125+ except ImportError :
126+ pass
127+
116128IS_INTERRUPTED = False
117129
118130
@@ -230,6 +242,11 @@ def __init__(self, test_class):
230242 self ._entities : dict [str , Any ] = {}
231243 self ._listeners : dict [str , EventListenerUtil ] = {}
232244 self ._session_lsids : dict [str , Mapping [str , Any ]] = {}
245+ # The id of the (at most one, today) client entity created with
246+ # observeTracingMessages. Spans carry no attribute identifying which
247+ # client emitted them, so multi-client tracing correlation isn't
248+ # supported; _create_entity fails loudly if a second one appears.
249+ self ._tracing_client_id : Optional [str ] = None
233250 self .test : UnifiedSpecTestMixinV1 = test_class
234251
235252 def __contains__ (self , item ):
@@ -311,6 +328,25 @@ async def _create_entity(self, entity_spec, uri=None):
311328 )
312329 self ._listeners [spec ["id" ]] = listener
313330 kwargs ["event_listeners" ] = [listener ]
331+
332+ observe_tracing = spec .get ("observeTracingMessages" )
333+ if observe_tracing is not None :
334+ if self ._tracing_client_id is not None :
335+ self .test .fail (
336+ "Multiple clients with observeTracingMessages are not supported "
337+ f"by the unified test format runner (already tracking "
338+ f"{ self ._tracing_client_id !r} , got { spec ['id' ]!r} )"
339+ )
340+ self ._tracing_client_id = spec ["id" ]
341+ enable_payload = observe_tracing .get ("enableCommandPayload" , False )
342+ kwargs ["tracing" ] = {
343+ "enabled" : True ,
344+ # Tests asserting db.query.text match the full, untruncated
345+ # command, so an effectively-unlimited length avoids
346+ # truncating and failing that assertion.
347+ "query_text_max_length" : 1_000_000 if enable_payload else None ,
348+ }
349+
314350 if spec .get ("useMultipleMongoses" ):
315351 if async_client_context .load_balancer :
316352 kwargs ["h" ] = async_client_context .MULTI_MONGOS_LB_URI
@@ -482,6 +518,8 @@ class UnifiedSpecTestMixinV1(AsyncIntegrationTest):
482518 TEST_SPEC : Any
483519 TEST_PATH = "" # This gets filled in by generate_test_classes
484520 mongos_clients : list [AsyncMongoClient ] = []
521+ # Set in setUpClass, only for test files that use observeTracingMessages.
522+ _tracing_exporter : Optional [Any ] = None
485523
486524 @staticmethod
487525 async def should_run_on (run_on_spec ):
@@ -526,6 +564,23 @@ async def insert_initial_data(self, initial_data):
526564
527565 @classmethod
528566 def setUpClass (cls ) -> None :
567+ # Only register a span exporter (and the shared SDK TracerProvider it
568+ # depends on) for test files that actually use observeTracingMessages,
569+ # to avoid needlessly accumulating span processors on the process-wide
570+ # provider for the (vast majority of) unified-format suites that don't.
571+ cls ._tracing_exporter = None
572+ uses_tracing = any (
573+ "observeTracingMessages" in entity .get ("client" , {})
574+ for entity in cls .TEST_SPEC .get ("createEntities" , [])
575+ )
576+ if uses_tracing :
577+ if not _HAS_OTEL_TEST_DEPS :
578+ raise unittest .SkipTest (
579+ "observeTracingMessages requires opentelemetry-sdk to be installed"
580+ )
581+ cls ._tracing_exporter = InMemorySpanExporter ()
582+ _shared_test_provider ().add_span_processor (SimpleSpanProcessor (cls ._tracing_exporter ))
583+
529584 # Speed up the tests by decreasing the heartbeat frequency.
530585 cls .knobs = client_knobs (
531586 heartbeat_frequency = 0.1 ,
@@ -538,6 +593,14 @@ def setUpClass(cls) -> None:
538593 @classmethod
539594 def tearDownClass (cls ) -> None :
540595 cls .knobs .disable ()
596+ # The exporter's span processor can never be removed from the shared process-wide
597+ # TracerProvider (see _shared_test_provider), so without this, every span emitted by any
598+ # client anywhere in the process for the rest of the test run keeps getting appended to this
599+ # (otherwise dead) class's exporter: an unbounded memory leak across a full test run, and
600+ # needless per-span export overhead for every other tracing-enabled test class that runs
601+ # afterwards. shutdown() makes further export() calls into this exporter no-ops.
602+ if cls ._tracing_exporter is not None :
603+ cls ._tracing_exporter .shutdown ()
541604
542605 async def asyncSetUp (self ):
543606 # super call creates internal client cls.client
@@ -576,6 +639,14 @@ def maybe_skip_test(self, spec):
576639 self .skipTest ("PyMongo does not support the symbol type" )
577640 if "timeoutms applied to entire download" in description :
578641 self .skipTest ("PyMongo's open_download_stream does not cap the stream's lifetime" )
642+ # Removed API: PyMongo no longer exposes map_reduce/inline_map_reduce at
643+ # all (mapReduce is deprecated server-side), so there's no code path left
644+ # that could send this command; this operation can never be exercised.
645+ if class_name == "testoperationmapreduce" and description == "mapreduce" :
646+ self .skipTest (
647+ "PyMongo removed the map_reduce/inline_map_reduce Collection methods "
648+ "(mapReduce is deprecated server-side); this operation cannot be exercised"
649+ )
579650 if any (
580651 x in description
581652 for x in [
@@ -1463,6 +1534,77 @@ def format_logs(log_list):
14631534 self .match_evaluator .match_result (expected_data , actual_data )
14641535 self .match_evaluator .match_result (expected_msg , actual_msg )
14651536
1537+ async def check_tracing_messages (self , operations , spec ):
1538+ # Like expectLogMessages/expectEvents, expectTracingMessages is a list of
1539+ # per-client blocks (even though only one client with
1540+ # observeTracingMessages is currently supported, see entity.py above).
1541+ exporter = self ._tracing_exporter
1542+ if exporter is None :
1543+ self .fail (
1544+ "expectTracingMessages requires a client entity created with observeTracingMessages"
1545+ )
1546+
1547+ exporter .clear ()
1548+ await self .run_operations (operations )
1549+ finished_spans = exporter .get_finished_spans ()
1550+
1551+ # Reconstruct the parent/child span tree from the flat, finish-ordered
1552+ # list the in-memory exporter records, keyed by each span's parent id.
1553+ children_by_parent_id = defaultdict (list )
1554+ for span in finished_spans :
1555+ parent_id = span .parent .span_id if span .parent is not None else None
1556+ children_by_parent_id [parent_id ].append (span )
1557+
1558+ def check_span_list (expected_list , actual_list , ignore_extra_spans ):
1559+ if ignore_extra_spans :
1560+ # Per the unified-test-format spec, "additional unexpected spans
1561+ # are allowed". Unlike ignoreExtraEvents (which only tolerates
1562+ # a trailing tail), spans from concurrent/out-of-band activity
1563+ # (e.g. a testRunner-issued configureFailPoint command) can
1564+ # finish interleaved anywhere among the expected ones, not just
1565+ # at the end. Filter down to just the spans that line up (by
1566+ # name, in order) with the expected list, dropping anything
1567+ # else, instead of naively truncating the tail.
1568+ filtered = []
1569+ expected_iter = iter (expected_list )
1570+ current_expected = next (expected_iter , None )
1571+ for actual in actual_list :
1572+ if current_expected is not None and actual .name == current_expected ["name" ]:
1573+ filtered .append (actual )
1574+ current_expected = next (expected_iter , None )
1575+ actual_list = filtered
1576+ self .assertEqual (
1577+ len (expected_list ),
1578+ len (actual_list ),
1579+ f"expected spans { [e ['name' ] for e in expected_list ]} but got "
1580+ f"{ [a .name for a in actual_list ]} " ,
1581+ )
1582+ for expected , actual in zip (expected_list , actual_list ):
1583+ self .assertEqual (expected ["name" ], actual .name )
1584+ self .match_evaluator .match_span_attributes (
1585+ expected ["attributes" ], actual .attributes
1586+ )
1587+ expected_nested = expected .get ("nested" )
1588+ if expected_nested is not None :
1589+ actual_children = children_by_parent_id [actual .context .span_id ]
1590+ check_span_list (expected_nested , actual_children , ignore_extra_spans )
1591+
1592+ for client_spec in spec :
1593+ expected_client_id = client_spec ["client" ]
1594+ tracing_client_id = self .entity_map ._tracing_client_id
1595+ self .assertEqual (
1596+ expected_client_id ,
1597+ tracing_client_id ,
1598+ f"expectTracingMessages.client { expected_client_id !r} does not match the "
1599+ f"client with observeTracingMessages enabled ({ tracing_client_id !r} )" ,
1600+ )
1601+
1602+ ignore_extra_spans = client_spec .get ("ignoreExtraSpans" , False )
1603+ expected_spans = client_spec ["spans" ]
1604+ self .assertTrue (expected_spans , "expectTracingMessages spans must be non-empty" )
1605+
1606+ check_span_list (expected_spans , children_by_parent_id [None ], ignore_extra_spans )
1607+
14661608 async def verify_outcome (self , spec ):
14671609 for collection_data in spec :
14681610 coll_name = collection_data ["collectionName" ]
@@ -1551,6 +1693,10 @@ async def _run_scenario(self, spec, uri=None):
15511693 expect_log_messages = spec ["expectLogMessages" ]
15521694 self .assertTrue (expect_log_messages , "expectEvents must be non-empty" )
15531695 await self .check_log_messages (spec ["operations" ], expect_log_messages )
1696+ elif "expectTracingMessages" in spec :
1697+ expect_tracing_messages = spec ["expectTracingMessages" ]
1698+ self .assertTrue (expect_tracing_messages , "expectTracingMessages must be non-empty" )
1699+ await self .check_tracing_messages (spec ["operations" ], expect_tracing_messages )
15541700 else :
15551701 # process operations
15561702 await self .run_operations (spec ["operations" ])
0 commit comments