chore(spanner): test fixes for mux rw (#3649)

Author: harshachinta

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/it/ITTransactionRetryTest.java

@@ -1521,6 +1521,10 @@ public void testRetryLargeResultSet() {
   /** Test the successful retry of a transaction with a high chance of multiple aborts */
   @Test
   public void testRetryHighAbortRate() {
+    // TODO(sriharshach): Remove this skip once backend support empty transactions to commit.
+    assumeFalse(
+        "Skipping for multiplexed sessions since it does not allow empty transactions to commit",
+        env.getTestHelper().getOptions().getSessionPoolOptions().getUseMultiplexedSessionForRW());
     final int NUMBER_OF_TEST_RECORDS = 10000;
     final long UPDATED_RECORDS = 1000L;
     // abort on 25% of all statements

google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITTransactionManagerAsyncTest.java

@@ -161,7 +161,16 @@ public void testInvalidInsert() throws InterruptedException {
         } catch (ExecutionException e) {
           assertThat(e.getCause()).isInstanceOf(SpannerException.class);
           SpannerException se = (SpannerException) e.getCause();
-          assertThat(se.getErrorCode()).isEqualTo(ErrorCode.NOT_FOUND);
+          if (env.getTestHelper()
+              .getOptions()
+              .getSessionPoolOptions()
+              .getUseMultiplexedSessionForRW()) {
+            // Backend currently returns INVALID_ARGUMENT, however this will be changed to NOT_FOUND
+            // in future.
+            assertThat(se.getErrorCode()).isAnyOf(ErrorCode.NOT_FOUND, ErrorCode.INVALID_ARGUMENT);
+          } else {
+            assertThat(se.getErrorCode()).isEqualTo(ErrorCode.NOT_FOUND);
+          }
           // expected
           break;
         }

google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITWriteTest.java

@@ -1043,7 +1043,16 @@ public void tableNotFound() {
               .build());
       fail("Expected exception");
     } catch (SpannerException ex) {
-      assertThat(ex.getErrorCode()).isEqualTo(ErrorCode.NOT_FOUND);
+      if (env.getTestHelper()
+          .getOptions()
+          .getSessionPoolOptions()
+          .getUseMultiplexedSessionForRW()) {
+        // Backend currently returns INVALID_ARGUMENT, however this will be changed to NOT_FOUND in
+        // future.
+        assertThat(ex.getErrorCode()).isAnyOf(ErrorCode.NOT_FOUND, ErrorCode.INVALID_ARGUMENT);
+      } else {
+        assertThat(ex.getErrorCode()).isEqualTo(ErrorCode.NOT_FOUND);
+      }
     }
   }
 
@@ -1053,7 +1062,16 @@ public void columnNotFound() {
       write(baseInsert().set("ColumnThatDoesNotExist").to("V1").build());
       fail("Expected exception");
     } catch (SpannerException ex) {
-      assertThat(ex.getErrorCode()).isEqualTo(ErrorCode.NOT_FOUND);
+      if (env.getTestHelper()
+          .getOptions()
+          .getSessionPoolOptions()
+          .getUseMultiplexedSessionForRW()) {
+        // Backend currently returns INVALID_ARGUMENT, however this will be changed to NOT_FOUND in
+        // future.
+        assertThat(ex.getErrorCode()).isAnyOf(ErrorCode.NOT_FOUND, ErrorCode.INVALID_ARGUMENT);
+      } else {
+        assertThat(ex.getErrorCode()).isEqualTo(ErrorCode.NOT_FOUND);
+      }
     }
   }