From 2ec87681987c15f9e074ddbb9c4c02fcdd74bff6 Mon Sep 17 00:00:00 2001
From: Yiru Tang <yiru@google.com>
Date: Wed, 15 Nov 2023 00:34:07 +0000
Subject: [PATCH] fix: enable request limit e2e test.

---
 .../it/ITBigQueryWriteManualClientTest.java   | 83 +++++++++----------
 1 file changed, 40 insertions(+), 43 deletions(-)

diff --git a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java
index 676d9a5d85..36169b89b4 100644
--- a/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java
+++ b/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/it/ITBigQueryWriteManualClientTest.java
@@ -1559,49 +1559,46 @@ public void testMultiplexingMixedLocation()
     assertEquals("eu", streamWriter3.getLocation());
   }
 
-  // Tested locally but project config is frozen and we need to wait for a while to enable the
-  // test in automatic workflow.
-  // @Test
-  // public void testLargeRequest() throws IOException, InterruptedException, ExecutionException {
-  //   String tableName = "largeRequestTable";
-  //   TableId tableId = TableId.of(DATASET, tableName);
-  //   Field col1 = Field.newBuilder("col1", StandardSQLTypeName.STRING).build();
-  //   Schema originalSchema = Schema.of(col1);
-  //   TableInfo tableInfo =
-  //       TableInfo.newBuilder(tableId, StandardTableDefinition.of(originalSchema)).build();
-  //   bigquery.create(tableInfo);
-  //   TableName parent = TableName.of(ServiceOptions.getDefaultProjectId(), DATASET, tableName);
-  //   try (StreamWriter streamWriter =
-  //       StreamWriter.newBuilder(parent.toString() + "/_default")
-  //           .setWriterSchema(CreateProtoSchemaWithColField())
-  //           .build()) {
-  //     List<Integer> sizeSet = Arrays.asList(15 * 1024 * 1024, 1024);
-  //     List<ApiFuture<AppendRowsResponse>> responseList =
-  //         new ArrayList<ApiFuture<AppendRowsResponse>>();
-  //     Random r = new Random();
-  //     for (int i = 0; i < 50; i++) {
-  //       int size = sizeSet.get(r.nextInt(2));
-  //       LOG.info("Sending size: " + size);
-  //       responseList.add(
-  //           streamWriter.append(
-  //               CreateProtoRows(
-  //                   new String[] {
-  //                     new String(new char[size]).replace('\u0000', (char) (r.nextInt(26) + 'a'))
-  //                   })));
-  //     }
-  //     for (int i = 0; i < 50; i++) {
-  //       assertFalse(responseList.get(i).get().hasError());
-  //     }
-  //     TableResult queryResult =
-  //         bigquery.query(
-  //             QueryJobConfiguration.newBuilder("SELECT count(*) from " + DATASET + '.' +
-  // tableName)
-  //                 .build());
-  //     Iterator<FieldValueList> queryIter = queryResult.getValues().iterator();
-  //     assertTrue(queryIter.hasNext());
-  //     assertEquals("50", queryIter.next().get(0).getStringValue());
-  //   }
-  // }
+  @Test
+  public void testLargeRequest() throws IOException, InterruptedException, ExecutionException {
+    String tableName = "largeRequestTable";
+    TableId tableId = TableId.of(DATASET, tableName);
+    Field col1 = Field.newBuilder("col1", StandardSQLTypeName.STRING).build();
+    Schema originalSchema = Schema.of(col1);
+    TableInfo tableInfo =
+        TableInfo.newBuilder(tableId, StandardTableDefinition.of(originalSchema)).build();
+    bigquery.create(tableInfo);
+    TableName parent = TableName.of(ServiceOptions.getDefaultProjectId(), DATASET, tableName);
+    try (StreamWriter streamWriter =
+        StreamWriter.newBuilder(parent.toString() + "/_default")
+            .setWriterSchema(CreateProtoSchemaWithColField())
+            .build()) {
+      List<Integer> sizeSet = Arrays.asList(15 * 1024 * 1024, 1024);
+      List<ApiFuture<AppendRowsResponse>> responseList =
+          new ArrayList<ApiFuture<AppendRowsResponse>>();
+      Random r = new Random();
+      for (int i = 0; i < 50; i++) {
+        int size = sizeSet.get(r.nextInt(2));
+        LOG.info("Sending size: " + size);
+        responseList.add(
+            streamWriter.append(
+                CreateProtoRows(
+                    new String[] {
+                      new String(new char[size]).replace('\u0000', (char) (r.nextInt(26) + 'a'))
+                    })));
+      }
+      for (int i = 0; i < 50; i++) {
+        assertFalse(responseList.get(i).get().hasError());
+      }
+      TableResult queryResult =
+          bigquery.query(
+              QueryJobConfiguration.newBuilder("SELECT count(*) from " + DATASET + '.' + tableName)
+                  .build());
+      Iterator<FieldValueList> queryIter = queryResult.getValues().iterator();
+      assertTrue(queryIter.hasNext());
+      assertEquals("50", queryIter.next().get(0).getStringValue());
+    }
+  }
 
   @Test
   public void testDefaultRequestLimit()