chore(spanner): fix SpanTest unit test

Author: harshachinta

google-cloud-spanner/src/test/java/com/google/cloud/spanner/SpanTest.java

@@ -480,7 +480,23 @@ public void transactionRunner() {
             "Request for 2 sessions returned 2 sessions",
             "Request for 1 multiplexed session returned 1 session",
             "Creating 2 sessions");
-    if (isMultiplexedSessionsEnabled()) {
+    List<String> expectedAnnotationsForMultiplexedSessionsRW =
+        ImmutableList.of(
+            "Starting Transaction Attempt",
+            "Starting Commit",
+            "Commit Done",
+            "Transaction Attempt Succeeded",
+            "Requesting 2 sessions",
+            "Request for 2 sessions returned 2 sessions",
+            "Request for 1 multiplexed session returned 1 session",
+            "Creating 2 sessions");
+    if (isMultiplexedSessionsEnabledForRW()) {
+      verifyAnnotations(
+          failOnOverkillTraceComponent.getAnnotations().stream()
+              .distinct()
+              .collect(Collectors.toList()),
+          expectedAnnotationsForMultiplexedSessionsRW);
+    } else if (isMultiplexedSessionsEnabled()) {
       verifyAnnotations(
           failOnOverkillTraceComponent.getAnnotations().stream()
               .distinct()
@@ -538,7 +554,21 @@ public void transactionRunnerWithError() {
             "Request for 1 multiplexed session returned 1 session",
             "Request for 2 sessions returned 2 sessions",
             "Creating 2 sessions");
-    if (isMultiplexedSessionsEnabled()) {
+    List<String> expectedAnnotationsForMultiplexedSessionsRW =
+        ImmutableList.of(
+            "Starting Transaction Attempt",
+            "Transaction Attempt Failed in user operation",
+            "Requesting 2 sessions",
+            "Request for 1 multiplexed session returned 1 session",
+            "Request for 2 sessions returned 2 sessions",
+            "Creating 2 sessions");
+    if (isMultiplexedSessionsEnabledForRW()) {
+      verifyAnnotations(
+          failOnOverkillTraceComponent.getAnnotations().stream()
+              .distinct()
+              .collect(Collectors.toList()),
+          expectedAnnotationsForMultiplexedSessionsRW);
+    } else if (isMultiplexedSessionsEnabled()) {
       verifyAnnotations(
           failOnOverkillTraceComponent.getAnnotations().stream()
               .distinct()
@@ -565,4 +595,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();
+  }
 }