fix tests

Author: rahul2393

.github/workflows/integration-tests-against-emulator-with-multiplexed-session.yaml

@@ -31,3 +31,4 @@ jobs:
           GOOGLE_CLOUD_PROJECT: emulator-test-project
           GOOGLE_CLOUD_TESTS_CREATE_SPANNER_INSTANCE: true
           GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS: true
+          GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS: true

google/cloud/spanner_v1/session.py

@@ -174,8 +174,9 @@ def create(self):
             request.session.labels = self._labels
 
         observability_options = getattr(self._database, "observability_options", None)
+        span_name = "CloudSpanner.CreateMultiplexedSession" if self._is_multiplexed else "CloudSpanner.CreateSession"
         with trace_call(
-            "CloudSpanner.CreateSession",
+            span_name,
             self,
             self._labels,
             observability_options=observability_options,

tests/system/test_observability_options.py

@@ -392,6 +392,7 @@ def tx_update(txn):
     reason="Tracing requires OpenTelemetry",
 )
 def test_database_partitioned_error():
+    import os
     from opentelemetry.trace.status import StatusCode
 
     db, trace_exporter = create_db_trace_exporter()
@@ -402,43 +403,101 @@ def test_database_partitioned_error():
         pass
 
     got_statuses, got_events = finished_spans_statuses(trace_exporter)
-    # Check for the series of events
-    want_events = [
-        ("Acquiring session", {"kind": "BurstyPool"}),
-        ("Waiting for a session to become available", {"kind": "BurstyPool"}),
-        ("No sessions available in pool. Creating session", {"kind": "BurstyPool"}),
-        ("Creating Session", {}),
-        ("Starting BeginTransaction", {}),
-        (
+
+    # Check if multiplexed sessions are enabled for partitioned operations
+    multiplexed_partitioned_enabled = (
+        os.getenv("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS") == "true"
+    )
+
+    # Define expected events based on whether multiplexed sessions are enabled
+    if multiplexed_partitioned_enabled:
+        # When multiplexed sessions are enabled for partitioned operations,
+        # the execution path is different - sessions manager creates multiplexed sessions directly
+        expected_event_names = [
+            "Creating Session",
+            "Using session",
+            "Starting BeginTransaction",
+            "Returning session",
             "exception",
-            {
-                "exception.type": "google.api_core.exceptions.InvalidArgument",
-                "exception.message": "400 Table not found: NonExistent [at 1:8]\nUPDATE NonExistent SET name = 'foo' WHERE id > 1\n       ^",
-                "exception.stacktrace": "EPHEMERAL",
-                "exception.escaped": "False",
-            },
-        ),
-        (
             "exception",
-            {
-                "exception.type": "google.api_core.exceptions.InvalidArgument",
-                "exception.message": "400 Table not found: NonExistent [at 1:8]\nUPDATE NonExistent SET name = 'foo' WHERE id > 1\n       ^",
-                "exception.stacktrace": "EPHEMERAL",
-                "exception.escaped": "False",
-            },
-        ),
-    ]
-    assert got_events == want_events
+        ]
+        # Check that we have the expected events
+        assert len(got_events) == len(expected_event_names)
+        for i, expected_name in enumerate(expected_event_names):
+            assert got_events[i][0] == expected_name
+
+        # Verify session usage event shows multiplexed session
+        assert got_events[1][1]["multiplexed"] is True
+
+        # Verify session return event shows multiplexed session
+        assert got_events[3][1]["multiplexed"] is True
+
+        # Verify the exception details
+        for i in [4, 5]:  # Both exception events
+            assert (
+                got_events[i][1]["exception.type"]
+                == "google.api_core.exceptions.InvalidArgument"
+            )
+            assert (
+                "Table not found: NonExistent" in got_events[i][1]["exception.message"]
+            )
+    else:
+        # When multiplexed sessions are disabled, sessions manager still manages sessions
+        # but uses regular pool sessions instead of multiplexed sessions
+        expected_event_names = [
+            "Acquiring session",
+            "Waiting for a session to become available",
+            "No sessions available in pool. Creating session",
+            "Creating Session",
+            "Using session",
+            "Starting BeginTransaction",
+            "Returning session",
+            "exception",
+            "exception",
+        ]
 
-    # Check for the statues.
+        # Check that we have the expected events
+        assert len(got_events) == len(expected_event_names)
+        for i, expected_name in enumerate(expected_event_names):
+            assert got_events[i][0] == expected_name
+
+        # Verify pool-related events
+        assert got_events[0][1]["kind"] == "BurstyPool"
+        assert got_events[1][1]["kind"] == "BurstyPool"
+        assert got_events[2][1]["kind"] == "BurstyPool"
+
+        # Verify session usage event shows non-multiplexed session
+        assert got_events[4][1]["multiplexed"] is False
+
+        # Verify session return event shows non-multiplexed session
+        assert got_events[6][1]["multiplexed"] is False
+
+        # Verify the exception details
+        for i in [7, 8]:  # Both exception events
+            assert (
+                got_events[i][1]["exception.type"]
+                == "google.api_core.exceptions.InvalidArgument"
+            )
+            assert (
+                "Table not found: NonExistent" in got_events[i][1]["exception.message"]
+            )
+
+        # Check for the statuses.
     codes = StatusCode
+
+    # Determine expected session creation span name based on multiplexed sessions
+    expected_session_span_name = (
+        "CloudSpanner.CreateMultiplexedSession"
+        if multiplexed_partitioned_enabled
+        else "CloudSpanner.CreateSession"
+    )
     want_statuses = [
         (
             "CloudSpanner.Database.execute_partitioned_pdml",
             codes.ERROR,
             "InvalidArgument: 400 Table not found: NonExistent [at 1:8]\nUPDATE NonExistent SET name = 'foo' WHERE id > 1\n       ^",
         ),
-        ("CloudSpanner.CreateSession", codes.OK, None),
+        (expected_session_span_name, codes.OK, None),
         (
             "CloudSpanner.ExecuteStreamingSql",
             codes.ERROR,