fix tests

googleapis · Nov 14, 2024 · 4fa9b28 · 4fa9b28
1 parent 131966c
commit 4fa9b28
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SessionPoolTest.java
@@ -1729,10 +1729,13 @@ public void testSessionNotFoundPartitionedUpdate() {
         SpannerExceptionFactoryTest.newSessionNotFoundException(sessionName);
     Statement statement = Statement.of("UPDATE FOO SET BAR=1 WHERE 1=1");
     final SessionImpl closedSession = mockSession();
-    when(closedSession.executePartitionedUpdate(statement)).thenThrow(sessionNotFound);
+    when(closedSession.executePartitionedUpdate(any(Statement.class))).thenThrow(sessionNotFound);
+    when(closedSession.executePartitionedUpdate(any(Statement.class), any(), any()))
+        .thenThrow(sessionNotFound);
 
     final SessionImpl openSession = mockSession();
-    when(openSession.executePartitionedUpdate(statement)).thenReturn(1L);
+    when(openSession.executePartitionedUpdate(any(Statement.class))).thenReturn(1L);
+    when(openSession.executePartitionedUpdate(any(Statement.class), any(), any())).thenReturn(1L);
     doAnswer(
             invocation -> {
               executor.submit(