test: add test for auto_batch_dml

googleapis · Oct 12, 2024 · 3d87ec4 · 3d87ec4
1 parent 57df3f3
commit 3d87ec4
Showing 1 changed file with 26 additions and 0 deletions.
diff --git a/src/test/java/com/google/cloud/spanner/jdbc/ExecuteMockServerTest.java b/src/test/java/com/google/cloud/spanner/jdbc/ExecuteMockServerTest.java
@@ -36,6 +36,8 @@
 import com.google.protobuf.Value;
 import com.google.rpc.Code;
 import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
+import com.google.spanner.v1.CommitRequest;
+import com.google.spanner.v1.ExecuteBatchDmlRequest;
 import com.google.spanner.v1.ResultSetMetadata;
 import com.google.spanner.v1.ResultSetStats;
 import com.google.spanner.v1.StructType;
@@ -844,4 +846,28 @@ public void testInvalidExecuteUpdate_shouldNotLeakSession() throws SQLException
  }
  }
  }
+
+ private String getExtension() {
+ return dialect == Dialect.POSTGRESQL ? "spanner." : "";
+ }
+
+ @Test
+ public void testExecuteAutoBatchDml() throws SQLException {
+ try (Connection connection = createJdbcConnection();
+ Statement statement = connection.createStatement()) {
+ connection.setAutoCommit(false);
+
+ assertFalse(statement.execute(String.format("set %sauto_batch_dml = true", getExtension())));
+ for (int i = 0; i < 3; i++) {
+ assertFalse(statement.execute(dml));
+ assertEquals(1, statement.getUpdateCount());
+ }
+ connection.commit();
+ }
+ assertEquals(1, mockSpanner.countRequestsOfType(ExecuteBatchDmlRequest.class));
+ ExecuteBatchDmlRequest request =
+ mockSpanner.getRequestsOfType(ExecuteBatchDmlRequest.class).get(0);
+ assertEquals(3, request.getStatementsCount());
+ assertEquals(1, mockSpanner.countRequestsOfType(CommitRequest.class));
+ }
 }