feat: Multiplexed sessions - Add test helper for multiplexed env vars.

Signed-off-by: Taylor Curran <taylor.curran@improving.com>

Author: currantw

tests/_helpers.py

@@ -1,7 +1,10 @@
 import unittest
+from os import getenv
+
 import mock
 
 from google.cloud.spanner_v1 import gapic_version
+from google.cloud.spanner_v1.session_options import TransactionType
 
 LIB_VERSION = gapic_version.__version__
 
@@ -31,6 +34,33 @@
 _TEST_OT_EXPORTER = None
 _TEST_OT_PROVIDER_INITIALIZED = False
 
+# Environment variables for enabling multiplexed sessions
+"GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS"
+ENV_VAR_ENABLE_MULTIPLEXED_FOR_PARTITIONED = (
+    "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS"
+)
+ENV_VAR_ENABLE_MULTIPLEXED_FOR_READ_WRITE = (
+    "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW"
+)
+
+
+def is_multiplexed_enabled(transaction_type: TransactionType) -> bool:
+    """Returns whether multiplexed sessions are enabled for the given transaction type."""
+
+    env_var = "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS"
+    env_var_partitioned = "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS"
+    env_var_read_write = "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_FOR_RW"
+
+    def _getenv(val: str) -> bool:
+        return getenv(val, "false").lower() == "true"
+
+    if transaction_type is TransactionType.READ_ONLY:
+        return _getenv(env_var)
+    elif transaction_type is TransactionType.PARTITIONED:
+        return _getenv(env_var) and _getenv(env_var_partitioned)
+    else:
+        return _getenv(env_var) and _getenv(env_var_read_write)
+
 
 def get_test_ot_exporter():
     global _TEST_OT_EXPORTER

tests/system/test_observability_options.py

@@ -16,12 +16,15 @@
 from mock import PropertyMock, patch
 
 from google.cloud.spanner_v1.session import Session
+from google.cloud.spanner_v1.session_options import TransactionType
 from . import _helpers
 from google.cloud.spanner_v1 import Client
 from google.api_core.exceptions import Aborted
 from google.auth.credentials import AnonymousCredentials
 from google.rpc import code_pb2
 
+from .._helpers import is_multiplexed_enabled
+
 HAS_OTEL_INSTALLED = False
 
 try:
@@ -113,11 +116,7 @@ def test_propagation(enable_extended_tracing):
         gotNames = [span.name for span in from_inject_spans]
 
         # Check if multiplexed sessions are enabled
-        import os
-
-        multiplexed_enabled = (
-            os.getenv("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS", "").lower() == "true"
-        )
+        multiplexed_enabled = is_multiplexed_enabled(TransactionType.READ_ONLY)
 
         # Determine expected session span name based on multiplexed sessions
         expected_session_span_name = (
@@ -213,13 +212,11 @@ def create_db_trace_exporter():
     reason="Tracing requires OpenTelemetry",
 )
 @patch.object(Session, "session_id", new_callable=PropertyMock)
-@patch.object(Session, "is_multiplexed", new_callable=PropertyMock)
-def test_transaction_abort_then_retry_spans(mock_session_multiplexed, mock_session_id):
+def test_transaction_abort_then_retry_spans(mock_session_id):
     from opentelemetry.trace.status import StatusCode
 
-    # Mock session properties for testing.
-    mock_session_multiplexed.return_value = session_multiplexed = False
     mock_session_id.return_value = session_id = "session-id"
+    multiplexed = is_multiplexed_enabled(TransactionType.READ_WRITE)
 
     db, trace_exporter = create_db_trace_exporter()
 
@@ -247,8 +244,8 @@ def select_in_txn(txn):
         ("Waiting for a session to become available", {"kind": "BurstyPool"}),
         ("No sessions available in pool. Creating session", {"kind": "BurstyPool"}),
         ("Creating Session", {}),
