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 index 798b7476a3..fb4f5734f6 100644 --- 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 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 index 02938ff7e5..1102fb4759 100644 --- 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