Address comments

Author: sakthivelmanii

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java

@@ -752,9 +752,8 @@ public class ConnectionProperties {
   static final ConnectionProperty<Long> DML_BATCH_UPDATE_COUNT =
       create(
           DML_BATCH_UPDATE_COUNT_PROPERTY_NAME,
-          "DML statements that are executed in batch when "
-              + " 'start batch dml' statement is executed. This property determines"
-              + " the update count that is returned for these DML statements. The default is "
+          "The update count that is returned for DML statements that are executed in an "
+              + "explicit DML batch. The default is "
               + DEFAULT_DML_BATCH_UPDATE_COUNT,
           DEFAULT_DML_BATCH_UPDATE_COUNT,
           LongConverter.INSTANCE,

google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/DmlBatch.java

@@ -84,6 +84,7 @@ Builder setAutoBatchUpdateCountVerificationSupplier(Supplier<Boolean> verificati
     }
 
     Builder setDmlBatchUpdateCountSupplier(Supplier<Long> dmlbatchUpdateCountSupplier) {
+      Preconditions.checkNotNull(dmlbatchUpdateCountSupplier);
       this.dmlbatchUpdateCountSupplier = dmlbatchUpdateCountSupplier;
       return this;
     }

google-cloud-spanner/src/main/resources/com/google/cloud/spanner/connection/PG_ClientSideStatements.json

@@ -976,6 +976,7 @@
       "exampleStatements": [
         "set local spanner.dml_batch_update_count = 0",
         "set local spanner.dml_batch_update_count = 100",
+        "set local spanner.dml_batch_update_count to 1",
         "set local spanner.dml_batch_update_count to 1"
       ],
       "setStatement": {

google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionAsyncApiTest.java

@@ -342,45 +342,22 @@ public void testAutocommitRunBatch() {
   @Test
   public void testDmlBatchUpdateCount() {
     SpannerPool.closeSpannerPool();
-    mockSpanner.putStatementResult(
-        MockSpannerServiceImpl.StatementResult.detectDialectResult(Dialect.POSTGRESQL));
     try {
       try (Connection connection = createConnection()) {
         connection.execute(Statement.of("set local spanner.dml_batch_update_count = 1"));
         connection.execute(Statement.of("START BATCH DML"));
         List<Statement> statements = Arrays.asList(INSERT_STATEMENT, INSERT_STATEMENT);
         long[] updateCounts = connection.executeBatchUpdate(statements);
         assertThat(updateCounts).asList().containsExactly(1L, 1L);
-      }
-      try (Connection connection = createConnection()) {
-        connection.execute(Statement.of("START BATCH DML"));
-        List<Statement> statements = Arrays.asList(INSERT_STATEMENT, INSERT_STATEMENT);
-        long[] updateCounts = connection.executeBatchUpdate(statements);
-        assertThat(updateCounts).asList().containsExactly(-1L, -1L);
-      }
-    } finally {
-      SpannerPool.closeSpannerPool();
-      mockSpanner.putStatementResult(
-          MockSpannerServiceImpl.StatementResult.detectDialectResult(Dialect.GOOGLE_STANDARD_SQL));
-    }
-  }
+        connection.execute(Statement.of("RUN BATCH"));
+        connection.commit();
 
-  @Test
-  public void testDmlBatchUpdateCountGoogleSql() {
-    SpannerPool.closeSpannerPool();
-    try {
-      try (Connection connection = createConnection()) {
-        connection.execute(Statement.of("set local dml_batch_update_count = 1"));
-        connection.execute(Statement.of("START BATCH DML"));
-        List<Statement> statements = Arrays.asList(INSERT_STATEMENT, INSERT_STATEMENT);
-        long[] updateCounts = connection.executeBatchUpdate(statements);
-        assertThat(updateCounts).asList().containsExactly(1L, 1L);
-      }
-      try (Connection connection = createConnection()) {
         connection.execute(Statement.of("START BATCH DML"));
-        List<Statement> statements = Arrays.asList(INSERT_STATEMENT, INSERT_STATEMENT);
-        long[] updateCounts = connection.executeBatchUpdate(statements);
+        statements = Arrays.asList(INSERT_STATEMENT, INSERT_STATEMENT);
+        updateCounts = connection.executeBatchUpdate(statements);
         assertThat(updateCounts).asList().containsExactly(-1L, -1L);
+        connection.execute(Statement.of("RUN BATCH"));
+        connection.commit();
       }
     } finally {
       SpannerPool.closeSpannerPool();