chore(spanner): fix failing OpenTelemetrySpanTest

googleapis · Dec 24, 2024 · 7211d2f · 7211d2f
1 parent 138c58f
commit 7211d2f
Showing 1 changed file with 43 additions and 2 deletions.
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetrySpanTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/OpenTelemetrySpanTest.java
@@ -147,7 +147,7 @@ public class OpenTelemetrySpanTest {
           "Commit Done",
           "Transaction Attempt Succeeded");
 
-  private int expectedReadWriteTransactionCount = 7;
+  private int expectedReadWriteTransactionEventsCount = 7;
   private List<String> expectedReadWriteTransactionErrorWithBeginTransactionEvents =
       ImmutableList.of(
           "Acquiring session",
@@ -238,6 +238,7 @@ public void setUp() throws Exception {
                 SessionPoolOptions.newBuilder()
                     .setMinSessions(2)
                     .setWaitForMinSessionsDuration(Duration.ofSeconds(10))
+                    .setSkipVerifyingBeginTransactionForMuxRW(true)
                     .build());
 
     spanner = builder.build().getService();
@@ -450,6 +451,17 @@ public void transactionRunner() {
                 "CloudSpannerOperation.Commit",
                 "CloudSpannerOperation.BatchCreateSessions",
                 "CloudSpanner.ReadWriteTransaction");
+
+    if (isMultiplexedSessionsEnabledForRW()) {
+      expectedReadWriteTransactionEvents =
+          ImmutableList.of(
+              "Starting Transaction Attempt",
+              "Starting Commit",
+              "Commit Done",
+              "Transaction Attempt Succeeded");
+      expectedReadWriteTransactionEventsCount = 4;
+    }
+
     DatabaseClient client = getClient();
     TransactionRunner runner = client.readWriteTransaction();
     runner.run(transaction -> transaction.executeUpdate(UPDATE_STATEMENT));
@@ -494,7 +506,7 @@ public void transactionRunner() {
                   verifyRequestEvents(
                       spanItem,
                       expectedReadWriteTransactionEvents,
-                      expectedReadWriteTransactionCount);
+                      expectedReadWriteTransactionEventsCount);
                   break;
                 default:
                   assert false;
@@ -519,6 +531,16 @@ public void transactionRunnerWithError() {
                 "CloudSpannerOperation.BatchCreateSessions",
                 "CloudSpannerOperation.ExecuteUpdate",
                 "CloudSpanner.ReadWriteTransaction");
+
+    if (isMultiplexedSessionsEnabledForRW()) {
+      expectedReadWriteTransactionErrorEvents =
+          ImmutableList.of(
+              "Starting Transaction Attempt",
+              "Transaction Attempt Failed in user operation",
+              "exception");
+      expectedReadWriteTransactionErrorEventsCount = 3;
+    }
+
     DatabaseClient client = getClient();
     TransactionRunner runner = client.readWriteTransaction();
     SpannerException e =
@@ -579,6 +601,18 @@ public void transactionRunnerWithFailedAndBeginTransaction() {
             "CloudSpannerOperation.Commit",
             "CloudSpannerOperation.BatchCreateSessions",
             "CloudSpanner.ReadWriteTransaction");
+    if (isMultiplexedSessionsEnabledForRW()) {
+      expectedReadWriteTransactionErrorWithBeginTransactionEvents =
+          ImmutableList.of(
+              "Starting Transaction Attempt",
+              "Transaction Attempt Aborted in user operation. Retrying",
+              "Creating Transaction",
+              "Transaction Creation Done",
+              "Starting Commit",
+              "Commit Done",
+              "Transaction Attempt Succeeded");
+      expectedReadWriteTransactionErrorWithBeginTransactionEventsCount = 8;
+    }
     DatabaseClient client = getClient();
     assertEquals(
         Long.valueOf(1L),
@@ -837,4 +871,11 @@ private boolean isMultiplexedSessionsEnabled() {
     }
     return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSession();
   }
+
+  private boolean isMultiplexedSessionsEnabledForRW() {
+    if (spanner.getOptions() == null || spanner.getOptions().getSessionPoolOptions() == null) {
+      return false;
+    }
+    return spanner.getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW();
+  }
 }