adding unit test for batch request failue (#851)

GoogleCloudDataproc · Aug 16, 2022 · 782b3e9 · 782b3e9
1 parent 8669e68
commit 782b3e9
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 2 deletions.
diff --git a/gcsio/src/test/java/com/google/cloud/hadoop/gcsio/GoogleCloudStorageTest.java b/gcsio/src/test/java/com/google/cloud/hadoop/gcsio/GoogleCloudStorageTest.java
@@ -3137,6 +3137,42 @@ public void testGetItemInfos() throws IOException {
         .inOrder();
   }
 
+  @Test
+  public void testGetItemInfosWithRetries() throws IOException {
+    Bucket bucket = newBucket(BUCKET_NAME);
+    StorageObject storageObject = newStorageObject(BUCKET_NAME, OBJECT_NAME);
+
+    MockHttpTransport transport =
+        mockBatchTransport(
+            /* requestsPerBatch= */ 2,
+            jsonDataResponse(storageObject),
+            jsonErrorResponse(ErrorResponses.RATE_LIMITED),
+            jsonDataResponse(bucket));
+
+    GoogleCloudStorage gcs = mockedGcs(transport);
+
+    // Call in order of StorageObject, Bucket.
+    List<GoogleCloudStorageItemInfo> itemInfos =
+        gcs.getItemInfos(ImmutableList.of(RESOURCE_ID, new StorageResourceId(BUCKET_NAME)));
+
+    assertThat(itemInfos)
+        .containsExactly(
+            createItemInfoForStorageObject(RESOURCE_ID, storageObject),
+            createItemInfoForBucket(new StorageResourceId(BUCKET_NAME), bucket))
+        .inOrder();
+
+    assertThat(trackingRequestInitializerWithRetries.getAllRequestStrings())
+        .containsExactly(
+            // Request of 1st batch
+            batchRequestString(),
+            getRequestString(BUCKET_NAME, OBJECT_NAME),
+            getBucketRequestString(BUCKET_NAME),
+            // Request of 2nd batch
+            batchRequestString(),
+            getBucketRequestString(BUCKET_NAME))
+        .inOrder();
+  }
+
   @Test
   public void testGetItemInfosNotFound() throws IOException {
     MockHttpTransport transport =

diff --git a/util/src/main/java/com/google/cloud/hadoop/util/testing/MockHttpTransportHelper.java b/util/src/main/java/com/google/cloud/hadoop/util/testing/MockHttpTransportHelper.java
@@ -132,12 +132,14 @@ public static MockHttpTransport mockBatchTransport(
       int requestsPerBatch, LowLevelHttpResponse... responses) {
     return new MockHttpTransport() {
       int responsesIndex = 0;
+      int responseBatchCounter = 0;
 
       @Override
       public LowLevelHttpRequest buildRequest(String method, String url) {
         return new MockLowLevelHttpRequest() {
           @Override
           public LowLevelHttpResponse execute() {
+            responseBatchCounter++;
             String boundary = "batch_pK7JBAk73-E=_AA5eFwv4m2Q=";
             String contentId = "";
 
@@ -147,8 +149,10 @@ public LowLevelHttpResponse execute() {
                     .setContentType("multipart/mixed; boundary=" + boundary);
 
             StringBuilder batchResponse = new StringBuilder();
-
-            for (int i = 0; i < requestsPerBatch; i++) {
+            int remainingResponses = responses.length - responsesIndex;
+            int responsesInBatch =
+                Math.min((responseBatchCounter * requestsPerBatch), remainingResponses);
+            for (int i = 0; i < responsesInBatch; i++) {
               try {
                 LowLevelHttpResponse resp = responses[responsesIndex++];
                 batchResponse