From 84cd62f3e54986bfe29944ae5e5039bf4ece3976 Mon Sep 17 00:00:00 2001
From: Arpan Mishra <arpanmishra@google.com>
Date: Tue, 19 Sep 2023 16:47:52 +0530
Subject: [PATCH] test: add integration tests for large batch dml updates.
 (#2618)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* test: add integration tests for large batch dml updates.

* chore: incorporate review comments.

* fix: lint issues.

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

---------

Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
---
 .../cloud/spanner/it/ITBatchDmlTest.java      | 63 +++++++++++++++++++
 1 file changed, 63 insertions(+)

diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITBatchDmlTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITBatchDmlTest.java
index 3cd9917e1a8..b11e4f613ce 100644
--- a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITBatchDmlTest.java
+++ b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/it/ITBatchDmlTest.java
@@ -31,7 +31,9 @@
 import com.google.cloud.spanner.TransactionRunner.TransactionCallable;
 import com.google.spanner.admin.database.v1.UpdateDatabaseDdlMetadata;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
+import java.util.LinkedList;
 import java.util.List;
 import org.junit.After;
 import org.junit.Assert;
@@ -189,4 +191,65 @@ public void errorBatchDmlAlreadyExist() {
       }
     }
   }
+
+  @Test
+  public void largeBatchDml() {
+    List<Statement> stmts = new LinkedList<>();
+    for (int i = 0; i < 40; i++) {
+      stmts.add(Statement.of("INSERT INTO T (k, v) VALUES ('boo" + i + "', " + i + ");"));
+    }
+
+    for (int i = 0; i < 40; i++) {
+      stmts.add(Statement.of("DELETE FROM T WHERE T.K = 'boo" + i + "';"));
+    }
+    long[] expectedRowCounts = new long[stmts.size()];
+    Arrays.fill(expectedRowCounts, 1L);
+
+    final TransactionCallable<long[]> callable = transaction -> transaction.batchUpdate(stmts);
+    TransactionRunner runner = client.readWriteTransaction();
+    long[] actualRowCounts = runner.run(callable);
+    assertThat(actualRowCounts.length).isEqualTo(80);
+    assertThat(expectedRowCounts).isEqualTo(actualRowCounts);
+  }
+
+  @Test
+  public void largeBatchDml_withParameterisedStatements() {
+    List<Statement> stmts = new LinkedList<>();
+    String insertQuery = "INSERT INTO T(k, v) VALUES(@key, @val)";
+    for (int i = 0; i < 80; i++) {
+      stmts.add(
+          Statement.newBuilder(insertQuery)
+              .bind("key")
+              .to("'boo" + i + "'")
+              .bind("val")
+              .to(i)
+              .build());
+    }
+    long[] expectedRowCounts = new long[stmts.size()];
+    Arrays.fill(expectedRowCounts, 1L);
+
+    final TransactionCallable<long[]> callable = transaction -> transaction.batchUpdate(stmts);
+    TransactionRunner runner = client.readWriteTransaction();
+    long[] actualRowCounts = runner.run(callable);
+
+    assertThat(actualRowCounts.length).isEqualTo(80);
+    assertThat(expectedRowCounts).isEqualTo(actualRowCounts);
+  }
+
+  @Test
+  public void largeBatchDml_withNonParameterisedStatements() {
+    List<Statement> stmts = new LinkedList<>();
+    for (int i = 0; i < 80; i++) {
+      stmts.add(Statement.of("INSERT INTO T (k, v) VALUES ('boo" + i + "', " + i + ");"));
+    }
+    long[] expectedRowCounts = new long[stmts.size()];
+    Arrays.fill(expectedRowCounts, 1L);
+
+    final TransactionCallable<long[]> callable = transaction -> transaction.batchUpdate(stmts);
+    TransactionRunner runner = client.readWriteTransaction();
+    long[] actualRowCounts = runner.run(callable);
+
+    assertThat(actualRowCounts.length).isEqualTo(80);
+    assertThat(expectedRowCounts).isEqualTo(actualRowCounts);
+  }
 }