chore(spanner): add mocks for create session with multiplexed session…

… unit tests
googleapis · Dec 23, 2024 · 64a93e5 · 64a93e5
1 parent a1e72c4
commit 64a93e5
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 0 deletions.
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionManagerImplTest.java
@@ -234,6 +234,21 @@ public void usesPreparedTransaction() {
                             com.google.protobuf.Timestamp.newBuilder()
                                 .setSeconds(System.currentTimeMillis() * 1000))
                         .build()));
+    when(rpc.createSession(
+            Mockito.anyString(),
+            Mockito.anyString(),
+            Mockito.anyMap(),
+            Mockito.eq(null),
+            Mockito.eq(true)))
+        .thenAnswer(
+            invocation ->
+                Session.newBuilder()
+                    .setName(invocation.getArguments()[0] + "/sessions/1")
+                    .setMultiplexed(true)
+                    .setCreateTime(
+                        com.google.protobuf.Timestamp.newBuilder()
+                            .setSeconds(System.currentTimeMillis() * 1000))
+                    .build());
     when(rpc.beginTransactionAsync(
             Mockito.any(BeginTransactionRequest.class), Mockito.anyMap(), eq(true)))
         .thenAnswer(
@@ -300,6 +315,21 @@ public void inlineBegin() {
                             com.google.protobuf.Timestamp.newBuilder()
                                 .setSeconds(System.currentTimeMillis() * 1000))
                         .build()));
+    when(rpc.createSession(
+            Mockito.anyString(),
+            Mockito.anyString(),
+            Mockito.anyMap(),
+            Mockito.eq(null),
+            Mockito.eq(true)))
+        .thenAnswer(
+            invocation ->
+                Session.newBuilder()
+                    .setName(invocation.getArguments()[0] + "/sessions/1")
+                    .setMultiplexed(true)
+                    .setCreateTime(
+                        com.google.protobuf.Timestamp.newBuilder()
+                            .setSeconds(System.currentTimeMillis() * 1000))
+                    .build());
     when(rpc.beginTransactionAsync(
             Mockito.any(BeginTransactionRequest.class), Mockito.anyMap(), eq(true)))
         .thenAnswer(

diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/TransactionRunnerImplTest.java
@@ -186,6 +186,21 @@ public void usesPreparedTransaction() {
                         .setCreateTime(
                             Timestamp.newBuilder().setSeconds(System.currentTimeMillis() * 1000))
                         .build()));
+    when(rpc.createSession(
+            Mockito.anyString(),
+            Mockito.anyString(),
+            Mockito.anyMap(),
+            Mockito.eq(null),
+            Mockito.eq(true)))
+        .thenAnswer(
+            invocation ->
+                Session.newBuilder()
+                    .setName(invocation.getArguments()[0] + "/sessions/1")
+                    .setMultiplexed(true)
+                    .setCreateTime(
+                        com.google.protobuf.Timestamp.newBuilder()
+                            .setSeconds(System.currentTimeMillis() * 1000))
+                    .build());
     when(rpc.beginTransactionAsync(
             Mockito.any(BeginTransactionRequest.class), Mockito.anyMap(), eq(true)))
         .thenAnswer(