-        ("Using session", {"id": session_id, "multiplexed": session_multiplexed}),
-        ("Returning session", {"id": session_id, "multiplexed": session_multiplexed}),
+        ("Using session", {"id": session_id, "multiplexed": multiplexed}),
+        ("Returning session", {"id": session_id, "multiplexed": multiplexed}),
         (
             "Transaction was aborted in user operation, retrying",
             {"delay_seconds": "EPHEMERAL", "cause": "EPHEMERAL", "attempt": 1},
@@ -417,7 +414,6 @@ 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()
@@ -428,12 +424,9 @@ def test_database_partitioned_error():
         pass
 
     got_statuses, got_events = finished_spans_statuses(trace_exporter)
+    multiplexed_enabled = is_multiplexed_enabled(TransactionType.PARTITIONED)
 
-    multiplexed_partitioned_enabled = (
-        os.getenv("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS") == "true"
-    )
-
-    if multiplexed_partitioned_enabled:
+    if multiplexed_enabled:
         expected_event_names = [
             "Creating Session",
             "Using session",
@@ -496,7 +489,7 @@ def test_database_partitioned_error():
 
     expected_session_span_name = (
         "CloudSpanner.CreateMultiplexedSession"
-        if multiplexed_partitioned_enabled
+        if multiplexed_enabled
         else "CloudSpanner.CreateSession"
     )
     want_statuses = [

tests/system/test_session_api.py

@@ -29,6 +29,7 @@
 from google.cloud.spanner_admin_database_v1 import DatabaseDialect
 from google.cloud._helpers import UTC
 from google.cloud.spanner_v1.data_types import JsonObject
+from google.cloud.spanner_v1.session_options import TransactionType
 from .testdata import singer_pb2
 from tests import _helpers as ot_helpers
 from . import _helpers
@@ -37,7 +38,7 @@
     REQ_RAND_PROCESS_ID,
     parse_request_id,
 )
-
+from .._helpers import is_multiplexed_enabled
 
 SOME_DATE = datetime.date(2011, 1, 17)
 SOME_TIME = datetime.datetime(1989, 1, 17, 17, 59, 12, 345612)
@@ -430,8 +431,6 @@ def test_session_crud(sessions_database):
 
 
 def test_batch_insert_then_read(sessions_database, ot_exporter):
-    import os
-
     db_name = sessions_database.name
     sd = _sample_data
 
@@ -453,10 +452,7 @@ def test_batch_insert_then_read(sessions_database, ot_exporter):
         nth_req0 = sampling_req_id[-2]
 
         db = sessions_database
-
-        multiplexed_enabled = (
-            os.getenv("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS", "").lower() == "true"
-        )
+        multiplexed_enabled = is_multiplexed_enabled(TransactionType.READ_ONLY)
 
         # [A] Verify batch checkout spans
         # -------------------------------
@@ -690,12 +686,7 @@ def transaction_work(transaction):
     assert rows == []
 
     if ot_exporter is not None:
-        import os
-
-        multiplexed_enabled = (
-            os.getenv("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS", "").lower() == "true"
-        )
-
+        multiplexed_enabled = is_multiplexed_enabled(TransactionType.READ_ONLY)
         span_list = ot_exporter.get_finished_spans()
         got_span_names = [span.name for span in span_list]
 
@@ -3332,17 +3323,13 @@ def test_interval_array_cast(transaction):
 
 
 def test_session_id_and_multiplexed_flag_behavior(sessions_database, ot_exporter):
-    import os
-
     sd = _sample_data
 
     with sessions_database.batch() as batch:
         batch.delete(sd.TABLE, sd.ALL)
         batch.insert(sd.TABLE, sd.COLUMNS, sd.ROW_DATA)
 
-    multiplexed_enabled = (
-        os.getenv("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS", "").lower() == "true"
-    )
+    multiplexed_enabled = is_multiplexed_enabled(TransactionType.READ_ONLY)
 
     snapshot1_session_id = None
     snapshot2_session_id = None

tests/unit/test_database.py

@@ -38,6 +38,7 @@
 from google.cloud.spanner_v1.session import Session
 from google.cloud.spanner_v1.session_options import SessionOptions, TransactionType
 from tests._builders import build_spanner_api
+from tests._helpers import is_multiplexed_enabled
 
 DML_WO_PARAM = """
 DELETE FROM citizens
@@ -1527,7 +1528,6 @@ def test_session_factory_w_labels(self):
         self.assertEqual(session.labels, labels)
 
     def test_snapshot_defaults(self):
-        import os
         from google.cloud.spanner_v1.database import SnapshotCheckout
         from google.cloud.spanner_v1.snapshot import Snapshot
 
@@ -1539,9 +1539,7 @@ def test_snapshot_defaults(self):
         database = self._make_one(self.DATABASE_ID, instance, pool=pool)
 
         # Check if multiplexed sessions are enabled for read operations
-        multiplexed_enabled = (
-            os.getenv("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS") == "true"
-        )
+        multiplexed_enabled = is_multiplexed_enabled(TransactionType.READ_ONLY)
 
         if multiplexed_enabled:
             # When multiplexed sessions are enabled, configure the sessions manager
@@ -1575,7 +1573,6 @@ def test_snapshot_defaults(self):
 
     def test_snapshot_w_read_timestamp_and_multi_use(self):
         import datetime
-        import os
         from google.cloud._helpers import UTC
         from google.cloud.spanner_v1.database import SnapshotCheckout
         from google.cloud.spanner_v1.snapshot import Snapshot
@@ -1589,9 +1586,7 @@ def test_snapshot_w_read_timestamp_and_multi_use(self):
         database = self._make_one(self.DATABASE_ID, instance, pool=pool)
 
         # Check if multiplexed sessions are enabled for read operations
-        multiplexed_enabled = (
-            os.getenv("GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS") == "true"
-        )
+        multiplexed_enabled = is_multiplexed_enabled(TransactionType.READ_ONLY)
 
         if multiplexed_enabled:
             # When multiplexed sessions are enabled, configure the sessions manager