From 64a93e5fa5117a1ade3732472d39c90f5bb98cb2 Mon Sep 17 00:00:00 2001
From: Sri Harsha CH <sriharshach@google.com>
Date: Mon, 23 Dec 2024 11:07:39 +0000
Subject: [PATCH] chore(spanner): add mocks for create session with multiplexed
 session unit tests

---
 .../spanner/TransactionManagerImplTest.java   | 30 +++++++++++++++++++
 .../spanner/TransactionRunnerImplTest.java    | 15 ++++++++++
 2 files changed, 45 insertions(+)

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
index 10b1312515..aee3d5ed5b 100644
--- 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
index 1fd6817ea9..d8bd6ed448 100644
--- 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(