From c27b602bde4fd5bd2fb8d36c5b6a163f2350bf7f Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Fri, 25 Sep 2026 14:42:10 -0400 Subject: [PATCH 1/6] Split V1 tests out of GcsUtilTest into GcsUtilV1Test --- .../sdk/extensions/gcp/util/GcsUtilTest.java | 1789 +--------------- .../extensions/gcp/util/GcsUtilV1Test.java | 1850 +++++++++++++++++ .../extensions/gcp/util/GcsUtilV2Test.java | 4 +- 3 files changed, 1877 insertions(+), 1766 deletions(-) create mode 100644 sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1Test.java diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java index 4beb56d6c200..4a4c4eda28c8 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java @@ -17,11 +17,6 @@ */ package org.apache.beam.sdk.extensions.gcp.util; -import static org.apache.beam.sdk.options.ExperimentalOptions.hasExperiment; -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.contains; -import static org.hamcrest.Matchers.equalTo; -import static org.hamcrest.Matchers.greaterThan; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotNull; @@ -30,77 +25,30 @@ import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import com.google.api.client.googleapis.batch.json.JsonBatchCallback; -import com.google.api.client.googleapis.json.GoogleJsonError; -import com.google.api.client.googleapis.json.GoogleJsonError.ErrorInfo; -import com.google.api.client.googleapis.json.GoogleJsonResponseException; -import com.google.api.client.googleapis.services.json.AbstractGoogleJsonClientRequest; -import com.google.api.client.http.HttpRequest; -import com.google.api.client.http.HttpRequestInitializer; -import com.google.api.client.http.HttpResponse; -import com.google.api.client.http.HttpStatusCodes; -import com.google.api.client.http.HttpTransport; -import com.google.api.client.http.LowLevelHttpRequest; -import com.google.api.client.http.LowLevelHttpResponse; -import com.google.api.client.json.GenericJson; -import com.google.api.client.json.Json; -import com.google.api.client.json.JsonFactory; -import com.google.api.client.json.gson.GsonFactory; -import com.google.api.client.testing.http.HttpTesting; -import com.google.api.client.testing.http.MockHttpTransport; -import com.google.api.client.testing.http.MockLowLevelHttpRequest; -import com.google.api.client.testing.http.MockLowLevelHttpResponse; -import com.google.api.client.util.BackOff; import com.google.api.gax.paging.Page; -import com.google.api.services.storage.Storage; import com.google.api.services.storage.model.Bucket; import com.google.api.services.storage.model.Objects; -import com.google.api.services.storage.model.RewriteResponse; import com.google.api.services.storage.model.StorageObject; -import com.google.auth.Credentials; import com.google.cloud.WriteChannel; -import com.google.cloud.hadoop.gcsio.CreateObjectOptions; -import com.google.cloud.hadoop.gcsio.GoogleCloudStorage; -import com.google.cloud.hadoop.gcsio.GoogleCloudStorageImpl; -import com.google.cloud.hadoop.gcsio.GoogleCloudStorageOptions; -import com.google.cloud.hadoop.gcsio.GoogleCloudStorageReadOptions; -import com.google.cloud.hadoop.gcsio.StorageResourceId; import com.google.cloud.storage.BucketInfo; import com.google.cloud.storage.Storage.BucketGetOption; import com.google.cloud.storage.Storage.BucketTargetOption; import com.google.cloud.storage.Storage.PredefinedAcl; import com.google.cloud.storage.StorageClass; import com.google.cloud.storage.StorageException; -import java.io.ByteArrayInputStream; -import java.io.ByteArrayOutputStream; import java.io.FileNotFoundException; import java.io.IOException; -import java.io.InputStream; -import java.io.InvalidObjectException; import java.math.BigInteger; -import java.net.SocketTimeoutException; -import java.nio.ByteBuffer; -import java.nio.channels.Channels; import java.nio.channels.SeekableByteChannel; import java.nio.channels.WritableByteChannel; -import java.nio.charset.StandardCharsets; -import java.nio.file.AccessDeniedException; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.Collections; import java.util.HashMap; -import java.util.LinkedList; import java.util.List; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; import java.util.concurrent.TimeUnit; -import java.util.function.Supplier; -import org.apache.beam.repackaged.core.org.apache.commons.compress.utils.SeekableInMemoryByteChannel; import org.apache.beam.runners.core.metrics.GcpResourceIdentifiers; import org.apache.beam.runners.core.metrics.MetricsContainerImpl; import org.apache.beam.runners.core.metrics.MonitoringInfoConstants; @@ -109,19 +57,11 @@ import org.apache.beam.sdk.extensions.gcp.options.GcsOptions; import org.apache.beam.sdk.extensions.gcp.util.GcsUtil.CreateOptions; import org.apache.beam.sdk.extensions.gcp.util.GcsUtil.StorageObjectOrIOException; -import org.apache.beam.sdk.extensions.gcp.util.GcsUtilV1.BatchInterface; -import org.apache.beam.sdk.extensions.gcp.util.GcsUtilV1.RewriteOp; import org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath; import org.apache.beam.sdk.io.fs.MoveOptions.StandardMoveOptions; -import org.apache.beam.sdk.metrics.MetricName; import org.apache.beam.sdk.metrics.MetricsEnvironment; -import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; -import org.apache.beam.sdk.util.FastNanoClockAndSleeper; -import org.apache.beam.sdk.util.FluentBackoff; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; -import org.checkerframework.checker.nullness.qual.Nullable; import org.junit.After; import org.junit.Before; import org.junit.Rule; @@ -165,30 +105,6 @@ public void testCreationWithDefaultOptions() { assertNotNull(pipelineOptions.getGcpCredential()); } - @Test - public void testUploadBufferSizeDefault() { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil util = pipelineOptions.getGcsUtil(); - assertNull(util.delegate.getUploadBufferSizeBytes()); - } - - @Test - public void testUploadBufferSizeUserSpecified() { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - pipelineOptions.setGcsUploadBufferSizeBytes(12345); - GcsUtil util = pipelineOptions.getGcsUtil(); - assertEquals((Integer) 12345, util.delegate.getUploadBufferSizeBytes()); - } - - @Test - public void testCreationWithExecutorServiceProvided() { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - pipelineOptions.setExecutorService(Executors.newCachedThreadPool()); - assertSame( - pipelineOptions.getExecutorService(), - pipelineOptions.getGcsUtil().delegate.executorService); - } - @Test public void testCreationWithGcsUtilProvided() { GcsOptions pipelineOptions = PipelineOptionsFactory.as(GcsOptions.class); @@ -197,35 +113,6 @@ public void testCreationWithGcsUtilProvided() { assertSame(gcsUtil, pipelineOptions.getGcsUtil()); } - @Test - public void testCreationWithExplicitGoogleCloudStorageReadOptions() throws Exception { - GoogleCloudStorageReadOptions readOptions = - GoogleCloudStorageReadOptions.builder() - .setFadvise(GoogleCloudStorageReadOptions.Fadvise.AUTO) - .setGzipEncodingSupportEnabled(true) - .setFastFailOnNotFoundEnabled(false) - .build(); - - GcsOptions pipelineOptions = PipelineOptionsFactory.as(GcsOptions.class); - pipelineOptions.setGoogleCloudStorageReadOptions(readOptions); - - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - GoogleCloudStorage googleCloudStorageMock = Mockito.spy(GoogleCloudStorage.class); - Mockito.when( - googleCloudStorageMock.open( - Mockito.any(StorageResourceId.class), - Mockito.any(GoogleCloudStorageReadOptions.class))) - .thenReturn(Mockito.mock(SeekableByteChannel.class)); - gcsUtil.delegate.setCloudStorageImpl(googleCloudStorageMock); - - assertEquals(readOptions, pipelineOptions.getGoogleCloudStorageReadOptions()); - - // Assert read options are passed to GCS calls - pipelineOptions.getGcsUtil().open(GcsPath.fromUri("gs://bucket/path")); - Mockito.verify(googleCloudStorageMock, Mockito.times(1)) - .open(StorageResourceId.fromStringPath("gs://bucket/path"), readOptions); - } - @Test public void testMultipleThreadsCanCompleteOutOfOrderWithDefaultThreadPool() throws Exception { GcsOptions pipelineOptions = PipelineOptionsFactory.as(GcsOptions.class); @@ -258,1653 +145,6 @@ public void testMultipleThreadsCanCompleteOutOfOrderWithDefaultThreadPool() thro "Expected tasks to complete", executorService.awaitTermination(10, TimeUnit.SECONDS)); } - @Test - public void testGlobExpansion() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); - Storage.Objects.Get mockStorageGet = Mockito.mock(Storage.Objects.Get.class); - Storage.Objects.List mockStorageList = Mockito.mock(Storage.Objects.List.class); - - Objects modelObjects = new Objects(); - List items = new ArrayList<>(); - // A directory - items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/")); - - // Files within the directory - items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/file1name")); - items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/file2name")); - items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/file3name")); - items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/otherfile")); - items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/anotherfile")); - - modelObjects.setItems(items); - - when(mockStorage.objects()).thenReturn(mockStorageObjects); - when(mockStorageObjects.get("testbucket", "testdirectory/otherfile")) - .thenReturn(mockStorageGet); - when(mockStorageObjects.list("testbucket")).thenReturn(mockStorageList); - when(mockStorageGet.execute()) - .thenReturn(new StorageObject().setBucket("testbucket").setName("testdirectory/otherfile")); - when(mockStorageList.execute()).thenReturn(modelObjects); - - // Test a single file. - { - GcsPath pattern = GcsPath.fromUri("gs://testbucket/testdirectory/otherfile"); - List expectedFiles = - ImmutableList.of(GcsPath.fromUri("gs://testbucket/testdirectory/otherfile")); - - assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray())); - } - - // Test patterns. - { - GcsPath pattern = GcsPath.fromUri("gs://testbucket/testdirectory/file*"); - List expectedFiles = - ImmutableList.of( - GcsPath.fromUri("gs://testbucket/testdirectory/file1name"), - GcsPath.fromUri("gs://testbucket/testdirectory/file2name"), - GcsPath.fromUri("gs://testbucket/testdirectory/file3name")); - - assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray())); - } - - { - GcsPath pattern = GcsPath.fromUri("gs://testbucket/testdirectory/file[1-3]*"); - List expectedFiles = - ImmutableList.of( - GcsPath.fromUri("gs://testbucket/testdirectory/file1name"), - GcsPath.fromUri("gs://testbucket/testdirectory/file2name"), - GcsPath.fromUri("gs://testbucket/testdirectory/file3name")); - - assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray())); - } - - { - GcsPath pattern = GcsPath.fromUri("gs://testbucket/testdirectory/file?name"); - List expectedFiles = - ImmutableList.of( - GcsPath.fromUri("gs://testbucket/testdirectory/file1name"), - GcsPath.fromUri("gs://testbucket/testdirectory/file2name"), - GcsPath.fromUri("gs://testbucket/testdirectory/file3name")); - - assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray())); - } - - { - GcsPath pattern = GcsPath.fromUri("gs://testbucket/test*ectory/fi*name"); - List expectedFiles = - ImmutableList.of( - GcsPath.fromUri("gs://testbucket/testdirectory/file1name"), - GcsPath.fromUri("gs://testbucket/testdirectory/file2name"), - GcsPath.fromUri("gs://testbucket/testdirectory/file3name")); - - assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray())); - } - } - - @Test - public void testRecursiveGlobExpansion() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); - Storage.Objects.Get mockStorageGet = Mockito.mock(Storage.Objects.Get.class); - Storage.Objects.List mockStorageList = Mockito.mock(Storage.Objects.List.class); - - Objects modelObjects = new Objects(); - List items = new ArrayList<>(); - // A directory - items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/")); - - // Files within the directory - items.add(new StorageObject().setBucket("testbucket").setName("test/directory/file1.txt")); - items.add(new StorageObject().setBucket("testbucket").setName("test/directory/file2.txt")); - items.add(new StorageObject().setBucket("testbucket").setName("test/directory/file3.txt")); - items.add(new StorageObject().setBucket("testbucket").setName("test/directory/otherfile")); - items.add(new StorageObject().setBucket("testbucket").setName("test/directory/anotherfile")); - items.add(new StorageObject().setBucket("testbucket").setName("test/file4.txt")); - - modelObjects.setItems(items); - - when(mockStorage.objects()).thenReturn(mockStorageObjects); - when(mockStorageObjects.get("testbucket", "test/directory/otherfile")) - .thenReturn(mockStorageGet); - when(mockStorageObjects.list("testbucket")).thenReturn(mockStorageList); - when(mockStorageGet.execute()) - .thenReturn( - new StorageObject().setBucket("testbucket").setName("test/directory/otherfile")); - when(mockStorageList.execute()).thenReturn(modelObjects); - - { - GcsPath pattern = GcsPath.fromUri("gs://testbucket/test/**/*.txt"); - List expectedFiles = - ImmutableList.of( - GcsPath.fromUri("gs://testbucket/test/directory/file1.txt"), - GcsPath.fromUri("gs://testbucket/test/directory/file2.txt"), - GcsPath.fromUri("gs://testbucket/test/directory/file3.txt"), - GcsPath.fromUri("gs://testbucket/test/file4.txt")); - - assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray())); - } - } - - // GCSUtil.expand() should fail when matching a single object when that object does not exist. - // We should return the empty result since GCS get object is strongly consistent. - @Test - public void testNonExistentObjectReturnsEmptyResult() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); - Storage.Objects.Get mockStorageGet = Mockito.mock(Storage.Objects.Get.class); - - GcsPath pattern = GcsPath.fromUri("gs://testbucket/testdirectory/nonexistentfile"); - GoogleJsonResponseException expectedException = - googleJsonResponseException( - HttpStatusCodes.STATUS_CODE_NOT_FOUND, "It don't exist", "Nothing here to see"); - - when(mockStorage.objects()).thenReturn(mockStorageObjects); - when(mockStorageObjects.get(pattern.getBucket(), pattern.getObject())) - .thenReturn(mockStorageGet); - when(mockStorageGet.execute()).thenThrow(expectedException); - - assertEquals(Collections.emptyList(), gcsUtil.expand(pattern)); - } - - // GCSUtil.expand() should fail for other errors such as access denied. - @Test - public void testAccessDeniedObjectThrowsIOException() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); - Storage.Objects.Get mockStorageGet = Mockito.mock(Storage.Objects.Get.class); - - GcsPath pattern = GcsPath.fromUri("gs://testbucket/testdirectory/accessdeniedfile"); - GoogleJsonResponseException expectedException = - googleJsonResponseException( - HttpStatusCodes.STATUS_CODE_FORBIDDEN, - "Waves hand mysteriously", - "These aren't the buckets you're looking for"); - - when(mockStorage.objects()).thenReturn(mockStorageObjects); - when(mockStorageObjects.get(pattern.getBucket(), pattern.getObject())) - .thenReturn(mockStorageGet); - when(mockStorageGet.execute()).thenThrow(expectedException); - - thrown.expect(IOException.class); - thrown.expectMessage("Unable to get the file object for path"); - gcsUtil.expand(pattern); - } - - @Test - public void testFileSizeNonBatch() throws Exception { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); - Storage.Objects.Get mockStorageGet = Mockito.mock(Storage.Objects.Get.class); - - when(mockStorage.objects()).thenReturn(mockStorageObjects); - when(mockStorageObjects.get("testbucket", "testobject")).thenReturn(mockStorageGet); - when(mockStorageGet.execute()) - .thenReturn(new StorageObject().setSize(BigInteger.valueOf(1000))); - - assertEquals(1000, gcsUtil.fileSize(GcsPath.fromComponents("testbucket", "testobject"))); - } - - @Test - public void testFileSizeWhenFileNotFoundNonBatch() throws Exception { - MockLowLevelHttpResponse notFoundResponse = new MockLowLevelHttpResponse(); - notFoundResponse.setContent(""); - notFoundResponse.setStatusCode(HttpStatusCodes.STATUS_CODE_NOT_FOUND); - - MockHttpTransport mockTransport = - new MockHttpTransport.Builder().setLowLevelHttpResponse(notFoundResponse).build(); - - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - gcsUtil.delegate.setStorageClient(new Storage(mockTransport, Transport.getJsonFactory(), null)); - - thrown.expect(FileNotFoundException.class); - gcsUtil.fileSize(GcsPath.fromComponents("testbucket", "testobject")); - } - - @Test - public void testRetryFileSizeNonBatch() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); - Storage.Objects.Get mockStorageGet = Mockito.mock(Storage.Objects.Get.class); - - BackOff mockBackOff = - BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.withMaxRetries(2).backoff()); - - when(mockStorage.objects()).thenReturn(mockStorageObjects); - when(mockStorageObjects.get("testbucket", "testobject")).thenReturn(mockStorageGet); - when(mockStorageGet.execute()) - .thenThrow(new SocketTimeoutException("SocketException")) - .thenThrow(new SocketTimeoutException("SocketException")) - .thenReturn(new StorageObject().setSize(BigInteger.valueOf(1000))); - - assertEquals( - 1000, - gcsUtil - .delegate - .getObject( - GcsPath.fromComponents("testbucket", "testobject"), - mockBackOff, - new FastNanoClockAndSleeper()::sleep) - .getSize() - .longValue()); - assertEquals(BackOff.STOP, mockBackOff.nextBackOffMillis()); - } - - @Test - public void testGetSizeBytesWhenFileNotFoundBatch() throws Exception { - JsonFactory jsonFactory = new GsonFactory(); - - String contentBoundary = "batch_foobarbaz"; - String contentBoundaryLine = "--" + contentBoundary; - String endOfContentBoundaryLine = "--" + contentBoundary + "--"; - - GenericJson error = new GenericJson().set("error", new GenericJson().set("code", 404)); - error.setFactory(jsonFactory); - - String content = - contentBoundaryLine - + "\n" - + "Content-Type: application/http\n" - + "\n" - + "HTTP/1.1 404 Not Found\n" - + "Content-Length: -1\n" - + "\n" - + error.toString() - + "\n" - + "\n" - + contentBoundaryLine - + "\n" - + "Content-Type: application/http\n" - + "\n" - + "HTTP/1.1 404 Not Found\n" - + "Content-Length: -1\n" - + "\n" - + error.toString() - + "\n" - + "\n" - + endOfContentBoundaryLine - + "\n"; - thrown.expect(FileNotFoundException.class); - MockLowLevelHttpResponse notFoundResponse = - new MockLowLevelHttpResponse() - .setContentType("multipart/mixed; boundary=" + contentBoundary) - .setContent(content) - .setStatusCode(HttpStatusCodes.STATUS_CODE_OK); - - MockHttpTransport mockTransport = - new MockHttpTransport.Builder().setLowLevelHttpResponse(notFoundResponse).build(); - - GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); - - gcsUtil.delegate.setStorageClient(new Storage(mockTransport, Transport.getJsonFactory(), null)); - gcsUtil.delegate.fileSizes( - ImmutableList.of( - GcsPath.fromComponents("testbucket", "testobject"), - GcsPath.fromComponents("testbucket", "testobject2"))); - } - - @Test - public void testGetSizeBytesWhenFileNotFoundNoBatch() throws Exception { - thrown.expect(FileNotFoundException.class); - MockLowLevelHttpResponse notFoundResponse = - new MockLowLevelHttpResponse() - .setContentType("text/plain") - .setContent("error") - .setStatusCode(HttpStatusCodes.STATUS_CODE_NOT_FOUND); - - MockHttpTransport mockTransport = - new MockHttpTransport.Builder().setLowLevelHttpResponse(notFoundResponse).build(); - - GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); - - gcsUtil.delegate.setStorageClient(new Storage(mockTransport, Transport.getJsonFactory(), null)); - gcsUtil.delegate.fileSizes( - ImmutableList.of(GcsPath.fromComponents("testbucket", "testobject"))); - } - - @Test - public void testGetSizeBytesWhenFileNotFoundBatchRetry() throws Exception { - JsonFactory jsonFactory = new GsonFactory(); - - String contentBoundary = "batch_foobarbaz"; - String contentBoundaryLine = "--" + contentBoundary; - String endOfContentBoundaryLine = "--" + contentBoundary + "--"; - - GenericJson error = new GenericJson().set("error", new GenericJson().set("code", 404)); - error.setFactory(jsonFactory); - - String content = - contentBoundaryLine - + "\n" - + "Content-Type: application/http\n" - + "\n" - + "HTTP/1.1 404 Not Found\n" - + "Content-Length: -1\n" - + "\n" - + error.toString() - + "\n" - + "\n" - + contentBoundaryLine - + "\n" - + "Content-Type: application/http\n" - + "\n" - + "HTTP/1.1 404 Not Found\n" - + "Content-Length: -1\n" - + "\n" - + error.toString() - + "\n" - + "\n" - + endOfContentBoundaryLine; - - thrown.expect(FileNotFoundException.class); - - final LowLevelHttpResponse[] mockResponses = - new LowLevelHttpResponse[] { - Mockito.mock(LowLevelHttpResponse.class), Mockito.mock(LowLevelHttpResponse.class), - }; - when(mockResponses[0].getContentType()).thenReturn("text/plain"); - when(mockResponses[1].getContentType()) - .thenReturn("multipart/mixed; boundary=" + contentBoundary); - - // 429: Too many requests, then 200: OK. - when(mockResponses[0].getStatusCode()).thenReturn(429); - when(mockResponses[1].getStatusCode()).thenReturn(200); - when(mockResponses[0].getContent()).thenReturn(toStream("error")); - when(mockResponses[1].getContent()).thenReturn(toStream(content)); - - // A mock transport that lets us mock the API responses. - MockHttpTransport mockTransport = - new MockHttpTransport.Builder() - .setLowLevelHttpRequest( - new MockLowLevelHttpRequest() { - int index = 0; - - @Override - public LowLevelHttpResponse execute() throws IOException { - return mockResponses[index++]; - } - }) - .build(); - - GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); - - gcsUtil.delegate.setStorageClient( - new Storage(mockTransport, Transport.getJsonFactory(), new RetryHttpRequestInitializer())); - gcsUtil.delegate.fileSizes( - ImmutableList.of( - GcsPath.fromComponents("testbucket", "testobject"), - GcsPath.fromComponents("testbucket", "testobject2"))); - } - - @Test - public void testGetSizeBytesWhenFileNotFoundNoBatchRetry() throws Exception { - thrown.expect(FileNotFoundException.class); - - final LowLevelHttpResponse[] mockResponses = - new LowLevelHttpResponse[] { - Mockito.mock(LowLevelHttpResponse.class), Mockito.mock(LowLevelHttpResponse.class), - }; - when(mockResponses[0].getContentType()).thenReturn("text/plain"); - when(mockResponses[1].getContentType()).thenReturn("text/plain"); - - // 429: Too many requests, then 200: OK. - when(mockResponses[0].getStatusCode()).thenReturn(429); - when(mockResponses[1].getStatusCode()).thenReturn(404); - when(mockResponses[0].getContent()).thenReturn(toStream("error")); - when(mockResponses[1].getContent()).thenReturn(toStream("error")); - - // A mock transport that lets us mock the API responses. - MockHttpTransport mockTransport = - new MockHttpTransport.Builder() - .setLowLevelHttpRequest( - new MockLowLevelHttpRequest() { - int index = 0; - - @Override - public LowLevelHttpResponse execute() throws IOException { - return mockResponses[index++]; - } - }) - .build(); - - GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); - - gcsUtil.delegate.setStorageClient( - new Storage(mockTransport, Transport.getJsonFactory(), new RetryHttpRequestInitializer())); - gcsUtil.delegate.fileSizes( - ImmutableList.of(GcsPath.fromComponents("testbucket", "testobject"))); - } - - @Test - public void testRemoveWhenFileNotFound() throws Exception { - JsonFactory jsonFactory = new GsonFactory(); - - String contentBoundary = "batch_foobarbaz"; - String contentBoundaryLine = "--" + contentBoundary; - String endOfContentBoundaryLine = "--" + contentBoundary + "--"; - - GenericJson error = new GenericJson().set("error", new GenericJson().set("code", 404)); - error.setFactory(jsonFactory); - - String content = - contentBoundaryLine - + "\n" - + "Content-Type: application/http\n" - + "\n" - + "HTTP/1.1 404 Not Found\n" - + "Content-Length: -1\n" - + "\n" - + error.toString() - + "\n" - + "\n" - + endOfContentBoundaryLine - + "\n"; - - final LowLevelHttpResponse mockResponse = Mockito.mock(LowLevelHttpResponse.class); - when(mockResponse.getContentType()).thenReturn("multipart/mixed; boundary=" + contentBoundary); - when(mockResponse.getStatusCode()).thenReturn(200); - when(mockResponse.getContent()).thenReturn(toStream(content)); - - // A mock transport that lets us mock the API responses. - MockLowLevelHttpRequest request = - new MockLowLevelHttpRequest() { - @Override - public LowLevelHttpResponse execute() throws IOException { - return mockResponse; - } - }; - MockHttpTransport mockTransport = - new MockHttpTransport.Builder().setLowLevelHttpRequest(request).build(); - - GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); - gcsUtil.delegate.setStorageClient( - new Storage(mockTransport, Transport.getJsonFactory(), new RetryHttpRequestInitializer())); - gcsUtil.remove(Arrays.asList("gs://some-bucket/already-deleted")); - } - - @Test - public void testCreateBucket() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - Storage.Buckets.Insert mockStorageInsert = Mockito.mock(Storage.Buckets.Insert.class); - - BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); - - when(mockStorage.buckets()).thenReturn(mockStorageObjects); - when(mockStorageObjects.insert(any(String.class), any(Bucket.class))) - .thenReturn(mockStorageInsert); - when(mockStorageInsert.execute()) - .thenThrow(new SocketTimeoutException("SocketException")) - .thenReturn(new Bucket()); - - gcsUtil.delegate.createBucket( - "a", new Bucket(), mockBackOff, new FastNanoClockAndSleeper()::sleep); - } - - @Test - public void testCreateBucketAccessErrors() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); - Storage.Buckets.Insert mockStorageInsert = Mockito.mock(Storage.Buckets.Insert.class); - - BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); - GoogleJsonResponseException expectedException = - googleJsonResponseException( - HttpStatusCodes.STATUS_CODE_FORBIDDEN, - "Waves hand mysteriously", - "These aren't the buckets you're looking for"); - - when(mockStorage.buckets()).thenReturn(mockStorageObjects); - when(mockStorageObjects.insert(any(String.class), any(Bucket.class))) - .thenReturn(mockStorageInsert); - when(mockStorageInsert.execute()).thenThrow(expectedException); - - thrown.expect(AccessDeniedException.class); - - gcsUtil.delegate.createBucket( - "a", new Bucket(), mockBackOff, new FastNanoClockAndSleeper()::sleep); - } - - @Test - public void testBucketAccessible() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); - Storage.Buckets.Get mockStorageGet = Mockito.mock(Storage.Buckets.Get.class); - - BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); - - when(mockStorage.buckets()).thenReturn(mockStorageObjects); - when(mockStorageObjects.get("testbucket")).thenReturn(mockStorageGet); - when(mockStorageGet.execute()) - .thenThrow(new SocketTimeoutException("SocketException")) - .thenReturn(new Bucket()); - - assertTrue( - gcsUtil.delegate.bucketAccessible( - GcsPath.fromComponents("testbucket", "testobject"), - mockBackOff, - new FastNanoClockAndSleeper()::sleep)); - } - - @Test - public void testBucketDoesNotExistBecauseOfAccessError() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); - Storage.Buckets.Get mockStorageGet = Mockito.mock(Storage.Buckets.Get.class); - - BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); - GoogleJsonResponseException expectedException = - googleJsonResponseException( - HttpStatusCodes.STATUS_CODE_FORBIDDEN, - "Waves hand mysteriously", - "These aren't the buckets you're looking for"); - - when(mockStorage.buckets()).thenReturn(mockStorageObjects); - when(mockStorageObjects.get("testbucket")).thenReturn(mockStorageGet); - when(mockStorageGet.execute()).thenThrow(expectedException); - - assertFalse( - gcsUtil.delegate.bucketAccessible( - GcsPath.fromComponents("testbucket", "testobject"), - mockBackOff, - new FastNanoClockAndSleeper()::sleep)); - } - - @Test - public void testBucketDoesNotExist() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); - Storage.Buckets.Get mockStorageGet = Mockito.mock(Storage.Buckets.Get.class); - - BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); - - when(mockStorage.buckets()).thenReturn(mockStorageObjects); - when(mockStorageObjects.get("testbucket")).thenReturn(mockStorageGet); - when(mockStorageGet.execute()) - .thenThrow( - googleJsonResponseException( - HttpStatusCodes.STATUS_CODE_NOT_FOUND, "It don't exist", "Nothing here to see")); - - assertFalse( - gcsUtil.delegate.bucketAccessible( - GcsPath.fromComponents("testbucket", "testobject"), - mockBackOff, - new FastNanoClockAndSleeper()::sleep)); - } - - @Test - public void testVerifyBucketAccessible() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); - Storage.Buckets.Get mockStorageGet = Mockito.mock(Storage.Buckets.Get.class); - - BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); - - when(mockStorage.buckets()).thenReturn(mockStorageObjects); - when(mockStorageObjects.get("testbucket")).thenReturn(mockStorageGet); - when(mockStorageGet.execute()) - .thenThrow(new SocketTimeoutException("SocketException")) - .thenReturn(new Bucket()); - - gcsUtil.delegate.verifyBucketAccessible( - GcsPath.fromComponents("testbucket", "testobject"), - mockBackOff, - new FastNanoClockAndSleeper()::sleep); - } - - @Test(expected = AccessDeniedException.class) - public void testVerifyBucketAccessibleAccessError() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); - Storage.Buckets.Get mockStorageGet = Mockito.mock(Storage.Buckets.Get.class); - - BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); - GoogleJsonResponseException expectedException = - googleJsonResponseException( - HttpStatusCodes.STATUS_CODE_FORBIDDEN, - "Waves hand mysteriously", - "These aren't the buckets you're looking for"); - - when(mockStorage.buckets()).thenReturn(mockStorageObjects); - when(mockStorageObjects.get("testbucket")).thenReturn(mockStorageGet); - when(mockStorageGet.execute()).thenThrow(expectedException); - - gcsUtil.delegate.verifyBucketAccessible( - GcsPath.fromComponents("testbucket", "testobject"), - mockBackOff, - new FastNanoClockAndSleeper()::sleep); - } - - @Test(expected = FileNotFoundException.class) - public void testVerifyBucketAccessibleDoesNotExist() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); - Storage.Buckets.Get mockStorageGet = Mockito.mock(Storage.Buckets.Get.class); - - BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); - - when(mockStorage.buckets()).thenReturn(mockStorageObjects); - when(mockStorageObjects.get("testbucket")).thenReturn(mockStorageGet); - when(mockStorageGet.execute()) - .thenThrow( - googleJsonResponseException( - HttpStatusCodes.STATUS_CODE_NOT_FOUND, "It don't exist", "Nothing here to see")); - - gcsUtil.delegate.verifyBucketAccessible( - GcsPath.fromComponents("testbucket", "testobject"), - mockBackOff, - new FastNanoClockAndSleeper()::sleep); - } - - @Test - public void testGetBucket() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); - Storage.Buckets.Get mockStorageGet = Mockito.mock(Storage.Buckets.Get.class); - - BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); - - when(mockStorage.buckets()).thenReturn(mockStorageObjects); - when(mockStorageObjects.get("testbucket")).thenReturn(mockStorageGet); - when(mockStorageGet.execute()) - .thenThrow(new SocketTimeoutException("SocketException")) - .thenReturn(new Bucket()); - - assertNotNull( - gcsUtil.delegate.getBucket( - GcsPath.fromComponents("testbucket", "testobject"), - mockBackOff, - new FastNanoClockAndSleeper()::sleep)); - } - - @Test - public void testGetBucketNotExists() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); - Storage.Buckets.Get mockStorageGet = Mockito.mock(Storage.Buckets.Get.class); - - BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); - - when(mockStorage.buckets()).thenReturn(mockStorageObjects); - when(mockStorageObjects.get("testbucket")).thenReturn(mockStorageGet); - when(mockStorageGet.execute()) - .thenThrow( - googleJsonResponseException( - HttpStatusCodes.STATUS_CODE_NOT_FOUND, "It don't exist", "Nothing here to see")); - - thrown.expect(FileNotFoundException.class); - thrown.expectMessage("It don't exist"); - gcsUtil.delegate.getBucket( - GcsPath.fromComponents("testbucket", "testobject"), - mockBackOff, - new FastNanoClockAndSleeper()::sleep); - } - - @Test - public void testGCSChannelCloseIdempotent() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - GoogleCloudStorageReadOptions readOptions = - GoogleCloudStorageReadOptions.builder().setFastFailOnNotFoundEnabled(false).build(); - - gcsUtil.delegate.setCloudStorageImpl( - GoogleCloudStorageOptions.builder() - .setAppName("Beam") - .setGrpcEnabled(true) - .setProjectId("my_project") - .setReadChannelOptions(readOptions) - .build()); - - SeekableByteChannel channel = gcsUtil.open(GcsPath.fromComponents("testbucket", "testobject")); - channel.close(); - channel.close(); - } - - @Test - public void testGCSReadMetricsIsSet() { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - GoogleCloudStorageReadOptions readOptions = - GoogleCloudStorageReadOptions.builder().setFastFailOnNotFoundEnabled(true).build(); - gcsUtil.delegate.setCloudStorageImpl( - GoogleCloudStorageOptions.builder() - .setAppName("Beam") - .setGrpcEnabled(true) - .setProjectId("my_project") - .setReadChannelOptions(readOptions) - .build()); - assertThrows( - IOException.class, () -> gcsUtil.open(GcsPath.fromComponents("testbucket", "testbucket"))); - verifyMetricWasSet("my_project", "testbucket", "GcsGet", "permission_denied", 1); - } - - @Test - public void testGCSWriteMetricsIsSet() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - GoogleCloudStorage mockStorage = Mockito.mock(GoogleCloudStorage.class); - gcsUtil.delegate.setCloudStorageImpl( - GoogleCloudStorageOptions.builder() - .setAppName("Beam") - .setGrpcEnabled(true) - .setProjectId("my_project") - .build()); - when(mockStorage.create( - new StorageResourceId("testbucket", "testobject"), - CreateObjectOptions.builder() - .setOverwriteExisting(true) - .setContentType("type") - .build())) - .thenThrow(IOException.class); - GcsPath gcsPath = GcsPath.fromComponents("testbucket", "testobject"); - assertThrows(IOException.class, () -> gcsUtil.create(gcsPath, "")); - verifyMetricWasSet("my_project", "testbucket", "GcsInsert", "permission_denied", 1); - } - - private void verifyMetricWasSet( - String projectId, String bucketId, String method, String status, long count) { - // Verify the metric as reported. - HashMap labels = new HashMap<>(); - labels.put(MonitoringInfoConstants.Labels.PTRANSFORM, ""); - labels.put(MonitoringInfoConstants.Labels.SERVICE, "Storage"); - labels.put(MonitoringInfoConstants.Labels.METHOD, method); - labels.put(MonitoringInfoConstants.Labels.GCS_PROJECT_ID, projectId); - labels.put(MonitoringInfoConstants.Labels.GCS_BUCKET, bucketId); - labels.put( - MonitoringInfoConstants.Labels.RESOURCE, - GcpResourceIdentifiers.cloudStorageBucket(bucketId)); - labels.put(MonitoringInfoConstants.Labels.STATUS, status); - - MonitoringInfoMetricName name = - MonitoringInfoMetricName.named(MonitoringInfoConstants.Urns.API_REQUEST_COUNT, labels); - MetricsContainerImpl container = - (MetricsContainerImpl) MetricsEnvironment.getProcessWideContainer(); - assertEquals(count, (long) container.getCounter(name).getCumulative()); - } - - /** Builds a fake GoogleJsonResponseException for testing API error handling. */ - private static GoogleJsonResponseException googleJsonResponseException( - final int status, final String reason, final String message) throws IOException { - final JsonFactory jsonFactory = new GsonFactory(); - HttpTransport transport = - new MockHttpTransport() { - @Override - public LowLevelHttpRequest buildRequest(String method, String url) throws IOException { - ErrorInfo errorInfo = new ErrorInfo(); - errorInfo.setReason(reason); - errorInfo.setMessage(message); - errorInfo.setFactory(jsonFactory); - GenericJson error = new GenericJson(); - error.set("code", status); - error.set("errors", Arrays.asList(errorInfo)); - error.setFactory(jsonFactory); - GenericJson errorResponse = new GenericJson(); - errorResponse.set("error", error); - errorResponse.setFactory(jsonFactory); - return new MockLowLevelHttpRequest() - .setResponse( - new MockLowLevelHttpResponse() - .setContent(errorResponse.toPrettyString()) - .setContentType(Json.MEDIA_TYPE) - .setStatusCode(status)); - } - }; - HttpRequest request = - transport.createRequestFactory().buildGetRequest(HttpTesting.SIMPLE_GENERIC_URL); - request.setThrowExceptionOnExecuteError(false); - HttpResponse response = request.execute(); - return GoogleJsonResponseException.from(jsonFactory, response); - } - - private static List makeStrings(String s, int n) { - return makeStrings("bucket", s, n); - } - - private static List makeStrings(String bucket, String s, int n) { - ImmutableList.Builder ret = ImmutableList.builder(); - for (int i = 0; i < n; ++i) { - ret.add(String.format("gs://%s/%s%d", bucket, s, i)); - } - return ret.build(); - } - - private static List makeGcsPaths(String s, int n) { - ImmutableList.Builder ret = ImmutableList.builder(); - for (int i = 0; i < n; ++i) { - ret.add(GcsPath.fromUri(String.format("gs://bucket/%s%d", s, i))); - } - return ret.build(); - } - - private static int sumBatchSizes(List batches) { - int ret = 0; - for (BatchInterface b : batches) { - ret += b.size(); - assertThat(b.size(), greaterThan(0)); - } - return ret; - } - - @Test - public void testMakeRewriteOps() throws IOException { - GcsOptions gcsOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = gcsOptions.getGcsUtil(); - - LinkedList rewrites = - gcsUtil.delegate.makeRewriteOps( - makeStrings("s", 1), makeStrings("d", 1), false, false, false); - assertEquals(1, rewrites.size()); - - RewriteOp rewrite = rewrites.pop(); - assertTrue(rewrite.getReadyToEnqueue()); - Storage.Objects.Rewrite request = rewrite.rewriteRequest; - assertNull(request.getMaxBytesRewrittenPerCall()); - assertEquals("bucket", request.getSourceBucket()); - assertEquals("s0", request.getSourceObject()); - assertEquals("bucket", request.getDestinationBucket()); - assertEquals("d0", request.getDestinationObject()); - } - - @Test - public void testMakeRewriteOpsWithOptions() throws IOException { - GcsOptions gcsOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = gcsOptions.getGcsUtil(); - gcsUtil.delegate.maxBytesRewrittenPerCall = 1337L; - - LinkedList rewrites = - gcsUtil.delegate.makeRewriteOps( - makeStrings("s", 1), makeStrings("d", 1), false, false, false); - assertEquals(1, rewrites.size()); - - RewriteOp rewrite = rewrites.pop(); - assertTrue(rewrite.getReadyToEnqueue()); - Storage.Objects.Rewrite request = rewrite.rewriteRequest; - assertEquals(Long.valueOf(1337L), request.getMaxBytesRewrittenPerCall()); - } - - @Test - public void testMakeRewriteBatches() throws IOException { - GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); - - // Small number of files fits in 1 batch - List batches = - gcsUtil.delegate.makeRewriteBatches( - gcsUtil.delegate.makeRewriteOps( - makeStrings("s", 3), makeStrings("d", 3), false, false, false)); - assertThat(batches.size(), equalTo(1)); - assertThat(sumBatchSizes(batches), equalTo(3)); - - // 1 batch of files fits in 1 batch - batches = - gcsUtil.delegate.makeRewriteBatches( - gcsUtil.delegate.makeRewriteOps( - makeStrings("s", 100), makeStrings("d", 100), false, false, false)); - assertThat(batches.size(), equalTo(1)); - assertThat(sumBatchSizes(batches), equalTo(100)); - - // A little more than 5 batches of files fits in 6 batches - batches = - gcsUtil.delegate.makeRewriteBatches( - gcsUtil.delegate.makeRewriteOps( - makeStrings("s", 501), makeStrings("d", 501), false, false, false)); - assertThat(batches.size(), equalTo(6)); - assertThat(sumBatchSizes(batches), equalTo(501)); - } - - @Test - public void testMakeRewriteBatchesWithLowerDataOpLimit() throws IOException { - GcsOptions options = gcsOptionsWithTestCredential(); - options.setGcsRewriteDataOpBatchLimit(2); - GcsUtil gcsUtil = options.getGcsUtil(); - - // Small number of files in same bucket fits in 1 batch - List batches = - gcsUtil.delegate.makeRewriteBatches( - gcsUtil.delegate.makeRewriteOps( - makeStrings("s", 5), makeStrings("d", 5), false, false, false)); - assertThat(batches.size(), equalTo(1)); - assertThat(sumBatchSizes(batches), equalTo(5)); - - // Files copying between buckets use smaller batch size - batches = - gcsUtil.delegate.makeRewriteBatches( - gcsUtil.delegate.makeRewriteOps( - makeStrings("bucket1", "s", 5), - makeStrings("bucket2", "d", 5), - false, - false, - false)); - assertThat(batches.size(), equalTo(3)); - assertThat(sumBatchSizes(batches), equalTo(5)); - - // A mix of same bucket and different buckets uses large batches when possible. - List fromFiles = new ArrayList<>(makeStrings("bucket1", "s", 3)); - List toFiles = new ArrayList<>(makeStrings("bucket2", "d", 3)); - fromFiles.addAll(makeStrings("t", 90)); - toFiles.addAll(makeStrings("e", 90)); - fromFiles.addAll(makeStrings("bucket3", "u", 3)); - toFiles.addAll(makeStrings("bucket4", "f", 3)); - fromFiles.addAll(makeStrings("bucket5", "v", 1)); - toFiles.addAll(makeStrings("bucket5", "g", 1)); - - batches = - gcsUtil.delegate.makeRewriteBatches( - gcsUtil.delegate.makeRewriteOps(fromFiles, toFiles, false, false, false)); - assertThat(batches.size(), equalTo(4)); - assertThat(batches.get(0).size(), equalTo(91)); - assertThat(sumBatchSizes(batches), equalTo(97)); - } - - @Test - public void testMakeRewriteOpsInvalid() throws IOException { - GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); - thrown.expect(IllegalArgumentException.class); - thrown.expectMessage("Number of source files 3"); - - gcsUtil.delegate.makeRewriteOps(makeStrings("s", 3), makeStrings("d", 1), false, false, false); - } - - private class FakeBatcher implements BatchInterface { - ArrayList> requests = new ArrayList<>(); - - @Override - public void queue(AbstractGoogleJsonClientRequest request, JsonBatchCallback cb) { - assertNotNull(request); - assertNotNull(cb); - requests.add( - () -> { - try { - try { - T result = request.execute(); - cb.onSuccess(result, null); - } catch (FileNotFoundException e) { - GoogleJsonError error = new GoogleJsonError(); - error.setCode(HttpStatusCodes.STATUS_CODE_NOT_FOUND); - cb.onFailure(error, null); - } catch (GoogleJsonResponseException e) { - cb.onFailure(e.getDetails(), null); - } catch (SocketTimeoutException e) { - System.out.println("Propagating socket exception as batch processing error"); - throw e; - } catch (Exception e) { - System.out.println("Propagating exception as server error " + e); - e.printStackTrace(); - GoogleJsonError error = new GoogleJsonError(); - error.setCode(HttpStatusCodes.STATUS_CODE_SERVER_ERROR); - cb.onFailure(error, null); - } - } catch (IOException e) { - throw new RuntimeException(e); - } - return null; - }); - } - - @Override - public void execute() throws IOException { - RuntimeException lastException = null; - for (Supplier request : requests) { - try { - request.get(); - } catch (RuntimeException e) { - lastException = e; - } - } - if (lastException != null) { - throw lastException; - } - } - - @Override - public int size() { - return requests.size(); - } - } - - @Test - public void testRename() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); - - Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); - Storage.Objects.Rewrite mockStorageRewrite = Mockito.mock(Storage.Objects.Rewrite.class); - Storage.Objects.Delete mockStorageDelete1 = Mockito.mock(Storage.Objects.Delete.class); - Storage.Objects.Delete mockStorageDelete2 = Mockito.mock(Storage.Objects.Delete.class); - - when(mockStorage.objects()).thenReturn(mockStorageObjects); - when(mockStorageObjects.rewrite("bucket", "s0", "bucket", "d0", null)) - .thenReturn(mockStorageRewrite); - when(mockStorageRewrite.execute()) - .thenThrow(new InvalidObjectException("Test exception")) - .thenReturn(new RewriteResponse().setDone(true)); - when(mockStorageObjects.delete("bucket", "s0")) - .thenReturn(mockStorageDelete1) - .thenReturn(mockStorageDelete2); - - when(mockStorageDelete1.execute()).thenThrow(new InvalidObjectException("Test exception")); - - gcsUtil.rename(makeStrings("s", 1), makeStrings("d", 1)); - verify(mockStorageRewrite, times(2)).execute(); - verify(mockStorageDelete1, times(1)).execute(); - verify(mockStorageDelete2, times(1)).execute(); - } - - @Test - public void testRenameIgnoringMissing() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); - - Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); - Storage.Objects.Rewrite mockStorageRewrite1 = Mockito.mock(Storage.Objects.Rewrite.class); - Storage.Objects.Rewrite mockStorageRewrite2 = Mockito.mock(Storage.Objects.Rewrite.class); - Storage.Objects.Delete mockStorageDelete = Mockito.mock(Storage.Objects.Delete.class); - - when(mockStorage.objects()).thenReturn(mockStorageObjects); - when(mockStorageObjects.rewrite("bucket", "s0", "bucket", "d0", null)) - .thenReturn(mockStorageRewrite1); - when(mockStorageRewrite1.execute()).thenThrow(new FileNotFoundException()); - when(mockStorageObjects.rewrite("bucket", "s1", "bucket", "d1", null)) - .thenReturn(mockStorageRewrite2); - when(mockStorageRewrite2.execute()).thenReturn(new RewriteResponse().setDone(true)); - when(mockStorageObjects.delete("bucket", "s1")).thenReturn(mockStorageDelete); - - gcsUtil.rename( - makeStrings("s", 2), makeStrings("d", 2), StandardMoveOptions.IGNORE_MISSING_FILES); - verify(mockStorageRewrite1, times(1)).execute(); - verify(mockStorageRewrite2, times(1)).execute(); - verify(mockStorageDelete, times(1)).execute(); - } - - @Test - public void testRenamePropagateMissingException() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); - - Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); - Storage.Objects.Rewrite mockStorageRewrite = Mockito.mock(Storage.Objects.Rewrite.class); - - when(mockStorage.objects()).thenReturn(mockStorageObjects); - when(mockStorageObjects.rewrite("bucket", "s0", "bucket", "d0", null)) - .thenReturn(mockStorageRewrite); - when(mockStorageRewrite.execute()).thenThrow(new FileNotFoundException()); - - assertThrows(IOException.class, () -> gcsUtil.rename(makeStrings("s", 1), makeStrings("d", 1))); - verify(mockStorageRewrite, times(1)).execute(); - } - - @Test - public void testRenameSkipDestinationExistsSameBucket() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); - - Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); - Storage.Objects.Rewrite mockStorageRewrite = Mockito.mock(Storage.Objects.Rewrite.class); - Storage.Objects.Delete mockStorageDelete = Mockito.mock(Storage.Objects.Delete.class); - - when(mockStorage.objects()).thenReturn(mockStorageObjects); - when(mockStorageObjects.rewrite("bucket", "s0", "bucket", "d0", null)) - .thenReturn(mockStorageRewrite); - when(mockStorageRewrite.execute()).thenReturn(new RewriteResponse().setDone(true)); - when(mockStorageObjects.delete("bucket", "s0")).thenReturn(mockStorageDelete); - - gcsUtil.rename( - makeStrings("s", 1), makeStrings("d", 1), StandardMoveOptions.SKIP_IF_DESTINATION_EXISTS); - verify(mockStorageRewrite, times(1)).execute(); - verify(mockStorageDelete, times(1)).execute(); - } - - @Test - public void testRenameSkipDestinationExistsDifferentBucket() throws IOException { - GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - - assertThrows( - UnsupportedOperationException.class, - () -> - gcsUtil.rename( - Collections.singletonList("gs://bucket/source"), - Collections.singletonList("gs://different_bucket/dest"), - StandardMoveOptions.SKIP_IF_DESTINATION_EXISTS)); - } - - @Test - public void testThrowRetentionPolicyNotMetErrorWhenUnequalChecksum() throws IOException { - // ./gradlew sdks:java:extensions:google-cloud-platform-core:test --tests - // org.apache.beam.sdk.extensions.gcp.util.GcsUtilTest.testHanRetentionPolicyNotMetError - GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); - - Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); - Storage.Objects.Get mockGetRequest1 = Mockito.mock(Storage.Objects.Get.class); - Storage.Objects.Get mockGetRequest2 = Mockito.mock(Storage.Objects.Get.class); - Storage.Objects.Rewrite mockStorageRewrite = Mockito.mock(Storage.Objects.Rewrite.class); - - // Gcs object to be used when checking the hash of the files during rewrite fail. - StorageObject srcObject = new StorageObject().setMd5Hash("a"); - StorageObject destObject = new StorageObject().setMd5Hash("b"); - - when(mockStorage.objects()).thenReturn(mockStorageObjects); - when(mockStorageObjects.rewrite("bucket", "s0", "bucket", "d0", null)) - .thenReturn(mockStorageRewrite); - when(mockStorageRewrite.execute()) - .thenThrow(googleJsonResponseException(403, "retentionPolicyNotMet", "Too soon")); - when(mockStorageObjects.get("bucket", "s0")).thenReturn(mockGetRequest1); - when(mockGetRequest1.execute()).thenReturn(srcObject); - when(mockStorageObjects.get("bucket", "d0")).thenReturn(mockGetRequest2); - when(mockGetRequest2.execute()).thenReturn(destObject); - - assertThrows(IOException.class, () -> gcsUtil.rename(makeStrings("s", 1), makeStrings("d", 1))); - - verify(mockStorageRewrite, times(1)).execute(); - } - - @Test - public void testIgnoreRetentionPolicyNotMetErrorWhenEqualChecksum() throws IOException { - GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); - - Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); - Storage.Objects.Get mockGetRequest = Mockito.mock(Storage.Objects.Get.class); - Storage.Objects.Rewrite mockStorageRewrite1 = Mockito.mock(Storage.Objects.Rewrite.class); - Storage.Objects.Rewrite mockStorageRewrite2 = Mockito.mock(Storage.Objects.Rewrite.class); - Storage.Objects.Delete mockStorageDelete = Mockito.mock(Storage.Objects.Delete.class); - - // Gcs object to be used when checking the hash of the files during rewrite fail. - StorageObject gcsObject = new StorageObject().setMd5Hash("a"); - - when(mockStorage.objects()).thenReturn(mockStorageObjects); - // First rewrite with retentionPolicyNotMet error. - when(mockStorageObjects.rewrite("bucket", "s0", "bucket", "d0", null)) - .thenReturn(mockStorageRewrite1); - when(mockStorageRewrite1.execute()) - .thenThrow(googleJsonResponseException(403, "retentionPolicyNotMet", "Too soon")); - when(mockStorageObjects.get(any(), any())) // to access object hash during error handling - .thenReturn(mockGetRequest); - when(mockGetRequest.execute()) - .thenReturn(gcsObject); // both source and destination will get the same hash - when(mockStorageObjects.delete("bucket", "s0")).thenReturn(mockStorageDelete); - - // Second rewrite should not be affected. - when(mockStorageObjects.rewrite("bucket", "s1", "bucket", "d1", null)) - .thenReturn(mockStorageRewrite2); - when(mockStorageRewrite2.execute()).thenReturn(new RewriteResponse().setDone(true)); - when(mockStorageObjects.delete("bucket", "s1")).thenReturn(mockStorageDelete); - - gcsUtil.rename(makeStrings("s", 2), makeStrings("d", 2)); - - verify(mockStorageRewrite1, times(1)).execute(); - verify(mockStorageRewrite2, times(1)).execute(); - verify(mockStorageDelete, times(2)).execute(); - } - - @Test - public void testMakeRemoveBatches() throws IOException { - GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); - - // Small number of files fits in 1 batch - List batches = gcsUtil.delegate.makeRemoveBatches(makeStrings("s", 3)); - assertThat(batches.size(), equalTo(1)); - assertThat(sumBatchSizes(batches), equalTo(3)); - - // 1 batch of files fits in 1 batch - batches = gcsUtil.delegate.makeRemoveBatches(makeStrings("s", 100)); - assertThat(batches.size(), equalTo(1)); - assertThat(sumBatchSizes(batches), equalTo(100)); - - // A little more than 5 batches of files fits in 6 batches - batches = gcsUtil.delegate.makeRemoveBatches(makeStrings("s", 501)); - assertThat(batches.size(), equalTo(6)); - assertThat(sumBatchSizes(batches), equalTo(501)); - } - - @Test - public void testMakeGetBatches() throws IOException { - GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); - - // Small number of files fits in 1 batch - List results = Lists.newArrayList(); - List batches = gcsUtil.delegate.makeGetBatches(makeGcsPaths("s", 3), results); - assertThat(batches.size(), equalTo(1)); - assertThat(sumBatchSizes(batches), equalTo(3)); - assertEquals(3, results.size()); - - // 1 batch of files fits in 1 batch - results = Lists.newArrayList(); - batches = gcsUtil.delegate.makeGetBatches(makeGcsPaths("s", 100), results); - assertThat(batches.size(), equalTo(1)); - assertThat(sumBatchSizes(batches), equalTo(100)); - assertEquals(100, results.size()); - - // A little more than 5 batches of files fits in 6 batches - results = Lists.newArrayList(); - batches = gcsUtil.delegate.makeGetBatches(makeGcsPaths("s", 501), results); - assertThat(batches.size(), equalTo(6)); - assertThat(sumBatchSizes(batches), equalTo(501)); - assertEquals(501, results.size()); - } - - @Test - public void testGetObjects() throws IOException { - GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); - - Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); - Storage.Objects.Get mockGetRequest = Mockito.mock(Storage.Objects.Get.class); - StorageObject object = new StorageObject(); - when(mockGetRequest.execute()).thenReturn(object); - when(mockStorageObjects.get(any(), any())).thenReturn(mockGetRequest); - when(mockStorage.objects()).thenReturn(mockStorageObjects); - - List results = gcsUtil.getObjects(makeGcsPaths("s", 1)); - - assertEquals(object, results.get(0).storageObject()); - } - - @Test - public void testGetObjectsWithException() throws IOException { - GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); - - Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); - Storage.Objects.Get mockGetRequest = Mockito.mock(Storage.Objects.Get.class); - when(mockStorage.objects()).thenReturn(mockStorageObjects); - when(mockStorageObjects.get(any(), any())).thenReturn(mockGetRequest); - when(mockGetRequest.execute()).thenThrow(new RuntimeException("fakeException")); - - thrown.expect(IOException.class); - thrown.expectMessage("Error trying to get gs://bucket/s0"); - - List results = gcsUtil.getObjects(makeGcsPaths("s", 1)); - - for (StorageObjectOrIOException result : results) { - if (null != result.ioException()) { - throw result.ioException(); - } - } - } - - @Test - public void testListObjectsException() throws IOException { - GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); - - Storage mockStorage = Mockito.mock(Storage.class); - gcsUtil.delegate.setStorageClient(mockStorage); - gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); - - Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); - when(mockStorage.objects()).thenReturn(mockStorageObjects); - Storage.Objects.List mockStorageList = Mockito.mock(Storage.Objects.List.class); - when(mockStorageObjects.list(any())).thenReturn(mockStorageList); - when(mockStorageList.execute()).thenThrow(new RuntimeException("FakeException")); - - thrown.expect(IOException.class); - thrown.expectMessage("Unable to match files in bucket testBucket"); - - gcsUtil.listObjects("testBucket", "prefix", null); - } - - public static class GcsUtilMock extends GcsUtil { - - public static GcsUtilMock createMockWithMockStorage(PipelineOptions options, byte[] readPayload) - throws IOException { - GcsUtilMock gcsUtilMock = createMock(options); - - GcsUtilV1Mock mockLegacy = GcsUtilV1Mock.createMockWithMockStorage(options, readPayload); - gcsUtilMock.delegate = mockLegacy; - - return gcsUtilMock; - } - - public static GcsUtilMock createMock(PipelineOptions options) { - return new GcsUtilMock(options); - } - - private GcsUtilMock(PipelineOptions options) { - super(options); - } - } - - public static class GcsUtilV1Mock extends GcsUtilV1 { - - public GoogleCloudStorage googleCloudStorage; - - public static GcsUtilV1Mock createMockWithMockStorage( - PipelineOptions options, byte[] readPayload) throws IOException { - GcsUtilV1Mock gcsUtilMock = createMock(options); - GoogleCloudStorage googleCloudStorageMock = Mockito.mock(GoogleCloudStorage.class); - gcsUtilMock.googleCloudStorage = googleCloudStorageMock; - // set the mock in the super object as well - gcsUtilMock.setCloudStorageImpl(gcsUtilMock.googleCloudStorage); - - if (readPayload == null) { - Mockito.when(googleCloudStorageMock.create(Mockito.any(), Mockito.any())) - .thenReturn(Channels.newChannel(new ByteArrayOutputStream())); - } else { - SeekableByteChannel seekableByteChannel = new SeekableInMemoryByteChannel(readPayload); - Mockito.when(googleCloudStorageMock.open(Mockito.any(StorageResourceId.class))) - .thenReturn(seekableByteChannel); - Mockito.when( - googleCloudStorageMock.open(Mockito.any(StorageResourceId.class), Mockito.any())) - .thenReturn(seekableByteChannel); - } - return gcsUtilMock; - } - - public static GcsUtilV1Mock createMock(PipelineOptions options) { - GcsOptions gcsOptions = options.as(GcsOptions.class); - Storage.Builder storageBuilder = Transport.newStorageClient(gcsOptions); - return new GcsUtilV1Mock( - storageBuilder.build(), - storageBuilder.getHttpRequestInitializer(), - gcsOptions.getExecutorService(), - hasExperiment(options, "use_grpc_for_gcs"), - gcsOptions.getGcpCredential(), - gcsOptions.getGcsUploadBufferSizeBytes(), - gcsOptions.getGcsRewriteDataOpBatchLimit(), - GcsUtilV1.GcsCountersOptions.create( - gcsOptions.getEnableBucketReadMetricCounter() - ? gcsOptions.getGcsReadCounterPrefix() - : null, - gcsOptions.getEnableBucketWriteMetricCounter() - ? gcsOptions.getGcsWriteCounterPrefix() - : null, - Boolean.TRUE.equals(gcsOptions.getGcsPerformanceMetrics())), - gcsOptions.getGoogleCloudStorageReadOptions()); - } - - private GcsUtilV1Mock( - Storage storageClient, - HttpRequestInitializer httpRequestInitializer, - ExecutorService executorService, - Boolean shouldUseGrpc, - Credentials credentials, - @Nullable Integer uploadBufferSizeBytes, - @Nullable Integer rewriteDataOpBatchLimit, - GcsUtilV1.GcsCountersOptions gcsCountersOptions, - GoogleCloudStorageReadOptions gcsReadOptions) { - super( - storageClient, - httpRequestInitializer, - executorService, - shouldUseGrpc, - credentials, - uploadBufferSizeBytes, - rewriteDataOpBatchLimit, - gcsCountersOptions, - gcsReadOptions); - } - - @Override - GoogleCloudStorage createGoogleCloudStorage( - GoogleCloudStorageOptions options, - Storage storage, - Credentials credentials, - @Nullable HttpRequestInitializer httpRequestInitializer) { - return googleCloudStorage; - } - } - - @Test - public void testCreate() throws IOException { - GcsOptions gcsOptions = gcsOptionsWithTestCredential(); - - GcsUtilMock gcsUtil = GcsUtilMock.createMock(gcsOptions); - - GoogleCloudStorage mockStorage = Mockito.mock(GoogleCloudStorage.class); - WritableByteChannel mockChannel = Mockito.mock(WritableByteChannel.class); - - GcsUtilV1Mock mockLegacy = GcsUtilV1Mock.createMock(gcsOptions); - mockLegacy.googleCloudStorage = mockStorage; - gcsUtil.delegate = mockLegacy; - - when(mockStorage.create(any(), any())).thenReturn(mockChannel); - - GcsPath path = GcsPath.fromUri("gs://testbucket/testdirectory/otherfile"); - CreateOptions createOptions = CreateOptions.builder().build(); - - assertEquals(mockChannel, gcsUtil.create(path, createOptions)); - } - - @Test - public void testCreateWithException() throws IOException { - GcsOptions gcsOptions = gcsOptionsWithTestCredential(); - - GcsUtilMock gcsUtil = GcsUtilMock.createMock(gcsOptions); - - GoogleCloudStorage mockStorage = Mockito.mock(GoogleCloudStorage.class); - - GcsUtilV1Mock mockLegacy = GcsUtilV1Mock.createMock(gcsOptions); - mockLegacy.googleCloudStorage = mockStorage; - gcsUtil.delegate = mockLegacy; - - when(mockStorage.create(any(), any())).thenThrow(new RuntimeException("testException")); - - GcsPath path = GcsPath.fromUri("gs://testbucket/testdirectory/otherfile"); - CreateOptions createOptions = CreateOptions.builder().build(); - - thrown.expect(RuntimeException.class); - thrown.expectMessage("testException"); - - gcsUtil.create(path, createOptions); - } - - private void testWriteMetrics(boolean enabled) throws IOException { - // arrange - GcsOptions gcsOptions = PipelineOptionsFactory.create().as(GcsOptions.class); - gcsOptions.setEnableBucketWriteMetricCounter(enabled); - gcsOptions.setGcsWriteCounterPrefix("test_counter"); - GcsUtilMock gcsUtil = GcsUtilMock.createMockWithMockStorage(gcsOptions, null); - byte[] payload = "some_bytes".getBytes(StandardCharsets.UTF_8); - String bucketName = "some_bucket"; - - // act - try (WritableByteChannel byteChannel = - gcsUtil.create(new GcsPath(null, bucketName, "o1"), CreateOptions.builder().build())) { - int bytesWrittenReportedByChannel = byteChannel.write(ByteBuffer.wrap(payload)); - long bytesWrittenReportedByMetric = - testMetricsContainer - .getCounter( - MetricName.named( - GcsUtil.class, - String.format("%s_%s", gcsOptions.getGcsWriteCounterPrefix(), bucketName))) - .getCumulative(); - - // assert - assertEquals(payload.length, bytesWrittenReportedByChannel); - assertEquals(enabled ? payload.length : 0, bytesWrittenReportedByMetric); - } - } - - private void testReadMetrics(boolean enabled, GoogleCloudStorageReadOptions readOptions) - throws IOException { - // arrange - GcsOptions gcsOptions = PipelineOptionsFactory.create().as(GcsOptions.class); - gcsOptions.setEnableBucketReadMetricCounter(enabled); - gcsOptions.setGcsReadCounterPrefix("test_counter"); - if (readOptions != null) { - gcsOptions.setGoogleCloudStorageReadOptions(readOptions); - } - byte[] payload = "some_bytes".getBytes(StandardCharsets.UTF_8); - GcsUtilMock gcsUtil = GcsUtilMock.createMockWithMockStorage(gcsOptions, payload); - String bucketName = "some_bucket"; - GcsPath gcsPath = new GcsPath(null, bucketName, "o1"); - // act - try (SeekableByteChannel byteChannel = gcsUtil.open(gcsPath)) { - int bytesReadReportedByChannel = byteChannel.read(ByteBuffer.allocate(payload.length)); - long bytesReadReportedByMetric = - testMetricsContainer - .getCounter( - MetricName.named( - GcsUtil.class, - String.format("%s_%s", gcsOptions.getGcsReadCounterPrefix(), bucketName))) - .getCumulative(); - - // assert - assertEquals(payload.length, bytesReadReportedByChannel); - assertEquals(enabled ? payload.length : 0, bytesReadReportedByMetric); - } - } - - @Test - public void testWriteMetricsAreCorrectlyReportedWhenEnabled() throws Exception { - testWriteMetrics(true); - } - - @Test - public void testWriteMetricsAreNotCollectedWhenNotEnabled() throws Exception { - testWriteMetrics(false); - } - - @Test - public void testReadMetricsAreCorrectlyReportedWhenEnabled() throws Exception { - testReadMetrics(true, null); - } - - @Test - public void testReadMetricsAreNotCollectedWhenNotEnabled() throws Exception { - testReadMetrics(false, null); - } - - @Test - public void testReadMetricsAreCorrectlyReportedWhenEnabledOpenWithOptions() throws Exception { - testReadMetrics(true, GoogleCloudStorageReadOptions.DEFAULT); - } - - @Test - public void testReadMetricsAreNotCollectedWhenNotEnabledOpenWithOptions() throws Exception { - testReadMetrics(false, GoogleCloudStorageReadOptions.DEFAULT); - } - - @Test - public void testGcsEndpoint() throws IOException { - GcsOptions pipelineOptions = PipelineOptionsFactory.as(GcsOptions.class); - pipelineOptions.setGcsEndpoint("http://localhost:4443/storage/v1/"); - - GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); - GoogleCloudStorageImpl gcsImpl = - (GoogleCloudStorageImpl) gcsUtil.delegate.getGoogleCloudStorage(); - assertEquals("http://localhost:4443/", gcsImpl.getOptions().getStorageRootUrl()); - assertEquals("storage/v1/", gcsImpl.getOptions().getStorageServicePath()); - } - - /** A helper to wrap a {@link GenericJson} object in a content stream. */ - private static InputStream toStream(String content) throws IOException { - return new ByteArrayInputStream(content.getBytes(StandardCharsets.UTF_8)); - } - // The tests below cover the routing that the GcsUtil facade performs for the legacy typed // methods once the use_gcsutil_v2 experiment installs a GcsUtilV2 delegate. Both delegates are // mocked, so they assert both that V2 is used and that V1 is left alone. @@ -2366,9 +606,30 @@ private GcsUtil gcsUtilWithV2Storage(com.google.cloud.storage.Storage storage) { return gcsUtil; } + private void verifyMetricWasSet( + String projectId, String bucketId, String method, String status, long count) { + // Verify the metric as reported. + HashMap labels = new HashMap<>(); + labels.put(MonitoringInfoConstants.Labels.PTRANSFORM, ""); + labels.put(MonitoringInfoConstants.Labels.SERVICE, "Storage"); + labels.put(MonitoringInfoConstants.Labels.METHOD, method); + labels.put(MonitoringInfoConstants.Labels.GCS_PROJECT_ID, projectId); + labels.put(MonitoringInfoConstants.Labels.GCS_BUCKET, bucketId); + labels.put( + MonitoringInfoConstants.Labels.RESOURCE, + GcpResourceIdentifiers.cloudStorageBucket(bucketId)); + labels.put(MonitoringInfoConstants.Labels.STATUS, status); + + MonitoringInfoMetricName name = + MonitoringInfoMetricName.named(MonitoringInfoConstants.Urns.API_REQUEST_COUNT, labels); + MetricsContainerImpl container = + (MetricsContainerImpl) MetricsEnvironment.getProcessWideContainer(); + assertEquals(count, (long) container.getCounter(name).getCumulative()); + } + /** - * Mirrors {@link #testGCSReadMetricsIsSet} for V2: opening a missing object records a {@code - * not_found} request, even though V2 detects it through a lookup rather than a {@link + * Mirrors {@link GcsUtilV1Test#testGCSReadMetricsIsSet} for V2: opening a missing object records + * a {@code not_found} request, even though V2 detects it through a lookup rather than a {@link * StorageException}. */ @Test @@ -2408,7 +669,7 @@ public void testV2WriteChannelCloseTranslatesStorageException() throws IOExcepti assertSame(preconditionFailed, thrown.getCause()); } - /** Mirrors {@link #testBucketDoesNotExist} for V2. */ + /** Mirrors {@link GcsUtilV1Test#testBucketDoesNotExist} for V2. */ @Test public void testV2BucketAccessibleIsFalseWhenBucketDoesNotExist() throws IOException { com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); @@ -2418,7 +679,7 @@ public void testV2BucketAccessibleIsFalseWhenBucketDoesNotExist() throws IOExcep assertFalse(gcsUtil.bucketAccessible(GcsPath.fromComponents("testbucket", "testobject"))); } - /** Mirrors {@link #testBucketDoesNotExistBecauseOfAccessError} for V2. */ + /** Mirrors {@link GcsUtilV1Test#testBucketDoesNotExistBecauseOfAccessError} for V2. */ @Test public void testV2BucketAccessibleIsFalseWhenAccessIsDenied() throws IOException { com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1Test.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1Test.java new file mode 100644 index 000000000000..72d554da7528 --- /dev/null +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV1Test.java @@ -0,0 +1,1850 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.extensions.gcp.util; + +import static org.apache.beam.sdk.options.ExperimentalOptions.hasExperiment; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.contains; +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.greaterThan; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.api.client.googleapis.batch.json.JsonBatchCallback; +import com.google.api.client.googleapis.json.GoogleJsonError; +import com.google.api.client.googleapis.json.GoogleJsonError.ErrorInfo; +import com.google.api.client.googleapis.json.GoogleJsonResponseException; +import com.google.api.client.googleapis.services.json.AbstractGoogleJsonClientRequest; +import com.google.api.client.http.HttpRequest; +import com.google.api.client.http.HttpRequestInitializer; +import com.google.api.client.http.HttpResponse; +import com.google.api.client.http.HttpStatusCodes; +import com.google.api.client.http.HttpTransport; +import com.google.api.client.http.LowLevelHttpRequest; +import com.google.api.client.http.LowLevelHttpResponse; +import com.google.api.client.json.GenericJson; +import com.google.api.client.json.Json; +import com.google.api.client.json.JsonFactory; +import com.google.api.client.json.gson.GsonFactory; +import com.google.api.client.testing.http.HttpTesting; +import com.google.api.client.testing.http.MockHttpTransport; +import com.google.api.client.testing.http.MockLowLevelHttpRequest; +import com.google.api.client.testing.http.MockLowLevelHttpResponse; +import com.google.api.client.util.BackOff; +import com.google.api.services.storage.Storage; +import com.google.api.services.storage.model.Bucket; +import com.google.api.services.storage.model.Objects; +import com.google.api.services.storage.model.RewriteResponse; +import com.google.api.services.storage.model.StorageObject; +import com.google.auth.Credentials; +import com.google.cloud.hadoop.gcsio.CreateObjectOptions; +import com.google.cloud.hadoop.gcsio.GoogleCloudStorage; +import com.google.cloud.hadoop.gcsio.GoogleCloudStorageImpl; +import com.google.cloud.hadoop.gcsio.GoogleCloudStorageOptions; +import com.google.cloud.hadoop.gcsio.GoogleCloudStorageReadOptions; +import com.google.cloud.hadoop.gcsio.StorageResourceId; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.FileNotFoundException; +import java.io.IOException; +import java.io.InputStream; +import java.io.InvalidObjectException; +import java.math.BigInteger; +import java.net.SocketTimeoutException; +import java.nio.ByteBuffer; +import java.nio.channels.Channels; +import java.nio.channels.SeekableByteChannel; +import java.nio.channels.WritableByteChannel; +import java.nio.charset.StandardCharsets; +import java.nio.file.AccessDeniedException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedList; +import java.util.List; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.function.Supplier; +import org.apache.beam.repackaged.core.org.apache.commons.compress.utils.SeekableInMemoryByteChannel; +import org.apache.beam.runners.core.metrics.GcpResourceIdentifiers; +import org.apache.beam.runners.core.metrics.MetricsContainerImpl; +import org.apache.beam.runners.core.metrics.MonitoringInfoConstants; +import org.apache.beam.runners.core.metrics.MonitoringInfoMetricName; +import org.apache.beam.sdk.extensions.gcp.auth.TestCredential; +import org.apache.beam.sdk.extensions.gcp.options.GcsOptions; +import org.apache.beam.sdk.extensions.gcp.util.GcsUtil.CreateOptions; +import org.apache.beam.sdk.extensions.gcp.util.GcsUtil.StorageObjectOrIOException; +import org.apache.beam.sdk.extensions.gcp.util.GcsUtilV1.BatchInterface; +import org.apache.beam.sdk.extensions.gcp.util.GcsUtilV1.RewriteOp; +import org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath; +import org.apache.beam.sdk.io.fs.MoveOptions.StandardMoveOptions; +import org.apache.beam.sdk.metrics.MetricName; +import org.apache.beam.sdk.metrics.MetricsEnvironment; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.util.FastNanoClockAndSleeper; +import org.apache.beam.sdk.util.FluentBackoff; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.junit.After; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.Mockito; + +/** Tests for {@link GcsUtilV1}, mostly through the {@link GcsUtil} facade in its default mode. */ +@RunWith(JUnit4.class) +public class GcsUtilV1Test { + @Rule public ExpectedException thrown = ExpectedException.none(); + MetricsContainerImpl testMetricsContainer; + + @Before + public void setUp() { + // Setup the ProcessWideContainer for testing metrics are set. + testMetricsContainer = new MetricsContainerImpl(null); + MetricsEnvironment.setProcessWideContainer(testMetricsContainer); + MetricsEnvironment.setCurrentContainer(testMetricsContainer); + } + + @After + public void tearDown() { + // Don't leak the containers installed by setUp into later tests in the same JVM. + MetricsEnvironment.setProcessWideContainer(null); + MetricsEnvironment.setCurrentContainer(null); + } + + private static GcsOptions gcsOptionsWithTestCredential() { + GcsOptions pipelineOptions = PipelineOptionsFactory.as(GcsOptions.class); + pipelineOptions.setGcpCredential(new TestCredential()); + return pipelineOptions; + } + + @Test + public void testUploadBufferSizeDefault() { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil util = pipelineOptions.getGcsUtil(); + assertNull(util.delegate.getUploadBufferSizeBytes()); + } + + @Test + public void testUploadBufferSizeUserSpecified() { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + pipelineOptions.setGcsUploadBufferSizeBytes(12345); + GcsUtil util = pipelineOptions.getGcsUtil(); + assertEquals((Integer) 12345, util.delegate.getUploadBufferSizeBytes()); + } + + @Test + public void testCreationWithExecutorServiceProvided() { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + pipelineOptions.setExecutorService(Executors.newCachedThreadPool()); + assertSame( + pipelineOptions.getExecutorService(), + pipelineOptions.getGcsUtil().delegate.executorService); + } + + @Test + public void testCreationWithExplicitGoogleCloudStorageReadOptions() throws Exception { + GoogleCloudStorageReadOptions readOptions = + GoogleCloudStorageReadOptions.builder() + .setFadvise(GoogleCloudStorageReadOptions.Fadvise.AUTO) + .setGzipEncodingSupportEnabled(true) + .setFastFailOnNotFoundEnabled(false) + .build(); + + GcsOptions pipelineOptions = PipelineOptionsFactory.as(GcsOptions.class); + pipelineOptions.setGoogleCloudStorageReadOptions(readOptions); + + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + GoogleCloudStorage googleCloudStorageMock = Mockito.spy(GoogleCloudStorage.class); + Mockito.when( + googleCloudStorageMock.open( + Mockito.any(StorageResourceId.class), + Mockito.any(GoogleCloudStorageReadOptions.class))) + .thenReturn(Mockito.mock(SeekableByteChannel.class)); + gcsUtil.delegate.setCloudStorageImpl(googleCloudStorageMock); + + assertEquals(readOptions, pipelineOptions.getGoogleCloudStorageReadOptions()); + + // Assert read options are passed to GCS calls + pipelineOptions.getGcsUtil().open(GcsPath.fromUri("gs://bucket/path")); + Mockito.verify(googleCloudStorageMock, Mockito.times(1)) + .open(StorageResourceId.fromStringPath("gs://bucket/path"), readOptions); + } + + @Test + public void testGlobExpansion() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); + Storage.Objects.Get mockStorageGet = Mockito.mock(Storage.Objects.Get.class); + Storage.Objects.List mockStorageList = Mockito.mock(Storage.Objects.List.class); + + Objects modelObjects = new Objects(); + List items = new ArrayList<>(); + // A directory + items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/")); + + // Files within the directory + items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/file1name")); + items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/file2name")); + items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/file3name")); + items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/otherfile")); + items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/anotherfile")); + + modelObjects.setItems(items); + + when(mockStorage.objects()).thenReturn(mockStorageObjects); + when(mockStorageObjects.get("testbucket", "testdirectory/otherfile")) + .thenReturn(mockStorageGet); + when(mockStorageObjects.list("testbucket")).thenReturn(mockStorageList); + when(mockStorageGet.execute()) + .thenReturn(new StorageObject().setBucket("testbucket").setName("testdirectory/otherfile")); + when(mockStorageList.execute()).thenReturn(modelObjects); + + // Test a single file. + { + GcsPath pattern = GcsPath.fromUri("gs://testbucket/testdirectory/otherfile"); + List expectedFiles = + ImmutableList.of(GcsPath.fromUri("gs://testbucket/testdirectory/otherfile")); + + assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray())); + } + + // Test patterns. + { + GcsPath pattern = GcsPath.fromUri("gs://testbucket/testdirectory/file*"); + List expectedFiles = + ImmutableList.of( + GcsPath.fromUri("gs://testbucket/testdirectory/file1name"), + GcsPath.fromUri("gs://testbucket/testdirectory/file2name"), + GcsPath.fromUri("gs://testbucket/testdirectory/file3name")); + + assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray())); + } + + { + GcsPath pattern = GcsPath.fromUri("gs://testbucket/testdirectory/file[1-3]*"); + List expectedFiles = + ImmutableList.of( + GcsPath.fromUri("gs://testbucket/testdirectory/file1name"), + GcsPath.fromUri("gs://testbucket/testdirectory/file2name"), + GcsPath.fromUri("gs://testbucket/testdirectory/file3name")); + + assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray())); + } + + { + GcsPath pattern = GcsPath.fromUri("gs://testbucket/testdirectory/file?name"); + List expectedFiles = + ImmutableList.of( + GcsPath.fromUri("gs://testbucket/testdirectory/file1name"), + GcsPath.fromUri("gs://testbucket/testdirectory/file2name"), + GcsPath.fromUri("gs://testbucket/testdirectory/file3name")); + + assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray())); + } + + { + GcsPath pattern = GcsPath.fromUri("gs://testbucket/test*ectory/fi*name"); + List expectedFiles = + ImmutableList.of( + GcsPath.fromUri("gs://testbucket/testdirectory/file1name"), + GcsPath.fromUri("gs://testbucket/testdirectory/file2name"), + GcsPath.fromUri("gs://testbucket/testdirectory/file3name")); + + assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray())); + } + } + + @Test + public void testRecursiveGlobExpansion() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); + Storage.Objects.Get mockStorageGet = Mockito.mock(Storage.Objects.Get.class); + Storage.Objects.List mockStorageList = Mockito.mock(Storage.Objects.List.class); + + Objects modelObjects = new Objects(); + List items = new ArrayList<>(); + // A directory + items.add(new StorageObject().setBucket("testbucket").setName("testdirectory/")); + + // Files within the directory + items.add(new StorageObject().setBucket("testbucket").setName("test/directory/file1.txt")); + items.add(new StorageObject().setBucket("testbucket").setName("test/directory/file2.txt")); + items.add(new StorageObject().setBucket("testbucket").setName("test/directory/file3.txt")); + items.add(new StorageObject().setBucket("testbucket").setName("test/directory/otherfile")); + items.add(new StorageObject().setBucket("testbucket").setName("test/directory/anotherfile")); + items.add(new StorageObject().setBucket("testbucket").setName("test/file4.txt")); + + modelObjects.setItems(items); + + when(mockStorage.objects()).thenReturn(mockStorageObjects); + when(mockStorageObjects.get("testbucket", "test/directory/otherfile")) + .thenReturn(mockStorageGet); + when(mockStorageObjects.list("testbucket")).thenReturn(mockStorageList); + when(mockStorageGet.execute()) + .thenReturn( + new StorageObject().setBucket("testbucket").setName("test/directory/otherfile")); + when(mockStorageList.execute()).thenReturn(modelObjects); + + { + GcsPath pattern = GcsPath.fromUri("gs://testbucket/test/**/*.txt"); + List expectedFiles = + ImmutableList.of( + GcsPath.fromUri("gs://testbucket/test/directory/file1.txt"), + GcsPath.fromUri("gs://testbucket/test/directory/file2.txt"), + GcsPath.fromUri("gs://testbucket/test/directory/file3.txt"), + GcsPath.fromUri("gs://testbucket/test/file4.txt")); + + assertThat(expectedFiles, contains(gcsUtil.expand(pattern).toArray())); + } + } + + // GCSUtil.expand() should fail when matching a single object when that object does not exist. + // We should return the empty result since GCS get object is strongly consistent. + @Test + public void testNonExistentObjectReturnsEmptyResult() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); + Storage.Objects.Get mockStorageGet = Mockito.mock(Storage.Objects.Get.class); + + GcsPath pattern = GcsPath.fromUri("gs://testbucket/testdirectory/nonexistentfile"); + GoogleJsonResponseException expectedException = + googleJsonResponseException( + HttpStatusCodes.STATUS_CODE_NOT_FOUND, "It don't exist", "Nothing here to see"); + + when(mockStorage.objects()).thenReturn(mockStorageObjects); + when(mockStorageObjects.get(pattern.getBucket(), pattern.getObject())) + .thenReturn(mockStorageGet); + when(mockStorageGet.execute()).thenThrow(expectedException); + + assertEquals(Collections.emptyList(), gcsUtil.expand(pattern)); + } + + // GCSUtil.expand() should fail for other errors such as access denied. + @Test + public void testAccessDeniedObjectThrowsIOException() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); + Storage.Objects.Get mockStorageGet = Mockito.mock(Storage.Objects.Get.class); + + GcsPath pattern = GcsPath.fromUri("gs://testbucket/testdirectory/accessdeniedfile"); + GoogleJsonResponseException expectedException = + googleJsonResponseException( + HttpStatusCodes.STATUS_CODE_FORBIDDEN, + "Waves hand mysteriously", + "These aren't the buckets you're looking for"); + + when(mockStorage.objects()).thenReturn(mockStorageObjects); + when(mockStorageObjects.get(pattern.getBucket(), pattern.getObject())) + .thenReturn(mockStorageGet); + when(mockStorageGet.execute()).thenThrow(expectedException); + + thrown.expect(IOException.class); + thrown.expectMessage("Unable to get the file object for path"); + gcsUtil.expand(pattern); + } + + @Test + public void testFileSizeNonBatch() throws Exception { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); + Storage.Objects.Get mockStorageGet = Mockito.mock(Storage.Objects.Get.class); + + when(mockStorage.objects()).thenReturn(mockStorageObjects); + when(mockStorageObjects.get("testbucket", "testobject")).thenReturn(mockStorageGet); + when(mockStorageGet.execute()) + .thenReturn(new StorageObject().setSize(BigInteger.valueOf(1000))); + + assertEquals(1000, gcsUtil.fileSize(GcsPath.fromComponents("testbucket", "testobject"))); + } + + @Test + public void testFileSizeWhenFileNotFoundNonBatch() throws Exception { + MockLowLevelHttpResponse notFoundResponse = new MockLowLevelHttpResponse(); + notFoundResponse.setContent(""); + notFoundResponse.setStatusCode(HttpStatusCodes.STATUS_CODE_NOT_FOUND); + + MockHttpTransport mockTransport = + new MockHttpTransport.Builder().setLowLevelHttpResponse(notFoundResponse).build(); + + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + gcsUtil.delegate.setStorageClient(new Storage(mockTransport, Transport.getJsonFactory(), null)); + + thrown.expect(FileNotFoundException.class); + gcsUtil.fileSize(GcsPath.fromComponents("testbucket", "testobject")); + } + + @Test + public void testRetryFileSizeNonBatch() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); + Storage.Objects.Get mockStorageGet = Mockito.mock(Storage.Objects.Get.class); + + BackOff mockBackOff = + BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.withMaxRetries(2).backoff()); + + when(mockStorage.objects()).thenReturn(mockStorageObjects); + when(mockStorageObjects.get("testbucket", "testobject")).thenReturn(mockStorageGet); + when(mockStorageGet.execute()) + .thenThrow(new SocketTimeoutException("SocketException")) + .thenThrow(new SocketTimeoutException("SocketException")) + .thenReturn(new StorageObject().setSize(BigInteger.valueOf(1000))); + + assertEquals( + 1000, + gcsUtil + .delegate + .getObject( + GcsPath.fromComponents("testbucket", "testobject"), + mockBackOff, + new FastNanoClockAndSleeper()::sleep) + .getSize() + .longValue()); + assertEquals(BackOff.STOP, mockBackOff.nextBackOffMillis()); + } + + @Test + public void testGetSizeBytesWhenFileNotFoundBatch() throws Exception { + JsonFactory jsonFactory = new GsonFactory(); + + String contentBoundary = "batch_foobarbaz"; + String contentBoundaryLine = "--" + contentBoundary; + String endOfContentBoundaryLine = "--" + contentBoundary + "--"; + + GenericJson error = new GenericJson().set("error", new GenericJson().set("code", 404)); + error.setFactory(jsonFactory); + + String content = + contentBoundaryLine + + "\n" + + "Content-Type: application/http\n" + + "\n" + + "HTTP/1.1 404 Not Found\n" + + "Content-Length: -1\n" + + "\n" + + error.toString() + + "\n" + + "\n" + + contentBoundaryLine + + "\n" + + "Content-Type: application/http\n" + + "\n" + + "HTTP/1.1 404 Not Found\n" + + "Content-Length: -1\n" + + "\n" + + error.toString() + + "\n" + + "\n" + + endOfContentBoundaryLine + + "\n"; + thrown.expect(FileNotFoundException.class); + MockLowLevelHttpResponse notFoundResponse = + new MockLowLevelHttpResponse() + .setContentType("multipart/mixed; boundary=" + contentBoundary) + .setContent(content) + .setStatusCode(HttpStatusCodes.STATUS_CODE_OK); + + MockHttpTransport mockTransport = + new MockHttpTransport.Builder().setLowLevelHttpResponse(notFoundResponse).build(); + + GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); + + gcsUtil.delegate.setStorageClient(new Storage(mockTransport, Transport.getJsonFactory(), null)); + gcsUtil.delegate.fileSizes( + ImmutableList.of( + GcsPath.fromComponents("testbucket", "testobject"), + GcsPath.fromComponents("testbucket", "testobject2"))); + } + + @Test + public void testGetSizeBytesWhenFileNotFoundNoBatch() throws Exception { + thrown.expect(FileNotFoundException.class); + MockLowLevelHttpResponse notFoundResponse = + new MockLowLevelHttpResponse() + .setContentType("text/plain") + .setContent("error") + .setStatusCode(HttpStatusCodes.STATUS_CODE_NOT_FOUND); + + MockHttpTransport mockTransport = + new MockHttpTransport.Builder().setLowLevelHttpResponse(notFoundResponse).build(); + + GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); + + gcsUtil.delegate.setStorageClient(new Storage(mockTransport, Transport.getJsonFactory(), null)); + gcsUtil.delegate.fileSizes( + ImmutableList.of(GcsPath.fromComponents("testbucket", "testobject"))); + } + + @Test + public void testGetSizeBytesWhenFileNotFoundBatchRetry() throws Exception { + JsonFactory jsonFactory = new GsonFactory(); + + String contentBoundary = "batch_foobarbaz"; + String contentBoundaryLine = "--" + contentBoundary; + String endOfContentBoundaryLine = "--" + contentBoundary + "--"; + + GenericJson error = new GenericJson().set("error", new GenericJson().set("code", 404)); + error.setFactory(jsonFactory); + + String content = + contentBoundaryLine + + "\n" + + "Content-Type: application/http\n" + + "\n" + + "HTTP/1.1 404 Not Found\n" + + "Content-Length: -1\n" + + "\n" + + error.toString() + + "\n" + + "\n" + + contentBoundaryLine + + "\n" + + "Content-Type: application/http\n" + + "\n" + + "HTTP/1.1 404 Not Found\n" + + "Content-Length: -1\n" + + "\n" + + error.toString() + + "\n" + + "\n" + + endOfContentBoundaryLine; + + thrown.expect(FileNotFoundException.class); + + final LowLevelHttpResponse[] mockResponses = + new LowLevelHttpResponse[] { + Mockito.mock(LowLevelHttpResponse.class), Mockito.mock(LowLevelHttpResponse.class), + }; + when(mockResponses[0].getContentType()).thenReturn("text/plain"); + when(mockResponses[1].getContentType()) + .thenReturn("multipart/mixed; boundary=" + contentBoundary); + + // 429: Too many requests, then 200: OK. + when(mockResponses[0].getStatusCode()).thenReturn(429); + when(mockResponses[1].getStatusCode()).thenReturn(200); + when(mockResponses[0].getContent()).thenReturn(toStream("error")); + when(mockResponses[1].getContent()).thenReturn(toStream(content)); + + // A mock transport that lets us mock the API responses. + MockHttpTransport mockTransport = + new MockHttpTransport.Builder() + .setLowLevelHttpRequest( + new MockLowLevelHttpRequest() { + int index = 0; + + @Override + public LowLevelHttpResponse execute() throws IOException { + return mockResponses[index++]; + } + }) + .build(); + + GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); + + gcsUtil.delegate.setStorageClient( + new Storage(mockTransport, Transport.getJsonFactory(), new RetryHttpRequestInitializer())); + gcsUtil.delegate.fileSizes( + ImmutableList.of( + GcsPath.fromComponents("testbucket", "testobject"), + GcsPath.fromComponents("testbucket", "testobject2"))); + } + + @Test + public void testGetSizeBytesWhenFileNotFoundNoBatchRetry() throws Exception { + thrown.expect(FileNotFoundException.class); + + final LowLevelHttpResponse[] mockResponses = + new LowLevelHttpResponse[] { + Mockito.mock(LowLevelHttpResponse.class), Mockito.mock(LowLevelHttpResponse.class), + }; + when(mockResponses[0].getContentType()).thenReturn("text/plain"); + when(mockResponses[1].getContentType()).thenReturn("text/plain"); + + // 429: Too many requests, then 200: OK. + when(mockResponses[0].getStatusCode()).thenReturn(429); + when(mockResponses[1].getStatusCode()).thenReturn(404); + when(mockResponses[0].getContent()).thenReturn(toStream("error")); + when(mockResponses[1].getContent()).thenReturn(toStream("error")); + + // A mock transport that lets us mock the API responses. + MockHttpTransport mockTransport = + new MockHttpTransport.Builder() + .setLowLevelHttpRequest( + new MockLowLevelHttpRequest() { + int index = 0; + + @Override + public LowLevelHttpResponse execute() throws IOException { + return mockResponses[index++]; + } + }) + .build(); + + GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); + + gcsUtil.delegate.setStorageClient( + new Storage(mockTransport, Transport.getJsonFactory(), new RetryHttpRequestInitializer())); + gcsUtil.delegate.fileSizes( + ImmutableList.of(GcsPath.fromComponents("testbucket", "testobject"))); + } + + @Test + public void testRemoveWhenFileNotFound() throws Exception { + JsonFactory jsonFactory = new GsonFactory(); + + String contentBoundary = "batch_foobarbaz"; + String contentBoundaryLine = "--" + contentBoundary; + String endOfContentBoundaryLine = "--" + contentBoundary + "--"; + + GenericJson error = new GenericJson().set("error", new GenericJson().set("code", 404)); + error.setFactory(jsonFactory); + + String content = + contentBoundaryLine + + "\n" + + "Content-Type: application/http\n" + + "\n" + + "HTTP/1.1 404 Not Found\n" + + "Content-Length: -1\n" + + "\n" + + error.toString() + + "\n" + + "\n" + + endOfContentBoundaryLine + + "\n"; + + final LowLevelHttpResponse mockResponse = Mockito.mock(LowLevelHttpResponse.class); + when(mockResponse.getContentType()).thenReturn("multipart/mixed; boundary=" + contentBoundary); + when(mockResponse.getStatusCode()).thenReturn(200); + when(mockResponse.getContent()).thenReturn(toStream(content)); + + // A mock transport that lets us mock the API responses. + MockLowLevelHttpRequest request = + new MockLowLevelHttpRequest() { + @Override + public LowLevelHttpResponse execute() throws IOException { + return mockResponse; + } + }; + MockHttpTransport mockTransport = + new MockHttpTransport.Builder().setLowLevelHttpRequest(request).build(); + + GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); + gcsUtil.delegate.setStorageClient( + new Storage(mockTransport, Transport.getJsonFactory(), new RetryHttpRequestInitializer())); + gcsUtil.remove(Arrays.asList("gs://some-bucket/already-deleted")); + } + + @Test + public void testCreateBucket() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + Storage.Buckets.Insert mockStorageInsert = Mockito.mock(Storage.Buckets.Insert.class); + + BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); + + when(mockStorage.buckets()).thenReturn(mockStorageObjects); + when(mockStorageObjects.insert(any(String.class), any(Bucket.class))) + .thenReturn(mockStorageInsert); + when(mockStorageInsert.execute()) + .thenThrow(new SocketTimeoutException("SocketException")) + .thenReturn(new Bucket()); + + gcsUtil.delegate.createBucket( + "a", new Bucket(), mockBackOff, new FastNanoClockAndSleeper()::sleep); + } + + @Test + public void testCreateBucketAccessErrors() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); + Storage.Buckets.Insert mockStorageInsert = Mockito.mock(Storage.Buckets.Insert.class); + + BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); + GoogleJsonResponseException expectedException = + googleJsonResponseException( + HttpStatusCodes.STATUS_CODE_FORBIDDEN, + "Waves hand mysteriously", + "These aren't the buckets you're looking for"); + + when(mockStorage.buckets()).thenReturn(mockStorageObjects); + when(mockStorageObjects.insert(any(String.class), any(Bucket.class))) + .thenReturn(mockStorageInsert); + when(mockStorageInsert.execute()).thenThrow(expectedException); + + thrown.expect(AccessDeniedException.class); + + gcsUtil.delegate.createBucket( + "a", new Bucket(), mockBackOff, new FastNanoClockAndSleeper()::sleep); + } + + @Test + public void testBucketAccessible() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); + Storage.Buckets.Get mockStorageGet = Mockito.mock(Storage.Buckets.Get.class); + + BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); + + when(mockStorage.buckets()).thenReturn(mockStorageObjects); + when(mockStorageObjects.get("testbucket")).thenReturn(mockStorageGet); + when(mockStorageGet.execute()) + .thenThrow(new SocketTimeoutException("SocketException")) + .thenReturn(new Bucket()); + + assertTrue( + gcsUtil.delegate.bucketAccessible( + GcsPath.fromComponents("testbucket", "testobject"), + mockBackOff, + new FastNanoClockAndSleeper()::sleep)); + } + + @Test + public void testBucketDoesNotExistBecauseOfAccessError() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); + Storage.Buckets.Get mockStorageGet = Mockito.mock(Storage.Buckets.Get.class); + + BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); + GoogleJsonResponseException expectedException = + googleJsonResponseException( + HttpStatusCodes.STATUS_CODE_FORBIDDEN, + "Waves hand mysteriously", + "These aren't the buckets you're looking for"); + + when(mockStorage.buckets()).thenReturn(mockStorageObjects); + when(mockStorageObjects.get("testbucket")).thenReturn(mockStorageGet); + when(mockStorageGet.execute()).thenThrow(expectedException); + + assertFalse( + gcsUtil.delegate.bucketAccessible( + GcsPath.fromComponents("testbucket", "testobject"), + mockBackOff, + new FastNanoClockAndSleeper()::sleep)); + } + + @Test + public void testBucketDoesNotExist() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); + Storage.Buckets.Get mockStorageGet = Mockito.mock(Storage.Buckets.Get.class); + + BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); + + when(mockStorage.buckets()).thenReturn(mockStorageObjects); + when(mockStorageObjects.get("testbucket")).thenReturn(mockStorageGet); + when(mockStorageGet.execute()) + .thenThrow( + googleJsonResponseException( + HttpStatusCodes.STATUS_CODE_NOT_FOUND, "It don't exist", "Nothing here to see")); + + assertFalse( + gcsUtil.delegate.bucketAccessible( + GcsPath.fromComponents("testbucket", "testobject"), + mockBackOff, + new FastNanoClockAndSleeper()::sleep)); + } + + @Test + public void testVerifyBucketAccessible() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); + Storage.Buckets.Get mockStorageGet = Mockito.mock(Storage.Buckets.Get.class); + + BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); + + when(mockStorage.buckets()).thenReturn(mockStorageObjects); + when(mockStorageObjects.get("testbucket")).thenReturn(mockStorageGet); + when(mockStorageGet.execute()) + .thenThrow(new SocketTimeoutException("SocketException")) + .thenReturn(new Bucket()); + + gcsUtil.delegate.verifyBucketAccessible( + GcsPath.fromComponents("testbucket", "testobject"), + mockBackOff, + new FastNanoClockAndSleeper()::sleep); + } + + @Test(expected = AccessDeniedException.class) + public void testVerifyBucketAccessibleAccessError() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); + Storage.Buckets.Get mockStorageGet = Mockito.mock(Storage.Buckets.Get.class); + + BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); + GoogleJsonResponseException expectedException = + googleJsonResponseException( + HttpStatusCodes.STATUS_CODE_FORBIDDEN, + "Waves hand mysteriously", + "These aren't the buckets you're looking for"); + + when(mockStorage.buckets()).thenReturn(mockStorageObjects); + when(mockStorageObjects.get("testbucket")).thenReturn(mockStorageGet); + when(mockStorageGet.execute()).thenThrow(expectedException); + + gcsUtil.delegate.verifyBucketAccessible( + GcsPath.fromComponents("testbucket", "testobject"), + mockBackOff, + new FastNanoClockAndSleeper()::sleep); + } + + @Test(expected = FileNotFoundException.class) + public void testVerifyBucketAccessibleDoesNotExist() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); + Storage.Buckets.Get mockStorageGet = Mockito.mock(Storage.Buckets.Get.class); + + BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); + + when(mockStorage.buckets()).thenReturn(mockStorageObjects); + when(mockStorageObjects.get("testbucket")).thenReturn(mockStorageGet); + when(mockStorageGet.execute()) + .thenThrow( + googleJsonResponseException( + HttpStatusCodes.STATUS_CODE_NOT_FOUND, "It don't exist", "Nothing here to see")); + + gcsUtil.delegate.verifyBucketAccessible( + GcsPath.fromComponents("testbucket", "testobject"), + mockBackOff, + new FastNanoClockAndSleeper()::sleep); + } + + @Test + public void testGetBucket() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); + Storage.Buckets.Get mockStorageGet = Mockito.mock(Storage.Buckets.Get.class); + + BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); + + when(mockStorage.buckets()).thenReturn(mockStorageObjects); + when(mockStorageObjects.get("testbucket")).thenReturn(mockStorageGet); + when(mockStorageGet.execute()) + .thenThrow(new SocketTimeoutException("SocketException")) + .thenReturn(new Bucket()); + + assertNotNull( + gcsUtil.delegate.getBucket( + GcsPath.fromComponents("testbucket", "testobject"), + mockBackOff, + new FastNanoClockAndSleeper()::sleep)); + } + + @Test + public void testGetBucketNotExists() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + Storage.Buckets mockStorageObjects = Mockito.mock(Storage.Buckets.class); + Storage.Buckets.Get mockStorageGet = Mockito.mock(Storage.Buckets.Get.class); + + BackOff mockBackOff = BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()); + + when(mockStorage.buckets()).thenReturn(mockStorageObjects); + when(mockStorageObjects.get("testbucket")).thenReturn(mockStorageGet); + when(mockStorageGet.execute()) + .thenThrow( + googleJsonResponseException( + HttpStatusCodes.STATUS_CODE_NOT_FOUND, "It don't exist", "Nothing here to see")); + + thrown.expect(FileNotFoundException.class); + thrown.expectMessage("It don't exist"); + gcsUtil.delegate.getBucket( + GcsPath.fromComponents("testbucket", "testobject"), + mockBackOff, + new FastNanoClockAndSleeper()::sleep); + } + + @Test + public void testGCSChannelCloseIdempotent() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + GoogleCloudStorageReadOptions readOptions = + GoogleCloudStorageReadOptions.builder().setFastFailOnNotFoundEnabled(false).build(); + + gcsUtil.delegate.setCloudStorageImpl( + GoogleCloudStorageOptions.builder() + .setAppName("Beam") + .setGrpcEnabled(true) + .setProjectId("my_project") + .setReadChannelOptions(readOptions) + .build()); + + SeekableByteChannel channel = gcsUtil.open(GcsPath.fromComponents("testbucket", "testobject")); + channel.close(); + channel.close(); + } + + @Test + public void testGCSReadMetricsIsSet() { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + GoogleCloudStorageReadOptions readOptions = + GoogleCloudStorageReadOptions.builder().setFastFailOnNotFoundEnabled(true).build(); + gcsUtil.delegate.setCloudStorageImpl( + GoogleCloudStorageOptions.builder() + .setAppName("Beam") + .setGrpcEnabled(true) + .setProjectId("my_project") + .setReadChannelOptions(readOptions) + .build()); + assertThrows( + IOException.class, () -> gcsUtil.open(GcsPath.fromComponents("testbucket", "testbucket"))); + verifyMetricWasSet("my_project", "testbucket", "GcsGet", "permission_denied", 1); + } + + @Test + public void testGCSWriteMetricsIsSet() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + GoogleCloudStorage mockStorage = Mockito.mock(GoogleCloudStorage.class); + gcsUtil.delegate.setCloudStorageImpl( + GoogleCloudStorageOptions.builder() + .setAppName("Beam") + .setGrpcEnabled(true) + .setProjectId("my_project") + .build()); + when(mockStorage.create( + new StorageResourceId("testbucket", "testobject"), + CreateObjectOptions.builder() + .setOverwriteExisting(true) + .setContentType("type") + .build())) + .thenThrow(IOException.class); + GcsPath gcsPath = GcsPath.fromComponents("testbucket", "testobject"); + assertThrows(IOException.class, () -> gcsUtil.create(gcsPath, "")); + verifyMetricWasSet("my_project", "testbucket", "GcsInsert", "permission_denied", 1); + } + + private void verifyMetricWasSet( + String projectId, String bucketId, String method, String status, long count) { + // Verify the metric as reported. + HashMap labels = new HashMap<>(); + labels.put(MonitoringInfoConstants.Labels.PTRANSFORM, ""); + labels.put(MonitoringInfoConstants.Labels.SERVICE, "Storage"); + labels.put(MonitoringInfoConstants.Labels.METHOD, method); + labels.put(MonitoringInfoConstants.Labels.GCS_PROJECT_ID, projectId); + labels.put(MonitoringInfoConstants.Labels.GCS_BUCKET, bucketId); + labels.put( + MonitoringInfoConstants.Labels.RESOURCE, + GcpResourceIdentifiers.cloudStorageBucket(bucketId)); + labels.put(MonitoringInfoConstants.Labels.STATUS, status); + + MonitoringInfoMetricName name = + MonitoringInfoMetricName.named(MonitoringInfoConstants.Urns.API_REQUEST_COUNT, labels); + MetricsContainerImpl container = + (MetricsContainerImpl) MetricsEnvironment.getProcessWideContainer(); + assertEquals(count, (long) container.getCounter(name).getCumulative()); + } + + /** Builds a fake GoogleJsonResponseException for testing API error handling. */ + private static GoogleJsonResponseException googleJsonResponseException( + final int status, final String reason, final String message) throws IOException { + final JsonFactory jsonFactory = new GsonFactory(); + HttpTransport transport = + new MockHttpTransport() { + @Override + public LowLevelHttpRequest buildRequest(String method, String url) throws IOException { + ErrorInfo errorInfo = new ErrorInfo(); + errorInfo.setReason(reason); + errorInfo.setMessage(message); + errorInfo.setFactory(jsonFactory); + GenericJson error = new GenericJson(); + error.set("code", status); + error.set("errors", Arrays.asList(errorInfo)); + error.setFactory(jsonFactory); + GenericJson errorResponse = new GenericJson(); + errorResponse.set("error", error); + errorResponse.setFactory(jsonFactory); + return new MockLowLevelHttpRequest() + .setResponse( + new MockLowLevelHttpResponse() + .setContent(errorResponse.toPrettyString()) + .setContentType(Json.MEDIA_TYPE) + .setStatusCode(status)); + } + }; + HttpRequest request = + transport.createRequestFactory().buildGetRequest(HttpTesting.SIMPLE_GENERIC_URL); + request.setThrowExceptionOnExecuteError(false); + HttpResponse response = request.execute(); + return GoogleJsonResponseException.from(jsonFactory, response); + } + + private static List makeStrings(String s, int n) { + return makeStrings("bucket", s, n); + } + + private static List makeStrings(String bucket, String s, int n) { + ImmutableList.Builder ret = ImmutableList.builder(); + for (int i = 0; i < n; ++i) { + ret.add(String.format("gs://%s/%s%d", bucket, s, i)); + } + return ret.build(); + } + + private static List makeGcsPaths(String s, int n) { + ImmutableList.Builder ret = ImmutableList.builder(); + for (int i = 0; i < n; ++i) { + ret.add(GcsPath.fromUri(String.format("gs://bucket/%s%d", s, i))); + } + return ret.build(); + } + + private static int sumBatchSizes(List batches) { + int ret = 0; + for (BatchInterface b : batches) { + ret += b.size(); + assertThat(b.size(), greaterThan(0)); + } + return ret; + } + + @Test + public void testMakeRewriteOps() throws IOException { + GcsOptions gcsOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = gcsOptions.getGcsUtil(); + + LinkedList rewrites = + gcsUtil.delegate.makeRewriteOps( + makeStrings("s", 1), makeStrings("d", 1), false, false, false); + assertEquals(1, rewrites.size()); + + RewriteOp rewrite = rewrites.pop(); + assertTrue(rewrite.getReadyToEnqueue()); + Storage.Objects.Rewrite request = rewrite.rewriteRequest; + assertNull(request.getMaxBytesRewrittenPerCall()); + assertEquals("bucket", request.getSourceBucket()); + assertEquals("s0", request.getSourceObject()); + assertEquals("bucket", request.getDestinationBucket()); + assertEquals("d0", request.getDestinationObject()); + } + + @Test + public void testMakeRewriteOpsWithOptions() throws IOException { + GcsOptions gcsOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = gcsOptions.getGcsUtil(); + gcsUtil.delegate.maxBytesRewrittenPerCall = 1337L; + + LinkedList rewrites = + gcsUtil.delegate.makeRewriteOps( + makeStrings("s", 1), makeStrings("d", 1), false, false, false); + assertEquals(1, rewrites.size()); + + RewriteOp rewrite = rewrites.pop(); + assertTrue(rewrite.getReadyToEnqueue()); + Storage.Objects.Rewrite request = rewrite.rewriteRequest; + assertEquals(Long.valueOf(1337L), request.getMaxBytesRewrittenPerCall()); + } + + @Test + public void testMakeRewriteBatches() throws IOException { + GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); + + // Small number of files fits in 1 batch + List batches = + gcsUtil.delegate.makeRewriteBatches( + gcsUtil.delegate.makeRewriteOps( + makeStrings("s", 3), makeStrings("d", 3), false, false, false)); + assertThat(batches.size(), equalTo(1)); + assertThat(sumBatchSizes(batches), equalTo(3)); + + // 1 batch of files fits in 1 batch + batches = + gcsUtil.delegate.makeRewriteBatches( + gcsUtil.delegate.makeRewriteOps( + makeStrings("s", 100), makeStrings("d", 100), false, false, false)); + assertThat(batches.size(), equalTo(1)); + assertThat(sumBatchSizes(batches), equalTo(100)); + + // A little more than 5 batches of files fits in 6 batches + batches = + gcsUtil.delegate.makeRewriteBatches( + gcsUtil.delegate.makeRewriteOps( + makeStrings("s", 501), makeStrings("d", 501), false, false, false)); + assertThat(batches.size(), equalTo(6)); + assertThat(sumBatchSizes(batches), equalTo(501)); + } + + @Test + public void testMakeRewriteBatchesWithLowerDataOpLimit() throws IOException { + GcsOptions options = gcsOptionsWithTestCredential(); + options.setGcsRewriteDataOpBatchLimit(2); + GcsUtil gcsUtil = options.getGcsUtil(); + + // Small number of files in same bucket fits in 1 batch + List batches = + gcsUtil.delegate.makeRewriteBatches( + gcsUtil.delegate.makeRewriteOps( + makeStrings("s", 5), makeStrings("d", 5), false, false, false)); + assertThat(batches.size(), equalTo(1)); + assertThat(sumBatchSizes(batches), equalTo(5)); + + // Files copying between buckets use smaller batch size + batches = + gcsUtil.delegate.makeRewriteBatches( + gcsUtil.delegate.makeRewriteOps( + makeStrings("bucket1", "s", 5), + makeStrings("bucket2", "d", 5), + false, + false, + false)); + assertThat(batches.size(), equalTo(3)); + assertThat(sumBatchSizes(batches), equalTo(5)); + + // A mix of same bucket and different buckets uses large batches when possible. + List fromFiles = new ArrayList<>(makeStrings("bucket1", "s", 3)); + List toFiles = new ArrayList<>(makeStrings("bucket2", "d", 3)); + fromFiles.addAll(makeStrings("t", 90)); + toFiles.addAll(makeStrings("e", 90)); + fromFiles.addAll(makeStrings("bucket3", "u", 3)); + toFiles.addAll(makeStrings("bucket4", "f", 3)); + fromFiles.addAll(makeStrings("bucket5", "v", 1)); + toFiles.addAll(makeStrings("bucket5", "g", 1)); + + batches = + gcsUtil.delegate.makeRewriteBatches( + gcsUtil.delegate.makeRewriteOps(fromFiles, toFiles, false, false, false)); + assertThat(batches.size(), equalTo(4)); + assertThat(batches.get(0).size(), equalTo(91)); + assertThat(sumBatchSizes(batches), equalTo(97)); + } + + @Test + public void testMakeRewriteOpsInvalid() throws IOException { + GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); + thrown.expect(IllegalArgumentException.class); + thrown.expectMessage("Number of source files 3"); + + gcsUtil.delegate.makeRewriteOps(makeStrings("s", 3), makeStrings("d", 1), false, false, false); + } + + private class FakeBatcher implements BatchInterface { + ArrayList> requests = new ArrayList<>(); + + @Override + public void queue(AbstractGoogleJsonClientRequest request, JsonBatchCallback cb) { + assertNotNull(request); + assertNotNull(cb); + requests.add( + () -> { + try { + try { + T result = request.execute(); + cb.onSuccess(result, null); + } catch (FileNotFoundException e) { + GoogleJsonError error = new GoogleJsonError(); + error.setCode(HttpStatusCodes.STATUS_CODE_NOT_FOUND); + cb.onFailure(error, null); + } catch (GoogleJsonResponseException e) { + cb.onFailure(e.getDetails(), null); + } catch (SocketTimeoutException e) { + System.out.println("Propagating socket exception as batch processing error"); + throw e; + } catch (Exception e) { + System.out.println("Propagating exception as server error " + e); + e.printStackTrace(); + GoogleJsonError error = new GoogleJsonError(); + error.setCode(HttpStatusCodes.STATUS_CODE_SERVER_ERROR); + cb.onFailure(error, null); + } + } catch (IOException e) { + throw new RuntimeException(e); + } + return null; + }); + } + + @Override + public void execute() throws IOException { + RuntimeException lastException = null; + for (Supplier request : requests) { + try { + request.get(); + } catch (RuntimeException e) { + lastException = e; + } + } + if (lastException != null) { + throw lastException; + } + } + + @Override + public int size() { + return requests.size(); + } + } + + @Test + public void testRename() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); + + Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); + Storage.Objects.Rewrite mockStorageRewrite = Mockito.mock(Storage.Objects.Rewrite.class); + Storage.Objects.Delete mockStorageDelete1 = Mockito.mock(Storage.Objects.Delete.class); + Storage.Objects.Delete mockStorageDelete2 = Mockito.mock(Storage.Objects.Delete.class); + + when(mockStorage.objects()).thenReturn(mockStorageObjects); + when(mockStorageObjects.rewrite("bucket", "s0", "bucket", "d0", null)) + .thenReturn(mockStorageRewrite); + when(mockStorageRewrite.execute()) + .thenThrow(new InvalidObjectException("Test exception")) + .thenReturn(new RewriteResponse().setDone(true)); + when(mockStorageObjects.delete("bucket", "s0")) + .thenReturn(mockStorageDelete1) + .thenReturn(mockStorageDelete2); + + when(mockStorageDelete1.execute()).thenThrow(new InvalidObjectException("Test exception")); + + gcsUtil.rename(makeStrings("s", 1), makeStrings("d", 1)); + verify(mockStorageRewrite, times(2)).execute(); + verify(mockStorageDelete1, times(1)).execute(); + verify(mockStorageDelete2, times(1)).execute(); + } + + @Test + public void testRenameIgnoringMissing() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); + + Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); + Storage.Objects.Rewrite mockStorageRewrite1 = Mockito.mock(Storage.Objects.Rewrite.class); + Storage.Objects.Rewrite mockStorageRewrite2 = Mockito.mock(Storage.Objects.Rewrite.class); + Storage.Objects.Delete mockStorageDelete = Mockito.mock(Storage.Objects.Delete.class); + + when(mockStorage.objects()).thenReturn(mockStorageObjects); + when(mockStorageObjects.rewrite("bucket", "s0", "bucket", "d0", null)) + .thenReturn(mockStorageRewrite1); + when(mockStorageRewrite1.execute()).thenThrow(new FileNotFoundException()); + when(mockStorageObjects.rewrite("bucket", "s1", "bucket", "d1", null)) + .thenReturn(mockStorageRewrite2); + when(mockStorageRewrite2.execute()).thenReturn(new RewriteResponse().setDone(true)); + when(mockStorageObjects.delete("bucket", "s1")).thenReturn(mockStorageDelete); + + gcsUtil.rename( + makeStrings("s", 2), makeStrings("d", 2), StandardMoveOptions.IGNORE_MISSING_FILES); + verify(mockStorageRewrite1, times(1)).execute(); + verify(mockStorageRewrite2, times(1)).execute(); + verify(mockStorageDelete, times(1)).execute(); + } + + @Test + public void testRenamePropagateMissingException() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); + + Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); + Storage.Objects.Rewrite mockStorageRewrite = Mockito.mock(Storage.Objects.Rewrite.class); + + when(mockStorage.objects()).thenReturn(mockStorageObjects); + when(mockStorageObjects.rewrite("bucket", "s0", "bucket", "d0", null)) + .thenReturn(mockStorageRewrite); + when(mockStorageRewrite.execute()).thenThrow(new FileNotFoundException()); + + assertThrows(IOException.class, () -> gcsUtil.rename(makeStrings("s", 1), makeStrings("d", 1))); + verify(mockStorageRewrite, times(1)).execute(); + } + + @Test + public void testRenameSkipDestinationExistsSameBucket() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); + + Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); + Storage.Objects.Rewrite mockStorageRewrite = Mockito.mock(Storage.Objects.Rewrite.class); + Storage.Objects.Delete mockStorageDelete = Mockito.mock(Storage.Objects.Delete.class); + + when(mockStorage.objects()).thenReturn(mockStorageObjects); + when(mockStorageObjects.rewrite("bucket", "s0", "bucket", "d0", null)) + .thenReturn(mockStorageRewrite); + when(mockStorageRewrite.execute()).thenReturn(new RewriteResponse().setDone(true)); + when(mockStorageObjects.delete("bucket", "s0")).thenReturn(mockStorageDelete); + + gcsUtil.rename( + makeStrings("s", 1), makeStrings("d", 1), StandardMoveOptions.SKIP_IF_DESTINATION_EXISTS); + verify(mockStorageRewrite, times(1)).execute(); + verify(mockStorageDelete, times(1)).execute(); + } + + @Test + public void testRenameSkipDestinationExistsDifferentBucket() throws IOException { + GcsOptions pipelineOptions = gcsOptionsWithTestCredential(); + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + + assertThrows( + UnsupportedOperationException.class, + () -> + gcsUtil.rename( + Collections.singletonList("gs://bucket/source"), + Collections.singletonList("gs://different_bucket/dest"), + StandardMoveOptions.SKIP_IF_DESTINATION_EXISTS)); + } + + @Test + public void testThrowRetentionPolicyNotMetErrorWhenUnequalChecksum() throws IOException { + // ./gradlew sdks:java:extensions:google-cloud-platform-core:test --tests + // org.apache.beam.sdk.extensions.gcp.util.GcsUtilV1Test.testHanRetentionPolicyNotMetError + GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); + + Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); + Storage.Objects.Get mockGetRequest1 = Mockito.mock(Storage.Objects.Get.class); + Storage.Objects.Get mockGetRequest2 = Mockito.mock(Storage.Objects.Get.class); + Storage.Objects.Rewrite mockStorageRewrite = Mockito.mock(Storage.Objects.Rewrite.class); + + // Gcs object to be used when checking the hash of the files during rewrite fail. + StorageObject srcObject = new StorageObject().setMd5Hash("a"); + StorageObject destObject = new StorageObject().setMd5Hash("b"); + + when(mockStorage.objects()).thenReturn(mockStorageObjects); + when(mockStorageObjects.rewrite("bucket", "s0", "bucket", "d0", null)) + .thenReturn(mockStorageRewrite); + when(mockStorageRewrite.execute()) + .thenThrow(googleJsonResponseException(403, "retentionPolicyNotMet", "Too soon")); + when(mockStorageObjects.get("bucket", "s0")).thenReturn(mockGetRequest1); + when(mockGetRequest1.execute()).thenReturn(srcObject); + when(mockStorageObjects.get("bucket", "d0")).thenReturn(mockGetRequest2); + when(mockGetRequest2.execute()).thenReturn(destObject); + + assertThrows(IOException.class, () -> gcsUtil.rename(makeStrings("s", 1), makeStrings("d", 1))); + + verify(mockStorageRewrite, times(1)).execute(); + } + + @Test + public void testIgnoreRetentionPolicyNotMetErrorWhenEqualChecksum() throws IOException { + GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); + + Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); + Storage.Objects.Get mockGetRequest = Mockito.mock(Storage.Objects.Get.class); + Storage.Objects.Rewrite mockStorageRewrite1 = Mockito.mock(Storage.Objects.Rewrite.class); + Storage.Objects.Rewrite mockStorageRewrite2 = Mockito.mock(Storage.Objects.Rewrite.class); + Storage.Objects.Delete mockStorageDelete = Mockito.mock(Storage.Objects.Delete.class); + + // Gcs object to be used when checking the hash of the files during rewrite fail. + StorageObject gcsObject = new StorageObject().setMd5Hash("a"); + + when(mockStorage.objects()).thenReturn(mockStorageObjects); + // First rewrite with retentionPolicyNotMet error. + when(mockStorageObjects.rewrite("bucket", "s0", "bucket", "d0", null)) + .thenReturn(mockStorageRewrite1); + when(mockStorageRewrite1.execute()) + .thenThrow(googleJsonResponseException(403, "retentionPolicyNotMet", "Too soon")); + when(mockStorageObjects.get(any(), any())) // to access object hash during error handling + .thenReturn(mockGetRequest); + when(mockGetRequest.execute()) + .thenReturn(gcsObject); // both source and destination will get the same hash + when(mockStorageObjects.delete("bucket", "s0")).thenReturn(mockStorageDelete); + + // Second rewrite should not be affected. + when(mockStorageObjects.rewrite("bucket", "s1", "bucket", "d1", null)) + .thenReturn(mockStorageRewrite2); + when(mockStorageRewrite2.execute()).thenReturn(new RewriteResponse().setDone(true)); + when(mockStorageObjects.delete("bucket", "s1")).thenReturn(mockStorageDelete); + + gcsUtil.rename(makeStrings("s", 2), makeStrings("d", 2)); + + verify(mockStorageRewrite1, times(1)).execute(); + verify(mockStorageRewrite2, times(1)).execute(); + verify(mockStorageDelete, times(2)).execute(); + } + + @Test + public void testMakeRemoveBatches() throws IOException { + GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); + + // Small number of files fits in 1 batch + List batches = gcsUtil.delegate.makeRemoveBatches(makeStrings("s", 3)); + assertThat(batches.size(), equalTo(1)); + assertThat(sumBatchSizes(batches), equalTo(3)); + + // 1 batch of files fits in 1 batch + batches = gcsUtil.delegate.makeRemoveBatches(makeStrings("s", 100)); + assertThat(batches.size(), equalTo(1)); + assertThat(sumBatchSizes(batches), equalTo(100)); + + // A little more than 5 batches of files fits in 6 batches + batches = gcsUtil.delegate.makeRemoveBatches(makeStrings("s", 501)); + assertThat(batches.size(), equalTo(6)); + assertThat(sumBatchSizes(batches), equalTo(501)); + } + + @Test + public void testMakeGetBatches() throws IOException { + GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); + + // Small number of files fits in 1 batch + List results = Lists.newArrayList(); + List batches = gcsUtil.delegate.makeGetBatches(makeGcsPaths("s", 3), results); + assertThat(batches.size(), equalTo(1)); + assertThat(sumBatchSizes(batches), equalTo(3)); + assertEquals(3, results.size()); + + // 1 batch of files fits in 1 batch + results = Lists.newArrayList(); + batches = gcsUtil.delegate.makeGetBatches(makeGcsPaths("s", 100), results); + assertThat(batches.size(), equalTo(1)); + assertThat(sumBatchSizes(batches), equalTo(100)); + assertEquals(100, results.size()); + + // A little more than 5 batches of files fits in 6 batches + results = Lists.newArrayList(); + batches = gcsUtil.delegate.makeGetBatches(makeGcsPaths("s", 501), results); + assertThat(batches.size(), equalTo(6)); + assertThat(sumBatchSizes(batches), equalTo(501)); + assertEquals(501, results.size()); + } + + @Test + public void testGetObjects() throws IOException { + GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); + + Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); + Storage.Objects.Get mockGetRequest = Mockito.mock(Storage.Objects.Get.class); + StorageObject object = new StorageObject(); + when(mockGetRequest.execute()).thenReturn(object); + when(mockStorageObjects.get(any(), any())).thenReturn(mockGetRequest); + when(mockStorage.objects()).thenReturn(mockStorageObjects); + + List results = gcsUtil.getObjects(makeGcsPaths("s", 1)); + + assertEquals(object, results.get(0).storageObject()); + } + + @Test + public void testGetObjectsWithException() throws IOException { + GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); + + Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); + Storage.Objects.Get mockGetRequest = Mockito.mock(Storage.Objects.Get.class); + when(mockStorage.objects()).thenReturn(mockStorageObjects); + when(mockStorageObjects.get(any(), any())).thenReturn(mockGetRequest); + when(mockGetRequest.execute()).thenThrow(new RuntimeException("fakeException")); + + thrown.expect(IOException.class); + thrown.expectMessage("Error trying to get gs://bucket/s0"); + + List results = gcsUtil.getObjects(makeGcsPaths("s", 1)); + + for (StorageObjectOrIOException result : results) { + if (null != result.ioException()) { + throw result.ioException(); + } + } + } + + @Test + public void testListObjectsException() throws IOException { + GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil(); + + Storage mockStorage = Mockito.mock(Storage.class); + gcsUtil.delegate.setStorageClient(mockStorage); + gcsUtil.delegate.setBatchRequestSupplier(FakeBatcher::new); + + Storage.Objects mockStorageObjects = Mockito.mock(Storage.Objects.class); + when(mockStorage.objects()).thenReturn(mockStorageObjects); + Storage.Objects.List mockStorageList = Mockito.mock(Storage.Objects.List.class); + when(mockStorageObjects.list(any())).thenReturn(mockStorageList); + when(mockStorageList.execute()).thenThrow(new RuntimeException("FakeException")); + + thrown.expect(IOException.class); + thrown.expectMessage("Unable to match files in bucket testBucket"); + + gcsUtil.listObjects("testBucket", "prefix", null); + } + + public static class GcsUtilMock extends GcsUtil { + + public static GcsUtilMock createMockWithMockStorage(PipelineOptions options, byte[] readPayload) + throws IOException { + GcsUtilMock gcsUtilMock = createMock(options); + + GcsUtilV1Mock mockLegacy = GcsUtilV1Mock.createMockWithMockStorage(options, readPayload); + gcsUtilMock.delegate = mockLegacy; + + return gcsUtilMock; + } + + public static GcsUtilMock createMock(PipelineOptions options) { + return new GcsUtilMock(options); + } + + private GcsUtilMock(PipelineOptions options) { + super(options); + } + } + + public static class GcsUtilV1Mock extends GcsUtilV1 { + + public GoogleCloudStorage googleCloudStorage; + + public static GcsUtilV1Mock createMockWithMockStorage( + PipelineOptions options, byte[] readPayload) throws IOException { + GcsUtilV1Mock gcsUtilMock = createMock(options); + GoogleCloudStorage googleCloudStorageMock = Mockito.mock(GoogleCloudStorage.class); + gcsUtilMock.googleCloudStorage = googleCloudStorageMock; + // set the mock in the super object as well + gcsUtilMock.setCloudStorageImpl(gcsUtilMock.googleCloudStorage); + + if (readPayload == null) { + Mockito.when(googleCloudStorageMock.create(Mockito.any(), Mockito.any())) + .thenReturn(Channels.newChannel(new ByteArrayOutputStream())); + } else { + SeekableByteChannel seekableByteChannel = new SeekableInMemoryByteChannel(readPayload); + Mockito.when(googleCloudStorageMock.open(Mockito.any(StorageResourceId.class))) + .thenReturn(seekableByteChannel); + Mockito.when( + googleCloudStorageMock.open(Mockito.any(StorageResourceId.class), Mockito.any())) + .thenReturn(seekableByteChannel); + } + return gcsUtilMock; + } + + public static GcsUtilV1Mock createMock(PipelineOptions options) { + GcsOptions gcsOptions = options.as(GcsOptions.class); + Storage.Builder storageBuilder = Transport.newStorageClient(gcsOptions); + return new GcsUtilV1Mock( + storageBuilder.build(), + storageBuilder.getHttpRequestInitializer(), + gcsOptions.getExecutorService(), + hasExperiment(options, "use_grpc_for_gcs"), + gcsOptions.getGcpCredential(), + gcsOptions.getGcsUploadBufferSizeBytes(), + gcsOptions.getGcsRewriteDataOpBatchLimit(), + GcsUtilV1.GcsCountersOptions.create( + gcsOptions.getEnableBucketReadMetricCounter() + ? gcsOptions.getGcsReadCounterPrefix() + : null, + gcsOptions.getEnableBucketWriteMetricCounter() + ? gcsOptions.getGcsWriteCounterPrefix() + : null, + Boolean.TRUE.equals(gcsOptions.getGcsPerformanceMetrics())), + gcsOptions.getGoogleCloudStorageReadOptions()); + } + + private GcsUtilV1Mock( + Storage storageClient, + HttpRequestInitializer httpRequestInitializer, + ExecutorService executorService, + Boolean shouldUseGrpc, + Credentials credentials, + @Nullable Integer uploadBufferSizeBytes, + @Nullable Integer rewriteDataOpBatchLimit, + GcsUtilV1.GcsCountersOptions gcsCountersOptions, + GoogleCloudStorageReadOptions gcsReadOptions) { + super( + storageClient, + httpRequestInitializer, + executorService, + shouldUseGrpc, + credentials, + uploadBufferSizeBytes, + rewriteDataOpBatchLimit, + gcsCountersOptions, + gcsReadOptions); + } + + @Override + GoogleCloudStorage createGoogleCloudStorage( + GoogleCloudStorageOptions options, + Storage storage, + Credentials credentials, + @Nullable HttpRequestInitializer httpRequestInitializer) { + return googleCloudStorage; + } + } + + @Test + public void testCreate() throws IOException { + GcsOptions gcsOptions = gcsOptionsWithTestCredential(); + + GcsUtilMock gcsUtil = GcsUtilMock.createMock(gcsOptions); + + GoogleCloudStorage mockStorage = Mockito.mock(GoogleCloudStorage.class); + WritableByteChannel mockChannel = Mockito.mock(WritableByteChannel.class); + + GcsUtilV1Mock mockLegacy = GcsUtilV1Mock.createMock(gcsOptions); + mockLegacy.googleCloudStorage = mockStorage; + gcsUtil.delegate = mockLegacy; + + when(mockStorage.create(any(), any())).thenReturn(mockChannel); + + GcsPath path = GcsPath.fromUri("gs://testbucket/testdirectory/otherfile"); + CreateOptions createOptions = CreateOptions.builder().build(); + + assertEquals(mockChannel, gcsUtil.create(path, createOptions)); + } + + @Test + public void testCreateWithException() throws IOException { + GcsOptions gcsOptions = gcsOptionsWithTestCredential(); + + GcsUtilMock gcsUtil = GcsUtilMock.createMock(gcsOptions); + + GoogleCloudStorage mockStorage = Mockito.mock(GoogleCloudStorage.class); + + GcsUtilV1Mock mockLegacy = GcsUtilV1Mock.createMock(gcsOptions); + mockLegacy.googleCloudStorage = mockStorage; + gcsUtil.delegate = mockLegacy; + + when(mockStorage.create(any(), any())).thenThrow(new RuntimeException("testException")); + + GcsPath path = GcsPath.fromUri("gs://testbucket/testdirectory/otherfile"); + CreateOptions createOptions = CreateOptions.builder().build(); + + thrown.expect(RuntimeException.class); + thrown.expectMessage("testException"); + + gcsUtil.create(path, createOptions); + } + + private void testWriteMetrics(boolean enabled) throws IOException { + // arrange + GcsOptions gcsOptions = PipelineOptionsFactory.create().as(GcsOptions.class); + gcsOptions.setEnableBucketWriteMetricCounter(enabled); + gcsOptions.setGcsWriteCounterPrefix("test_counter"); + GcsUtilMock gcsUtil = GcsUtilMock.createMockWithMockStorage(gcsOptions, null); + byte[] payload = "some_bytes".getBytes(StandardCharsets.UTF_8); + String bucketName = "some_bucket"; + + // act + try (WritableByteChannel byteChannel = + gcsUtil.create(new GcsPath(null, bucketName, "o1"), CreateOptions.builder().build())) { + int bytesWrittenReportedByChannel = byteChannel.write(ByteBuffer.wrap(payload)); + long bytesWrittenReportedByMetric = + testMetricsContainer + .getCounter( + MetricName.named( + GcsUtil.class, + String.format("%s_%s", gcsOptions.getGcsWriteCounterPrefix(), bucketName))) + .getCumulative(); + + // assert + assertEquals(payload.length, bytesWrittenReportedByChannel); + assertEquals(enabled ? payload.length : 0, bytesWrittenReportedByMetric); + } + } + + private void testReadMetrics(boolean enabled, GoogleCloudStorageReadOptions readOptions) + throws IOException { + // arrange + GcsOptions gcsOptions = PipelineOptionsFactory.create().as(GcsOptions.class); + gcsOptions.setEnableBucketReadMetricCounter(enabled); + gcsOptions.setGcsReadCounterPrefix("test_counter"); + if (readOptions != null) { + gcsOptions.setGoogleCloudStorageReadOptions(readOptions); + } + byte[] payload = "some_bytes".getBytes(StandardCharsets.UTF_8); + GcsUtilMock gcsUtil = GcsUtilMock.createMockWithMockStorage(gcsOptions, payload); + String bucketName = "some_bucket"; + GcsPath gcsPath = new GcsPath(null, bucketName, "o1"); + // act + try (SeekableByteChannel byteChannel = gcsUtil.open(gcsPath)) { + int bytesReadReportedByChannel = byteChannel.read(ByteBuffer.allocate(payload.length)); + long bytesReadReportedByMetric = + testMetricsContainer + .getCounter( + MetricName.named( + GcsUtil.class, + String.format("%s_%s", gcsOptions.getGcsReadCounterPrefix(), bucketName))) + .getCumulative(); + + // assert + assertEquals(payload.length, bytesReadReportedByChannel); + assertEquals(enabled ? payload.length : 0, bytesReadReportedByMetric); + } + } + + @Test + public void testWriteMetricsAreCorrectlyReportedWhenEnabled() throws Exception { + testWriteMetrics(true); + } + + @Test + public void testWriteMetricsAreNotCollectedWhenNotEnabled() throws Exception { + testWriteMetrics(false); + } + + @Test + public void testReadMetricsAreCorrectlyReportedWhenEnabled() throws Exception { + testReadMetrics(true, null); + } + + @Test + public void testReadMetricsAreNotCollectedWhenNotEnabled() throws Exception { + testReadMetrics(false, null); + } + + @Test + public void testReadMetricsAreCorrectlyReportedWhenEnabledOpenWithOptions() throws Exception { + testReadMetrics(true, GoogleCloudStorageReadOptions.DEFAULT); + } + + @Test + public void testReadMetricsAreNotCollectedWhenNotEnabledOpenWithOptions() throws Exception { + testReadMetrics(false, GoogleCloudStorageReadOptions.DEFAULT); + } + + @Test + public void testGcsEndpoint() throws IOException { + GcsOptions pipelineOptions = PipelineOptionsFactory.as(GcsOptions.class); + pipelineOptions.setGcsEndpoint("http://localhost:4443/storage/v1/"); + + GcsUtil gcsUtil = pipelineOptions.getGcsUtil(); + GoogleCloudStorageImpl gcsImpl = + (GoogleCloudStorageImpl) gcsUtil.delegate.getGoogleCloudStorage(); + assertEquals("http://localhost:4443/", gcsImpl.getOptions().getStorageRootUrl()); + assertEquals("storage/v1/", gcsImpl.getOptions().getStorageServicePath()); + } + + /** A helper to wrap a {@link GenericJson} object in a content stream. */ + private static InputStream toStream(String content) throws IOException { + return new ByteArrayInputStream(content.getBytes(StandardCharsets.UTF_8)); + } +} diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java index 8bd79860cc28..7058efa46e06 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java @@ -152,7 +152,7 @@ public void testDefaultHostWithoutGcsEndpoint() { assertEquals(defaultHost, storageOptionsOf(gcsUtilV2(false)).getHost()); } - /** Mirrors {@code GcsUtilTest#testGcsEndpoint}: only the root of the endpoint applies to V2. */ + /** Mirrors {@code GcsUtilV1Test#testGcsEndpoint}: only the root of the endpoint applies to V2. */ @Test public void testGcsEndpointRootIsUsedAsHost() { GcsOptions options = gcsOptions(); @@ -273,7 +273,7 @@ public void testPerformanceMetricsFlagGatesTheContainer() { } // --------------------------------------------------------------------------------------------- - // Byte counters (mirrors GcsUtilTest#testReadMetrics / #testWriteMetrics for V1) + // Byte counters (mirrors GcsUtilV1Test#testReadMetrics / #testWriteMetrics for V1) // --------------------------------------------------------------------------------------------- private static @Nullable Long counterOrNull(MetricsContainerImpl container, MetricName name) { From a3beb0e05458e09a5dab43ff8c777363eda78b31 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Mon, 28 Sep 2026 10:37:39 -0400 Subject: [PATCH 2/6] Move V2 tests out of GcsUtilTest into GcsUtilV2Test --- .../sdk/extensions/gcp/util/GcsUtilTest.java | 131 ---------------- .../extensions/gcp/util/GcsUtilV2Test.java | 143 +++++++++++++++++- 2 files changed, 142 insertions(+), 132 deletions(-) diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java index 4a4c4eda28c8..24ee159653dc 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java @@ -24,7 +24,6 @@ import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; import static org.junit.Assert.assertTrue; -import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -32,30 +31,21 @@ import com.google.api.services.storage.model.Bucket; import com.google.api.services.storage.model.Objects; import com.google.api.services.storage.model.StorageObject; -import com.google.cloud.WriteChannel; import com.google.cloud.storage.BucketInfo; -import com.google.cloud.storage.Storage.BucketGetOption; import com.google.cloud.storage.Storage.BucketTargetOption; import com.google.cloud.storage.Storage.PredefinedAcl; import com.google.cloud.storage.StorageClass; -import com.google.cloud.storage.StorageException; import java.io.FileNotFoundException; import java.io.IOException; import java.math.BigInteger; import java.nio.channels.SeekableByteChannel; -import java.nio.channels.WritableByteChannel; -import java.util.HashMap; import java.util.List; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.TimeUnit; -import org.apache.beam.runners.core.metrics.GcpResourceIdentifiers; import org.apache.beam.runners.core.metrics.MetricsContainerImpl; -import org.apache.beam.runners.core.metrics.MonitoringInfoConstants; -import org.apache.beam.runners.core.metrics.MonitoringInfoMetricName; import org.apache.beam.sdk.extensions.gcp.auth.TestCredential; import org.apache.beam.sdk.extensions.gcp.options.GcsOptions; -import org.apache.beam.sdk.extensions.gcp.util.GcsUtil.CreateOptions; import org.apache.beam.sdk.extensions.gcp.util.GcsUtil.StorageObjectOrIOException; import org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath; import org.apache.beam.sdk.io.fs.MoveOptions.StandardMoveOptions; @@ -586,125 +576,4 @@ public void testBucketAccessibleIsRoutedToV2() throws IOException { assertFalse(gcsUtil.bucketAccessible(inaccessible)); Mockito.verifyNoMoreInteractions(mockDelegate); } - - // The tests below exercise a real GcsUtilV2 delegate whose java-storage client is mocked, to - // cover behavior that GcsUtilV2 must share with GcsUtilV1. - // TODO: Move these to a parity test that runs against both delegates. - - /** - * Returns a {@link GcsUtil} backed by a real {@link GcsUtilV2} that issues every call to {@code - * storage}. Performance metrics are off by default, so the per-operation clients of {@link - * GcsUtilV2#storageWithHttpMetrics} resolve to this one as well. - */ - private GcsUtil gcsUtilWithV2Storage(com.google.cloud.storage.Storage storage) { - GcsOptions options = gcsOptionsWithTestCredential(); - options.setProject("my_project"); - GcsUtil gcsUtil = options.getGcsUtil(); - GcsUtilV2 delegateV2 = Mockito.spy(new GcsUtilV2(options)); - Mockito.doReturn(storage).when(delegateV2).storage(); - gcsUtil.delegateV2 = delegateV2; - return gcsUtil; - } - - private void verifyMetricWasSet( - String projectId, String bucketId, String method, String status, long count) { - // Verify the metric as reported. - HashMap labels = new HashMap<>(); - labels.put(MonitoringInfoConstants.Labels.PTRANSFORM, ""); - labels.put(MonitoringInfoConstants.Labels.SERVICE, "Storage"); - labels.put(MonitoringInfoConstants.Labels.METHOD, method); - labels.put(MonitoringInfoConstants.Labels.GCS_PROJECT_ID, projectId); - labels.put(MonitoringInfoConstants.Labels.GCS_BUCKET, bucketId); - labels.put( - MonitoringInfoConstants.Labels.RESOURCE, - GcpResourceIdentifiers.cloudStorageBucket(bucketId)); - labels.put(MonitoringInfoConstants.Labels.STATUS, status); - - MonitoringInfoMetricName name = - MonitoringInfoMetricName.named(MonitoringInfoConstants.Urns.API_REQUEST_COUNT, labels); - MetricsContainerImpl container = - (MetricsContainerImpl) MetricsEnvironment.getProcessWideContainer(); - assertEquals(count, (long) container.getCounter(name).getCumulative()); - } - - /** - * Mirrors {@link GcsUtilV1Test#testGCSReadMetricsIsSet} for V2: opening a missing object records - * a {@code not_found} request, even though V2 detects it through a lookup rather than a {@link - * StorageException}. - */ - @Test - public void testV2OpenMissingObjectRecordsNotFoundMetric() { - com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); - // An unstubbed get() returns null, which is how java-storage reports a missing object. - GcsUtil gcsUtil = gcsUtilWithV2Storage(storage); - - assertThrows( - FileNotFoundException.class, - () -> gcsUtil.open(GcsPath.fromComponents("testbucket", "testobject"))); - - verifyMetricWasSet("my_project", "testbucket", "GcsGet", "not_found", 1); - verifyMetricWasSet("my_project", "testbucket", "GcsGet", "ok", 0); - } - - /** - * An upload is finalized when its channel is closed, so a failed precondition surfaces there. V2 - * must report it as an {@link IOException}, as V1 does, rather than an unchecked {@link - * StorageException}. - */ - @Test - public void testV2WriteChannelCloseTranslatesStorageException() throws IOException { - com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); - WriteChannel writer = Mockito.mock(WriteChannel.class); - StorageException preconditionFailed = new StorageException(412, "Precondition Failed"); - Mockito.doThrow(preconditionFailed).when(writer).close(); - when(storage.writer(any(com.google.cloud.storage.BlobInfo.class), any())).thenReturn(writer); - GcsUtil gcsUtil = gcsUtilWithV2Storage(storage); - - WritableByteChannel channel = - gcsUtil.create( - GcsPath.fromComponents("testbucket", "testobject"), - CreateOptions.builder().setExpectFileToNotExist(true).build()); - - IOException thrown = assertThrows(IOException.class, channel::close); - assertSame(preconditionFailed, thrown.getCause()); - } - - /** Mirrors {@link GcsUtilV1Test#testBucketDoesNotExist} for V2. */ - @Test - public void testV2BucketAccessibleIsFalseWhenBucketDoesNotExist() throws IOException { - com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); - // An unstubbed get() returns null, which is how java-storage reports a missing bucket. - GcsUtil gcsUtil = gcsUtilWithV2Storage(storage); - - assertFalse(gcsUtil.bucketAccessible(GcsPath.fromComponents("testbucket", "testobject"))); - } - - /** Mirrors {@link GcsUtilV1Test#testBucketDoesNotExistBecauseOfAccessError} for V2. */ - @Test - public void testV2BucketAccessibleIsFalseWhenAccessIsDenied() throws IOException { - com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); - when(storage.get(Mockito.eq("testbucket"), any(BucketGetOption.class))) - .thenThrow(new StorageException(403, "Forbidden")); - GcsUtil gcsUtil = gcsUtilWithV2Storage(storage); - - assertFalse(gcsUtil.bucketAccessible(GcsPath.fromComponents("testbucket", "testobject"))); - } - - /** - * Any other failure (e.g. a 5xx) says nothing about whether the bucket is accessible, so it must - * propagate rather than be reported as an inaccessible bucket, as V1 does. - */ - @Test - public void testV2BucketAccessiblePropagatesOtherFailures() { - com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); - StorageException serverError = new StorageException(503, "Service Unavailable"); - when(storage.get(Mockito.eq("testbucket"), any(BucketGetOption.class))).thenThrow(serverError); - GcsUtil gcsUtil = gcsUtilWithV2Storage(storage); - - IOException thrown = - assertThrows( - IOException.class, - () -> gcsUtil.bucketAccessible(GcsPath.fromComponents("testbucket", "testobject"))); - assertSame(serverError, thrown.getCause()); - } } diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java index 7058efa46e06..858422f2e0ec 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java @@ -19,10 +19,13 @@ import static java.nio.charset.StandardCharsets.UTF_8; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNotSame; import static org.junit.Assert.assertNull; import static org.junit.Assert.assertSame; import static org.junit.Assert.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.when; import com.google.api.client.http.GenericUrl; import com.google.api.client.http.HttpRequestInitializer; @@ -30,22 +33,32 @@ import com.google.api.client.testing.http.MockLowLevelHttpResponse; import com.google.auth.Credentials; import com.google.cloud.NoCredentials; +import com.google.cloud.WriteChannel; import com.google.cloud.hadoop.util.AsyncWriteChannelOptions; import com.google.cloud.http.HttpTransportOptions; import com.google.cloud.storage.Storage; +import com.google.cloud.storage.Storage.BucketGetOption; +import com.google.cloud.storage.StorageException; import com.google.cloud.storage.StorageOptions; import java.io.ByteArrayOutputStream; +import java.io.FileNotFoundException; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.channels.Channels; import java.nio.channels.SeekableByteChannel; import java.nio.channels.WritableByteChannel; +import java.util.HashMap; import org.apache.beam.repackaged.core.org.apache.commons.compress.utils.SeekableInMemoryByteChannel; import org.apache.beam.runners.core.metrics.CounterCell; +import org.apache.beam.runners.core.metrics.GcpResourceIdentifiers; import org.apache.beam.runners.core.metrics.MetricsContainerImpl; +import org.apache.beam.runners.core.metrics.MonitoringInfoConstants; +import org.apache.beam.runners.core.metrics.MonitoringInfoMetricName; import org.apache.beam.sdk.extensions.gcp.auth.NoopCredentialFactory; import org.apache.beam.sdk.extensions.gcp.auth.TestCredential; import org.apache.beam.sdk.extensions.gcp.options.GcsOptions; +import org.apache.beam.sdk.extensions.gcp.util.GcsUtil.CreateOptions; +import org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath; import org.apache.beam.sdk.metrics.MetricName; import org.apache.beam.sdk.metrics.MetricsEnvironment; import org.apache.beam.sdk.options.PipelineOptionsFactory; @@ -54,11 +67,12 @@ import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; +import org.mockito.Mockito; /** * Unit tests for {@link GcsUtilV2}. * - *

Covers two things without making any request to GCS: + *

Covers three things without making any request to GCS: * *

    *
  • The storage client is configured from the pipeline options (project, credentials, @@ -66,6 +80,8 @@ *
  • HTTP metrics and byte counters are installed only when their flags ask for them, and only * when there is a container to report into. The counting logic itself is covered by {@link * TransportTest}. + *
  • The results and exceptions that callers depend on match those of {@link GcsUtilV1}, for the + * responses of a mocked java-storage client. *
* *

End-to-end parity with {@link GcsUtilV1} against real GCS is covered by {@link @@ -81,6 +97,7 @@ public class GcsUtilV2Test { @After public void tearDown() { + MetricsEnvironment.setProcessWideContainer(null); MetricsEnvironment.setCurrentContainer(null); } @@ -427,4 +444,128 @@ public void testDefaultUploadChunkSizeMatchesV1() { AsyncWriteChannelOptions.DEFAULT.getUploadChunkSize(), GcsUtilV2.DEFAULT_UPLOAD_CHUNK_SIZE_BYTES); } + + // --------------------------------------------------------------------------------------------- + // Behavior shared with GcsUtilV1, through a mocked java-storage client. Each test mirrors the + // GcsUtilV1Test case it names; end-to-end parity is covered by GcsUtilParameterizedIT. + // --------------------------------------------------------------------------------------------- + + /** + * Returns a {@link GcsUtil} backed by a real {@link GcsUtilV2} that issues every call to {@code + * storage}. Performance metrics are off by default, so the per-operation clients of {@link + * GcsUtilV2#storageWithHttpMetrics} resolve to this one as well. + */ + private GcsUtil gcsUtilWithV2Storage(com.google.cloud.storage.Storage storage) { + GcsOptions options = gcsOptions(); + options.setProject("my_project"); + GcsUtil gcsUtil = options.getGcsUtil(); + GcsUtilV2 delegateV2 = Mockito.spy(new GcsUtilV2(options)); + Mockito.doReturn(storage).when(delegateV2).storage(); + gcsUtil.delegateV2 = delegateV2; + return gcsUtil; + } + + private void verifyMetricWasSet( + String projectId, String bucketId, String method, String status, long count) { + // Verify the metric as reported. + HashMap labels = new HashMap<>(); + labels.put(MonitoringInfoConstants.Labels.PTRANSFORM, ""); + labels.put(MonitoringInfoConstants.Labels.SERVICE, "Storage"); + labels.put(MonitoringInfoConstants.Labels.METHOD, method); + labels.put(MonitoringInfoConstants.Labels.GCS_PROJECT_ID, projectId); + labels.put(MonitoringInfoConstants.Labels.GCS_BUCKET, bucketId); + labels.put( + MonitoringInfoConstants.Labels.RESOURCE, + GcpResourceIdentifiers.cloudStorageBucket(bucketId)); + labels.put(MonitoringInfoConstants.Labels.STATUS, status); + + MonitoringInfoMetricName name = + MonitoringInfoMetricName.named(MonitoringInfoConstants.Urns.API_REQUEST_COUNT, labels); + MetricsContainerImpl container = + (MetricsContainerImpl) MetricsEnvironment.getProcessWideContainer(); + assertEquals(count, (long) container.getCounter(name).getCumulative()); + } + + /** + * Mirrors {@link GcsUtilV1Test#testGCSReadMetricsIsSet} for V2: opening a missing object records + * a {@code not_found} request, even though V2 detects it through a lookup rather than a {@link + * StorageException}. + */ + @Test + public void testV2OpenMissingObjectRecordsNotFoundMetric() { + MetricsContainerImpl container = new MetricsContainerImpl(null); + MetricsEnvironment.setProcessWideContainer(container); + com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); + // An unstubbed get() returns null, which is how java-storage reports a missing object. + GcsUtil gcsUtil = gcsUtilWithV2Storage(storage); + + assertThrows( + FileNotFoundException.class, + () -> gcsUtil.open(GcsPath.fromComponents("testbucket", "testobject"))); + + verifyMetricWasSet("my_project", "testbucket", "GcsGet", "not_found", 1); + verifyMetricWasSet("my_project", "testbucket", "GcsGet", "ok", 0); + } + + /** + * An upload is finalized when its channel is closed, so a failed precondition surfaces there. V2 + * must report it as an {@link IOException}, as V1 does, rather than an unchecked {@link + * StorageException}. + */ + @Test + public void testV2WriteChannelCloseTranslatesStorageException() throws IOException { + com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); + WriteChannel writer = Mockito.mock(WriteChannel.class); + StorageException preconditionFailed = new StorageException(412, "Precondition Failed"); + Mockito.doThrow(preconditionFailed).when(writer).close(); + when(storage.writer(any(com.google.cloud.storage.BlobInfo.class), any())).thenReturn(writer); + GcsUtil gcsUtil = gcsUtilWithV2Storage(storage); + + WritableByteChannel channel = + gcsUtil.create( + GcsPath.fromComponents("testbucket", "testobject"), + CreateOptions.builder().setExpectFileToNotExist(true).build()); + + IOException thrown = assertThrows(IOException.class, channel::close); + assertSame(preconditionFailed, thrown.getCause()); + } + + /** Mirrors {@link GcsUtilV1Test#testBucketDoesNotExist} for V2. */ + @Test + public void testV2BucketAccessibleIsFalseWhenBucketDoesNotExist() throws IOException { + com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); + // An unstubbed get() returns null, which is how java-storage reports a missing bucket. + GcsUtil gcsUtil = gcsUtilWithV2Storage(storage); + + assertFalse(gcsUtil.bucketAccessible(GcsPath.fromComponents("testbucket", "testobject"))); + } + + /** Mirrors {@link GcsUtilV1Test#testBucketDoesNotExistBecauseOfAccessError} for V2. */ + @Test + public void testV2BucketAccessibleIsFalseWhenAccessIsDenied() throws IOException { + com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); + when(storage.get(Mockito.eq("testbucket"), any(BucketGetOption.class))) + .thenThrow(new StorageException(403, "Forbidden")); + GcsUtil gcsUtil = gcsUtilWithV2Storage(storage); + + assertFalse(gcsUtil.bucketAccessible(GcsPath.fromComponents("testbucket", "testobject"))); + } + + /** + * Any other failure (e.g. a 5xx) says nothing about whether the bucket is accessible, so it must + * propagate rather than be reported as an inaccessible bucket, as V1 does. + */ + @Test + public void testV2BucketAccessiblePropagatesOtherFailures() { + com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); + StorageException serverError = new StorageException(503, "Service Unavailable"); + when(storage.get(Mockito.eq("testbucket"), any(BucketGetOption.class))).thenThrow(serverError); + GcsUtil gcsUtil = gcsUtilWithV2Storage(storage); + + IOException thrown = + assertThrows( + IOException.class, + () -> gcsUtil.bucketAccessible(GcsPath.fromComponents("testbucket", "testobject"))); + assertSame(serverError, thrown.getCause()); + } } From 2e35f6dfc8f6f8b3cfb03243c70774d1e717108d Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Fri, 25 Sep 2026 21:23:53 -0400 Subject: [PATCH 3/6] Move GcsUtilIT tests into GcsUtilParameterizedIT --- .../sdk/extensions/gcp/util/GcsUtilIT.java | 145 ------------------ .../gcp/util/GcsUtilParameterizedIT.java | 123 ++++++++++++++- 2 files changed, 118 insertions(+), 150 deletions(-) delete mode 100644 sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilIT.java diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilIT.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilIT.java deleted file mode 100644 index be62e07e5085..000000000000 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilIT.java +++ /dev/null @@ -1,145 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.apache.beam.sdk.extensions.gcp.util; - -import static org.hamcrest.MatcherAssert.assertThat; -import static org.hamcrest.Matchers.equalTo; -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertNotNull; - -import com.google.protobuf.ByteString; -import java.io.IOException; -import java.nio.ByteBuffer; -import java.nio.channels.ReadableByteChannel; -import java.nio.channels.WritableByteChannel; -import java.time.LocalDateTime; -import java.time.ZoneId; -import java.util.Collections; -import java.util.concurrent.atomic.AtomicInteger; -import org.apache.beam.sdk.extensions.gcp.options.GcsOptions; -import org.apache.beam.sdk.extensions.gcp.util.GcsUtil.CreateOptions; -import org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath; -import org.apache.beam.sdk.io.FileSystems; -import org.apache.beam.sdk.options.ExperimentalOptions; -import org.apache.beam.sdk.testing.TestPipeline; -import org.apache.beam.sdk.testing.TestPipelineOptions; -import org.apache.beam.sdk.testing.UsesKms; -import org.apache.beam.sdk.util.ByteStringOutputStream; -import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; -import org.junit.Test; -import org.junit.experimental.categories.Category; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; - -/** - * Integration tests for {@link GcsUtil}. These tests are designed to run against production Google - * Cloud Storage. - * - *

This is a runnerless integration test, even though the Beam IT framework assumes one. Thus, - * this test should only be run against single runner (such as DirectRunner). - */ -@RunWith(JUnit4.class) -@Category(UsesKms.class) -public class GcsUtilIT { - /** Tests a rewrite operation that requires multiple API calls (using a continuation token). */ - @Test - public void testRewriteMultiPart() throws IOException { - TestPipelineOptions options = - TestPipeline.testingPipelineOptions().as(TestPipelineOptions.class); - // Using a KMS key is necessary to trigger multi-part rewrites (bucket is created - // with a bucket default key). - assertNotNull(options.getTempRoot()); - options.setTempLocation( - FileSystems.matchNewDirectory(options.getTempRoot(), "testRewriteMultiPart").toString()); - - GcsOptions gcsOptions = options.as(GcsOptions.class); - GcsUtil gcsUtil = gcsOptions.getGcsUtil(); - String srcFilename = "gs://dataflow-samples/wikipedia_edits/wiki_data-000000000000.json"; - String dstFilename = - gcsOptions.getGcpTempLocation() - + String.format( - "/GcsUtilIT-%tF-% Date: Fri, 25 Sep 2026 21:26:22 -0400 Subject: [PATCH 4/6] Rename GcsUtilParameterizedIT to GcsUtilIT --- .../gcp/util/{GcsUtilParameterizedIT.java => GcsUtilIT.java} | 2 +- .../apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java | 5 ++--- 2 files changed, 3 insertions(+), 4 deletions(-) rename sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/{GcsUtilParameterizedIT.java => GcsUtilIT.java} (99%) diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilParameterizedIT.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilIT.java similarity index 99% rename from sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilParameterizedIT.java rename to sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilIT.java index 0e3a08b6ee16..852ec2977aaf 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilParameterizedIT.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilIT.java @@ -95,7 +95,7 @@ */ @RunWith(Parameterized.class) @Category(UsesKms.class) -public class GcsUtilParameterizedIT { +public class GcsUtilIT { private static final String READ_COUNTER_PREFIX = "it_read_bytes"; private static final String WRITE_COUNTER_PREFIX = "it_write_bytes"; diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java index 858422f2e0ec..43ed32733cb1 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java @@ -84,8 +84,7 @@ * responses of a mocked java-storage client. * * - *

End-to-end parity with {@link GcsUtilV1} against real GCS is covered by {@link - * GcsUtilParameterizedIT}. + *

End-to-end parity with {@link GcsUtilV1} against real GCS is covered by {@link GcsUtilIT}. */ @RunWith(JUnit4.class) public class GcsUtilV2Test { @@ -447,7 +446,7 @@ public void testDefaultUploadChunkSizeMatchesV1() { // --------------------------------------------------------------------------------------------- // Behavior shared with GcsUtilV1, through a mocked java-storage client. Each test mirrors the - // GcsUtilV1Test case it names; end-to-end parity is covered by GcsUtilParameterizedIT. + // GcsUtilV1Test case it names; end-to-end parity is covered by GcsUtilIT. // --------------------------------------------------------------------------------------------- /** From 221ac4b7a7b02622493ab2cf6dd45ccec4d88046 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Tue, 29 Sep 2026 21:26:21 -0400 Subject: [PATCH 5/6] Address reviewer comments --- .../beam/sdk/extensions/gcp/util/GcsUtilIT.java | 4 ++-- .../beam/sdk/extensions/gcp/util/GcsUtilTest.java | 3 --- .../sdk/extensions/gcp/util/GcsUtilV2Test.java | 15 ++++++++------- 3 files changed, 10 insertions(+), 12 deletions(-) diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilIT.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilIT.java index 852ec2977aaf..ef8f11c17159 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilIT.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilIT.java @@ -684,7 +684,7 @@ public void testRead() throws IOException, NoSuchAlgorithmException { public void testWriteAndRead() throws IOException { final String bucketName = randomBucketName(); final GcsPath targetPath = - GcsPath.fromComponents(bucketName, "test-object-" + java.util.UUID.randomUUID() + ".txt"); + GcsPath.fromComponents(bucketName, "test-object-" + UUID.randomUUID() + ".txt"); final byte[] content = "Hello, GCS!".getBytes(StandardCharsets.UTF_8); try { @@ -817,7 +817,7 @@ public void testReadMetrics() throws IOException { @Test public void testWriteMetrics() throws IOException { final String bucket = - "apache-beam-temp-metrics-" + java.util.UUID.randomUUID().toString().substring(0, 8); + "apache-beam-temp-metrics-" + UUID.randomUUID().toString().substring(0, 8); final GcsPath targetPath = GcsPath.fromComponents(bucket, "test-object.txt"); final byte[] content = "Hello, GCS metrics!".getBytes(StandardCharsets.UTF_8); GcsUtil metricsGcsUtil = gcsUtilWithAllMetrics(); diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java index 24ee159653dc..ccadd3580e42 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java @@ -54,10 +54,8 @@ import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.junit.After; import org.junit.Before; -import org.junit.Rule; import org.junit.Test; import org.junit.function.ThrowingRunnable; -import org.junit.rules.ExpectedException; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; import org.mockito.Mockito; @@ -65,7 +63,6 @@ /** Test case for {@link GcsUtil}. */ @RunWith(JUnit4.class) public class GcsUtilTest { - @Rule public ExpectedException thrown = ExpectedException.none(); MetricsContainerImpl testMetricsContainer; @Before diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java index 43ed32733cb1..67c189e0c262 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java @@ -36,6 +36,7 @@ import com.google.cloud.WriteChannel; import com.google.cloud.hadoop.util.AsyncWriteChannelOptions; import com.google.cloud.http.HttpTransportOptions; +import com.google.cloud.storage.BlobInfo; import com.google.cloud.storage.Storage; import com.google.cloud.storage.Storage.BucketGetOption; import com.google.cloud.storage.StorageException; @@ -454,7 +455,7 @@ public void testDefaultUploadChunkSizeMatchesV1() { * storage}. Performance metrics are off by default, so the per-operation clients of {@link * GcsUtilV2#storageWithHttpMetrics} resolve to this one as well. */ - private GcsUtil gcsUtilWithV2Storage(com.google.cloud.storage.Storage storage) { + private GcsUtil gcsUtilWithV2Storage(Storage storage) { GcsOptions options = gcsOptions(); options.setProject("my_project"); GcsUtil gcsUtil = options.getGcsUtil(); @@ -494,7 +495,7 @@ private void verifyMetricWasSet( public void testV2OpenMissingObjectRecordsNotFoundMetric() { MetricsContainerImpl container = new MetricsContainerImpl(null); MetricsEnvironment.setProcessWideContainer(container); - com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); + Storage storage = Mockito.mock(Storage.class); // An unstubbed get() returns null, which is how java-storage reports a missing object. GcsUtil gcsUtil = gcsUtilWithV2Storage(storage); @@ -513,11 +514,11 @@ public void testV2OpenMissingObjectRecordsNotFoundMetric() { */ @Test public void testV2WriteChannelCloseTranslatesStorageException() throws IOException { - com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); + Storage storage = Mockito.mock(Storage.class); WriteChannel writer = Mockito.mock(WriteChannel.class); StorageException preconditionFailed = new StorageException(412, "Precondition Failed"); Mockito.doThrow(preconditionFailed).when(writer).close(); - when(storage.writer(any(com.google.cloud.storage.BlobInfo.class), any())).thenReturn(writer); + when(storage.writer(any(BlobInfo.class), any())).thenReturn(writer); GcsUtil gcsUtil = gcsUtilWithV2Storage(storage); WritableByteChannel channel = @@ -532,7 +533,7 @@ public void testV2WriteChannelCloseTranslatesStorageException() throws IOExcepti /** Mirrors {@link GcsUtilV1Test#testBucketDoesNotExist} for V2. */ @Test public void testV2BucketAccessibleIsFalseWhenBucketDoesNotExist() throws IOException { - com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); + Storage storage = Mockito.mock(Storage.class); // An unstubbed get() returns null, which is how java-storage reports a missing bucket. GcsUtil gcsUtil = gcsUtilWithV2Storage(storage); @@ -542,7 +543,7 @@ public void testV2BucketAccessibleIsFalseWhenBucketDoesNotExist() throws IOExcep /** Mirrors {@link GcsUtilV1Test#testBucketDoesNotExistBecauseOfAccessError} for V2. */ @Test public void testV2BucketAccessibleIsFalseWhenAccessIsDenied() throws IOException { - com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); + Storage storage = Mockito.mock(Storage.class); when(storage.get(Mockito.eq("testbucket"), any(BucketGetOption.class))) .thenThrow(new StorageException(403, "Forbidden")); GcsUtil gcsUtil = gcsUtilWithV2Storage(storage); @@ -556,7 +557,7 @@ public void testV2BucketAccessibleIsFalseWhenAccessIsDenied() throws IOException */ @Test public void testV2BucketAccessiblePropagatesOtherFailures() { - com.google.cloud.storage.Storage storage = Mockito.mock(com.google.cloud.storage.Storage.class); + Storage storage = Mockito.mock(Storage.class); StorageException serverError = new StorageException(503, "Service Unavailable"); when(storage.get(Mockito.eq("testbucket"), any(BucketGetOption.class))).thenThrow(serverError); GcsUtil gcsUtil = gcsUtilWithV2Storage(storage); From 63a59518cb0dd8cbbf0b4252ac782ae98347189a Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Tue, 29 Sep 2026 21:41:49 -0400 Subject: [PATCH 6/6] Trim long qualified name --- .../sdk/extensions/gcp/util/GcsUtilTest.java | 36 ++++++++++--------- 1 file changed, 20 insertions(+), 16 deletions(-) diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java index ccadd3580e42..f3f894ceec3e 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilTest.java @@ -31,6 +31,7 @@ import com.google.api.services.storage.model.Bucket; import com.google.api.services.storage.model.Objects; import com.google.api.services.storage.model.StorageObject; +import com.google.cloud.storage.Blob; import com.google.cloud.storage.BucketInfo; import com.google.cloud.storage.Storage.BucketTargetOption; import com.google.cloud.storage.Storage.PredefinedAcl; @@ -39,6 +40,9 @@ import java.io.IOException; import java.math.BigInteger; import java.nio.channels.SeekableByteChannel; +import java.time.Duration; +import java.time.Instant; +import java.time.ZoneOffset; import java.util.List; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; @@ -148,8 +152,8 @@ private GcsUtil gcsUtilRoutingToV2() { return gcsUtil; } - private static com.google.cloud.storage.Blob mockBlob(String bucket, String object) { - com.google.cloud.storage.Blob blob = Mockito.mock(com.google.cloud.storage.Blob.class); + private static Blob mockBlob(String bucket, String object) { + Blob blob = Mockito.mock(Blob.class); when(blob.getBucket()).thenReturn(bucket); when(blob.getName()).thenReturn(object); return blob; @@ -224,7 +228,7 @@ public void testRemoveIsRoutedToV2AndToleratesMissingFiles() throws IOException @Test public void testGetObjectsIsRoutedToV2AndConvertsBlobs() throws IOException { GcsUtil gcsUtil = gcsUtilRoutingToV2(); - com.google.cloud.storage.Blob blob = mockBlob("bucket", "found"); + Blob blob = mockBlob("bucket", "found"); FileNotFoundException notFound = new FileNotFoundException("gs://bucket/missing"); List paths = ImmutableList.of(GcsPath.fromUri("gs://bucket/found"), GcsPath.fromUri("gs://bucket/miss")); @@ -248,11 +252,11 @@ public void testGetObjectsIsRoutedToV2AndConvertsBlobs() throws IOException { @Test public void testListObjectsIsRoutedToV2AndConvertsAPage() throws IOException { GcsUtil gcsUtil = gcsUtilRoutingToV2(); - com.google.cloud.storage.Blob object = mockBlob("bucket", "prefix/object"); - com.google.cloud.storage.Blob directory = mockBlob("bucket", "prefix/dir/"); + Blob object = mockBlob("bucket", "prefix/object"); + Blob directory = mockBlob("bucket", "prefix/dir/"); when(directory.isDirectory()).thenReturn(true); @SuppressWarnings("unchecked") - Page page = Mockito.mock(Page.class); + Page page = Mockito.mock(Page.class); when(page.getValues()).thenReturn(ImmutableList.of(object, directory)); when(page.hasNextPage()).thenReturn(true); when(page.getNextPageToken()).thenReturn("next"); @@ -271,7 +275,7 @@ public void testListObjectsIsRoutedToV2AndConvertsAPage() throws IOException { public void testListObjectsReportsTheLastPageWithANullToken() throws IOException { GcsUtil gcsUtil = gcsUtilRoutingToV2(); @SuppressWarnings("unchecked") - Page page = Mockito.mock(Page.class); + Page page = Mockito.mock(Page.class); when(page.getValues()).thenReturn(ImmutableList.of()); // A gax page reports an empty token rather than a null one once it is exhausted. Callers of // listObjects loop until the token is null, so it has to be normalized. @@ -312,7 +316,7 @@ public void testCreateBucketIsRoutedToV2WithProjectPrivateAcls() throws IOExcept .setStorageClass(StorageClass.NEARLINE) .setSoftDeletePolicy( BucketInfo.SoftDeletePolicy.newBuilder() - .setRetentionDuration(java.time.Duration.ZERO) + .setRetentionDuration(Duration.ZERO) .build()) .build(), BucketTargetOption.predefinedAcl(PredefinedAcl.PROJECT_PRIVATE), @@ -344,7 +348,7 @@ public void testBucketOwnerIsRoutedToV2BucketProject() throws IOException { public void testGetObjectIsRoutedToV2AndKeepsAllFields() throws IOException { GcsUtil gcsUtil = gcsUtilRoutingToV2(); GcsPath path = GcsPath.fromUri("gs://bucket/object"); - com.google.cloud.storage.Blob blob = mockBlob("bucket", "object"); + Blob blob = mockBlob("bucket", "object"); when(blob.getSize()).thenReturn(42L); when(blob.getGeneration()).thenReturn(7L); when(blob.getMetageneration()).thenReturn(3L); @@ -354,9 +358,9 @@ public void testGetObjectIsRoutedToV2AndKeepsAllFields() throws IOException { when(blob.getCrc32c()).thenReturn("crc=="); when(blob.getEtag()).thenReturn("etag"); when(blob.getUpdateTimeOffsetDateTime()) - .thenReturn(java.time.Instant.ofEpochMilli(1234L).atOffset(java.time.ZoneOffset.UTC)); + .thenReturn(Instant.ofEpochMilli(1234L).atOffset(ZoneOffset.UTC)); when(blob.getCreateTimeOffsetDateTime()) - .thenReturn(java.time.Instant.ofEpochMilli(1000L).atOffset(java.time.ZoneOffset.UTC)); + .thenReturn(Instant.ofEpochMilli(1000L).atOffset(ZoneOffset.UTC)); when(mockDelegateV2.getBlob(path)).thenReturn(blob); StorageObject object = gcsUtil.getObject(path); @@ -381,7 +385,7 @@ public void testGetObjectIsRoutedToV2AndKeepsAllFields() throws IOException { public void testGetObjectLeavesMissingFieldsUnsetForV2() throws IOException { GcsUtil gcsUtil = gcsUtilRoutingToV2(); GcsPath path = GcsPath.fromUri("gs://bucket/object"); - com.google.cloud.storage.Blob blob = mockBlob("bucket", "object"); + Blob blob = mockBlob("bucket", "object"); // Mockito would otherwise answer 0 for the boxed size. when(blob.getSize()).thenReturn(null); when(mockDelegateV2.getBlob(path)).thenReturn(blob); @@ -427,7 +431,7 @@ public void testGetBucketIsRoutedToV2AndConvertsTheBucket() throws IOException { when(bucket.getSoftDeletePolicy()) .thenReturn( BucketInfo.SoftDeletePolicy.newBuilder() - .setRetentionDuration(java.time.Duration.ofDays(7)) + .setRetentionDuration(Duration.ofDays(7)) .build()); when(mockDelegateV2.getBucket(path)).thenReturn(bucket); @@ -439,7 +443,7 @@ public void testGetBucketIsRoutedToV2AndConvertsTheBucket() throws IOException { assertEquals(BigInteger.valueOf(123L), converted.getProjectNumber()); assertEquals("NEARLINE", converted.getStorageClass()); assertEquals( - Long.valueOf(java.time.Duration.ofDays(7).getSeconds()), + Long.valueOf(Duration.ofDays(7).getSeconds()), converted.getSoftDeletePolicy().getRetentionDurationSeconds()); Mockito.verifyNoMoreInteractions(mockDelegate); } @@ -527,9 +531,9 @@ public void testFileSizeIsRoutedToV2() throws IOException { @Test public void testListObjectsWithDelimiterIsRoutedToV2() throws IOException { GcsUtil gcsUtil = gcsUtilRoutingToV2(); - com.google.cloud.storage.Blob object = mockBlob("bucket", "prefix/object"); + Blob object = mockBlob("bucket", "prefix/object"); @SuppressWarnings("unchecked") - Page page = Mockito.mock(Page.class); + Page page = Mockito.mock(Page.class); when(page.getValues()).thenReturn(ImmutableList.of(object)); when(mockDelegateV2.listBlobs("bucket", "prefix/", "token", "/")).thenReturn(page);