GcsUtliTest's mock not to supply content stream for error

Since google-http-client 1.39.1, it does not read content stream when a response has error status. The related pull request is googleapis/google-http-java-client#1315
suztomo · Apr 14, 2021 · 855a216 · 855a216
1 parent 2b36069
commit 855a216
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 11 deletions.
diff --git a/...loud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java b/...loud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java
@@ -502,7 +502,9 @@ public void testGetSizeBytesWhenFileNotFoundBatchRetry() throws Exception {
 
     // 429: Too many requests, then 200: OK.
     when(mockResponse.getStatusCode()).thenReturn(429, 200);
-    when(mockResponse.getContent()).thenReturn(toStream("error"), toStream(content));
+    // Google-http-client:1.39.1 and above does not read content stream of the error response.
+    // Therefore the mock only supplies the content of success case.
+    when(mockResponse.getContent()).thenReturn(toStream(content));
 
     // A mock transport that lets us mock the API responses.
     MockHttpTransport mockTransport =

diff --git a/...a/org/apache/beam/sdk/extensions/gcp/util/LatencyRecordingHttpRequestInitializerTest.java b/...a/org/apache/beam/sdk/extensions/gcp/util/LatencyRecordingHttpRequestInitializerTest.java
@@ -103,7 +103,7 @@ public void testOkResponse() throws IOException {
     verify(mockLowLevelRequest).setTimeout(anyInt(), anyInt());
     verify(mockLowLevelRequest).setWriteTimeout(anyInt());
     verify(mockLowLevelRequest).execute();
-    verify(mockLowLevelResponse).getStatusCode();
+    verify(mockLowLevelResponse, atLeastOnce()).getStatusCode();
   }
 
   @Test
@@ -124,6 +124,6 @@ public void testErrorResponse() throws IOException {
     verify(mockLowLevelRequest).setTimeout(anyInt(), anyInt());
     verify(mockLowLevelRequest).setWriteTimeout(anyInt());
     verify(mockLowLevelRequest).execute();
-    verify(mockLowLevelResponse).getStatusCode();
+    verify(mockLowLevelResponse, atLeastOnce()).getStatusCode();
   }
 }
diff --git a/...rc/test/java/org/apache/beam/sdk/extensions/gcp/util/RetryHttpRequestInitializerTest.java b/...rc/test/java/org/apache/beam/sdk/extensions/gcp/util/RetryHttpRequestInitializerTest.java
@@ -138,7 +138,7 @@ public void testBasicOperation() throws IOException {
     verify(mockLowLevelRequest).setTimeout(anyInt(), anyInt());
     verify(mockLowLevelRequest).setWriteTimeout(anyInt());
     verify(mockLowLevelRequest).execute();
-    verify(mockLowLevelResponse).getStatusCode();
+    verify(mockLowLevelResponse, atLeastOnce()).getStatusCode();
     expectedLogs.verifyNotLogged("Request failed");
   }
 
@@ -209,7 +209,7 @@ public void testThrowIOException() throws IOException {
     verify(mockLowLevelRequest, times(2)).setTimeout(anyInt(), anyInt());
     verify(mockLowLevelRequest, times(2)).setWriteTimeout(anyInt());
     verify(mockLowLevelRequest, times(2)).execute();
-    verify(mockLowLevelResponse).getStatusCode();
+    verify(mockLowLevelResponse, atLeastOnce()).getStatusCode();
     expectedLogs.verifyDebug("Request failed with IOException");
   }
 

diff --git a/...-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java b/...-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java
@@ -191,8 +191,8 @@ public void testStartLoadJobRetry() throws IOException, InterruptedException {
     // First response is 403 rate limited, second response has valid payload.
     when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
     when(response.getStatusCode()).thenReturn(403).thenReturn(200);
+    // Google-http-client 1.39.1 and above only reads content of successful responses
     when(response.getContent())
-        .thenReturn(toStream(errorWithReasonAndStatus("rateLimitExceeded", 403)))
         .thenReturn(toStream(testJob));
 
     Sleeper sleeper = new FastNanoClockAndSleeper();
@@ -204,8 +204,9 @@ public void testStartLoadJobRetry() throws IOException, InterruptedException {
         BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()));
 
     verify(response, times(2)).getStatusCode();
-    verify(response, times(2)).getContent();
-    verify(response, times(2)).getContentType();
+    // Only the successful responses is read
+    verify(response, times(1)).getContent();
+    verify(response, times(1)).getContentType();
   }
 
   /** Tests that {@link BigQueryServicesImpl.JobServiceImpl#pollJob} succeeds. */
@@ -334,8 +335,8 @@ public void testGetTableSucceeds() throws Exception {
 
     when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
     when(response.getStatusCode()).thenReturn(403).thenReturn(200);
+    // Google-http-client 1.39.1 and above only reads content of successful responses
     when(response.getContent())
-        .thenReturn(toStream(errorWithReasonAndStatus("rateLimitExceeded", 403)))
         .thenReturn(toStream(testTable));
 
     BigQueryServicesImpl.DatasetServiceImpl datasetService =
@@ -346,8 +347,9 @@ public void testGetTableSucceeds() throws Exception {
 
     assertEquals(testTable, table);
     verify(response, times(2)).getStatusCode();
-    verify(response, times(2)).getContent();
-    verify(response, times(2)).getContentType();
+    // Only the successful response is read
+    verify(response, times(1)).getContent();
+    verify(response, times(1)).getContentType();
   }
 
   @Test