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 be62e07e5085..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 @@ -19,32 +19,72 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.equalTo; +import static org.junit.Assert.assertArrayEquals; 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.assertThrows; +import static org.junit.Assert.assertTrue; +import static org.junit.Assume.assumeTrue; +import com.google.api.gax.paging.Page; +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.StorageChannelUtils; import com.google.protobuf.ByteString; +import java.io.ByteArrayOutputStream; +import java.io.FileNotFoundException; import java.io.IOException; import java.nio.ByteBuffer; import java.nio.channels.ReadableByteChannel; +import java.nio.channels.SeekableByteChannel; import java.nio.channels.WritableByteChannel; +import java.nio.charset.StandardCharsets; +import java.nio.file.AccessDeniedException; +import java.nio.file.FileAlreadyExistsException; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; import java.time.LocalDateTime; import java.time.ZoneId; +import java.util.Arrays; import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.UUID; import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; +import org.apache.beam.runners.core.metrics.CounterCell; +import org.apache.beam.runners.core.metrics.GcpResourceIdentifiers; +import org.apache.beam.runners.core.metrics.MetricUpdates.MetricUpdate; +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.options.GcsOptions; import org.apache.beam.sdk.extensions.gcp.util.GcsUtil.CreateOptions; +import org.apache.beam.sdk.extensions.gcp.util.GcsUtilV2.MissingStrategy; +import org.apache.beam.sdk.extensions.gcp.util.GcsUtilV2.OverwriteStrategy; import org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath; import org.apache.beam.sdk.io.FileSystems; +import org.apache.beam.sdk.io.fs.MoveOptions; +import org.apache.beam.sdk.metrics.MetricName; +import org.apache.beam.sdk.metrics.MetricsEnvironment; 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.Before; import org.junit.Test; import org.junit.experimental.categories.Category; import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameter; +import org.junit.runners.Parameterized.Parameters; /** * Integration tests for {@link GcsUtil}. These tests are designed to run against production Google @@ -53,12 +93,778 @@ *

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) +@RunWith(Parameterized.class) @Category(UsesKms.class) public class GcsUtilIT { + + private static final String READ_COUNTER_PREFIX = "it_read_bytes"; + private static final String WRITE_COUNTER_PREFIX = "it_write_bytes"; + + @Parameters(name = "{0}") + public static Iterable data() { + return Arrays.asList("use_gcsutil_v1", "use_gcsutil_v2"); + } + + @Parameter public String experiment; + + private TestPipelineOptions options; + private GcsUtil gcsUtil; + + @Before + public void setUp() { + options = TestPipeline.testingPipelineOptions().as(TestPipelineOptions.class); + + // set the experimental flag. + ExperimentalOptions experimentalOptions = options.as(ExperimentalOptions.class); + experimentalOptions.setExperiments(Collections.singletonList(experiment)); + + GcsOptions gcsOptions = options.as(GcsOptions.class); + gcsUtil = gcsOptions.getGcsUtil(); + } + + /** Returns a bucket name unique to this test run, so concurrent runs don't collide. */ + private static String randomBucketName() { + return "apache-beam-temp-bucket-" + UUID.randomUUID(); + } + + @Test + public void testFileSize() throws IOException { + final GcsPath gcsPath = GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kinglear.txt"); + final long expectedSize = 157283L; + + assertEquals(expectedSize, gcsUtil.fileSize(gcsPath)); + } + + @Test + public void testGetObjectOrGetBlob() throws IOException { + final GcsPath existingPath = + GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kinglear.txt"); + final String expectedCRC = "s0a3Tg=="; + + String crc; + if (experiment.equals("use_gcsutil_v2")) { + Blob blob = gcsUtil.getBlob(existingPath); + crc = blob.getCrc32c(); + } else { + StorageObject obj = gcsUtil.getObject(existingPath); + crc = obj.getCrc32c(); + } + assertEquals(expectedCRC, crc); + + final GcsPath nonExistentPath = + GcsPath.fromUri("gs://my-random-test-bucket-12345/unknown-12345.txt"); + final GcsPath forbiddenPath = GcsPath.fromUri("gs://test-bucket/unknown-12345.txt"); + + if (experiment.equals("use_gcsutil_v2")) { + assertThrows(FileNotFoundException.class, () -> gcsUtil.getBlob(nonExistentPath)); + // For V2, we are returning AccessDeniedException (a subclass of IOException) for forbidden + // paths. + assertThrows(AccessDeniedException.class, () -> gcsUtil.getBlob(forbiddenPath)); + } else { + assertThrows(FileNotFoundException.class, () -> gcsUtil.getObject(nonExistentPath)); + assertThrows(IOException.class, () -> gcsUtil.getObject(forbiddenPath)); + } + } + + @Test + public void testGetObjectsOrGetBlobs() throws IOException { + final GcsPath existingPath = + GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kinglear.txt"); + final GcsPath nonExistentPath = + GcsPath.fromUri("gs://my-random-test-bucket-12345/unknown-12345.txt"); + final List paths = Arrays.asList(existingPath, nonExistentPath); + + if (experiment.equals("use_gcsutil_v2")) { + List results = gcsUtil.getBlobs(paths); + assertEquals(2, results.size()); + assertTrue(results.get(0).blob() != null); + assertTrue(results.get(0).ioException() == null); + assertTrue(results.get(1).blob() == null); + assertTrue(results.get(1).ioException() != null); + } else { + List results = gcsUtil.getObjects(paths); + assertEquals(2, results.size()); + assertTrue(results.get(0).storageObject() != null); + assertTrue(results.get(0).ioException() == null); + assertTrue(results.get(1).storageObject() == null); + assertTrue(results.get(1).ioException() != null); + } + } + + @Test + public void testListObjectsOrListBlobs() throws IOException { + final String bucket = "apache-beam-samples"; + final String prefix = "shakespeare/kingrichard"; + + List names; + if (experiment.equals("use_gcsutil_v2")) { + Page blobs = gcsUtil.listBlobs(bucket, prefix, null); + names = blobs.streamAll().map(blob -> blob.getName()).collect(Collectors.toList()); + } else { + Objects objs = gcsUtil.listObjects(bucket, prefix, null); + names = objs.getItems().stream().map(obj -> obj.getName()).collect(Collectors.toList()); + } + assertEquals( + Arrays.asList("shakespeare/kingrichardii.txt", "shakespeare/kingrichardiii.txt"), names); + + final String randomPrefix = "my-random-prefix/random"; + if (experiment.equals("use_gcsutil_v2")) { + Page blobs = gcsUtil.listBlobs(bucket, randomPrefix, null); + assertEquals(0, blobs.streamAll().count()); + } else { + Objects objs = gcsUtil.listObjects(bucket, randomPrefix, null); + assertEquals(null, objs.getItems()); + } + } + + @Test + public void testExpand() throws IOException { + final GcsPath existingPattern = + GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kingrichardii*.txt"); + List paths = gcsUtil.expand(existingPattern); + + assertEquals( + Arrays.asList( + GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kingrichardii.txt"), + GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kingrichardiii.txt")), + paths); + + final GcsPath nonExistentPattern1 = + GcsPath.fromUri("gs://apache-beam-samples/my_random_folder/random*.txt"); + assertTrue(gcsUtil.expand(nonExistentPattern1).isEmpty()); + + final GcsPath nonExistentPattern2 = + GcsPath.fromUri("gs://apache-beam-samples/shakespeare/king*.csv"); + assertTrue(gcsUtil.expand(nonExistentPattern2).isEmpty()); + } + + @Test + public void testGetBucketOrGetBucketWithOptions() throws IOException { + final GcsPath existingPath = GcsPath.fromUri("gs://apache-beam-samples"); + + String bucket; + if (experiment.equals("use_gcsutil_v2")) { + bucket = gcsUtil.getBucketWithOptions(existingPath).getName(); + } else { + bucket = gcsUtil.getBucket(existingPath).getName(); + } + assertEquals("apache-beam-samples", bucket); + + final GcsPath nonExistentPath = GcsPath.fromUri("gs://my-random-test-bucket-12345"); + final GcsPath forbiddenPath = GcsPath.fromUri("gs://test-bucket"); + + if (experiment.equals("use_gcsutil_v2")) { + assertThrows( + FileNotFoundException.class, () -> gcsUtil.getBucketWithOptions(nonExistentPath)); + assertThrows(AccessDeniedException.class, () -> gcsUtil.getBucketWithOptions(forbiddenPath)); + } else { + assertThrows(FileNotFoundException.class, () -> gcsUtil.getBucket(nonExistentPath)); + assertThrows(AccessDeniedException.class, () -> gcsUtil.getBucket(forbiddenPath)); + } + } + + @Test + public void testBucketAccessible() throws IOException { + final GcsPath existingPath = GcsPath.fromUri("gs://apache-beam-samples"); + final GcsPath nonExistentPath = GcsPath.fromUri("gs://my-random-test-bucket-12345"); + final GcsPath forbiddenPath = GcsPath.fromUri("gs://test-bucket"); + + assertEquals(true, gcsUtil.bucketAccessible(existingPath)); + assertEquals(false, gcsUtil.bucketAccessible(nonExistentPath)); + assertEquals(false, gcsUtil.bucketAccessible(forbiddenPath)); + } + + @Test + public void testBucketOwner() throws IOException { + final GcsPath existingPath = GcsPath.fromUri("gs://apache-beam-samples"); + final long expectedProjectNumber = 844138762903L; // apache-beam-testing + assertEquals(expectedProjectNumber, gcsUtil.bucketOwner(existingPath)); + + final GcsPath nonExistentPath = GcsPath.fromUri("gs://my-random-test-bucket-12345"); + final GcsPath forbiddenPath = GcsPath.fromUri("gs://test-bucket"); + assertThrows(FileNotFoundException.class, () -> gcsUtil.bucketOwner(nonExistentPath)); + assertThrows(AccessDeniedException.class, () -> gcsUtil.bucketOwner(forbiddenPath)); + } + + @Test + public void testCreateAndRemoveBucket() throws IOException { + final GcsPath gcsPath = GcsPath.fromUri("gs://" + randomBucketName()); + + if (experiment.equals("use_gcsutil_v2")) { + BucketInfo bucketInfo = BucketInfo.of(gcsPath.getBucket()); + try { + assertFalse(gcsUtil.bucketAccessible(gcsPath)); + gcsUtil.createBucket(bucketInfo); + assertTrue(gcsUtil.bucketAccessible(gcsPath)); + + // raise exception when the bucket already exists during creation + assertThrows(FileAlreadyExistsException.class, () -> gcsUtil.createBucket(bucketInfo)); + + assertTrue(gcsUtil.bucketAccessible(gcsPath)); + gcsUtil.removeBucket(bucketInfo); + assertFalse(gcsUtil.bucketAccessible(gcsPath)); + + // raise exception when the bucket does not exist during removal + assertThrows(FileNotFoundException.class, () -> gcsUtil.removeBucket(bucketInfo)); + } finally { + // clean up and ignore errors no matter what + try { + gcsUtil.removeBucket(bucketInfo); + } catch (IOException e) { + } + } + } else { + Bucket bucket = new Bucket().setName(gcsPath.getBucket()); + GcsOptions gcsOptions = options.as(GcsOptions.class); + String projectId = gcsOptions.getProject(); + try { + assertFalse(gcsUtil.bucketAccessible(gcsPath)); + gcsUtil.createBucket(projectId, bucket); + assertTrue(gcsUtil.bucketAccessible(gcsPath)); + + // raise exception when the bucket already exists during creation + assertThrows( + FileAlreadyExistsException.class, () -> gcsUtil.createBucket(projectId, bucket)); + + assertTrue(gcsUtil.bucketAccessible(gcsPath)); + gcsUtil.removeBucket(bucket); + assertFalse(gcsUtil.bucketAccessible(gcsPath)); + + // raise exception when the bucket does not exist during removal + assertThrows(FileNotFoundException.class, () -> gcsUtil.removeBucket(bucket)); + } finally { + // clean up and ignore errors no matter what + try { + gcsUtil.removeBucket(bucket); + } catch (IOException e) { + } + } + } + } + + private List createTestBucketHelper(String bucketName, boolean copyData) + throws IOException { + final List originPaths = + Arrays.asList( + GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kingrichardii.txt"), + GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kingrichardiii.txt")); + + final List testPaths = + originPaths.stream() + .map(o -> GcsPath.fromComponents(bucketName, o.getObject())) + .collect(Collectors.toList()); + + // create bucket and copy some initial files into there + if (experiment.equals("use_gcsutil_v2")) { + gcsUtil.createBucket(BucketInfo.of(bucketName)); + + if (copyData) { + gcsUtil.copyV2(originPaths, testPaths); + } else { + return Collections.emptyList(); + } + } else { + GcsOptions gcsOptions = options.as(GcsOptions.class); + gcsUtil.createBucket(gcsOptions.getProject(), new Bucket().setName(bucketName)); + + if (copyData) { + final List originList = + originPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); + final List testList = + testPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); + gcsUtil.copy(originList, testList); + } else { + return Collections.emptyList(); + } + } + + return testPaths; + } + + private void tearDownTestBucketHelper(String bucketName) { + try { + // use "**" in the pattern to match any characters including "/". + final List paths = + gcsUtil.expand(GcsPath.fromUri(String.format("gs://%s/**", bucketName))); + if (experiment.equals("use_gcsutil_v2")) { + gcsUtil.remove(paths, MissingStrategy.SKIP_IF_MISSING); + gcsUtil.removeBucket(BucketInfo.of(bucketName)); + } else { + gcsUtil.remove(paths.stream().map(GcsPath::toString).collect(Collectors.toList())); + gcsUtil.removeBucket(new Bucket().setName(bucketName)); + } + } catch (IOException e) { + System.err.println( + "Error during tear down of test bucket " + bucketName + ": " + e.getMessage()); + } + } + + @Test + public void testCopy() throws IOException { + final String existingBucket = randomBucketName(); + final String nonExistentBucket = "my-random-test-bucket-12345"; + + try { + final List srcPaths = createTestBucketHelper(existingBucket, true); + final List dstPaths = + srcPaths.stream() + .map(o -> GcsPath.fromComponents(existingBucket, o.getObject() + ".bak")) + .collect(Collectors.toList()); + final List errPaths = + srcPaths.stream() + .map(o -> GcsPath.fromComponents(nonExistentBucket, o.getObject())) + .collect(Collectors.toList()); + + assertNotExists(dstPaths.get(0)); + assertNotExists(dstPaths.get(1)); + + if (experiment.equals("use_gcsutil_v2")) { + // (1) when the target files do not exist + gcsUtil.copyV2(srcPaths, dstPaths); + assertExists(dstPaths.get(0)); + assertExists(dstPaths.get(1)); + + // (2) when the target files exist + // (2a) no exception on SAFE_OVERWRITE, ALWAYS_OVERWRITE, SKIP_IF_EXISTS + gcsUtil.copyV2(srcPaths, dstPaths); + gcsUtil.copy(srcPaths, dstPaths, OverwriteStrategy.ALWAYS_OVERWRITE); + gcsUtil.copy(srcPaths, dstPaths, OverwriteStrategy.SKIP_IF_EXISTS); + + // (2b) raise exception on FAIL_IF_EXISTS + assertThrows( + FileAlreadyExistsException.class, + () -> gcsUtil.copy(srcPaths, dstPaths, OverwriteStrategy.FAIL_IF_EXISTS)); + + // (3) raise exception when the target bucket is nonexistent. + assertThrows(FileNotFoundException.class, () -> gcsUtil.copyV2(srcPaths, errPaths)); + + // (4) raise exception when the source files are nonexistent. + assertThrows(FileNotFoundException.class, () -> gcsUtil.copyV2(errPaths, dstPaths)); + } else { + final List srcList = + srcPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); + final List dstList = + dstPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); + final List errList = + errPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); + + // (1) when the target files do not exist + gcsUtil.copy(srcList, dstList); + assertExists(dstPaths.get(0)); + assertExists(dstPaths.get(1)); + + // (2) when the target files exist, no exception + gcsUtil.copy(srcList, dstList); + + // (3) raise exception when the target bucket is nonexistent. + assertThrows(FileNotFoundException.class, () -> gcsUtil.copy(srcList, errList)); + + // (4) raise exception when the source files are nonexistent. + assertThrows(FileNotFoundException.class, () -> gcsUtil.copy(errList, dstList)); + } + } finally { + tearDownTestBucketHelper(existingBucket); + } + } + + @Test + public void testRemove() throws IOException { + final String existingBucket = randomBucketName(); + final String nonExistentBucket = "my-random-test-bucket-12345"; + + try { + final List srcPaths = createTestBucketHelper(existingBucket, true); + final List errPaths = + srcPaths.stream() + .map(o -> GcsPath.fromComponents(nonExistentBucket, o.getObject())) + .collect(Collectors.toList()); + + assertExists(srcPaths.get(0)); + assertExists(srcPaths.get(1)); + + if (experiment.equals("use_gcsutil_v2")) { + // (1) when the files to remove exist + gcsUtil.removeV2(srcPaths); + assertNotExists(srcPaths.get(0)); + assertNotExists(srcPaths.get(1)); + + // (2) when the files to remove have been deleted + // (2a) no exception on SKIP_IF_MISSING + gcsUtil.removeV2(srcPaths); + gcsUtil.remove(srcPaths, MissingStrategy.SKIP_IF_MISSING); + + // (2b) raise exception on FAIL_IF_MISSING + assertThrows( + FileNotFoundException.class, + () -> gcsUtil.remove(srcPaths, MissingStrategy.FAIL_IF_MISSING)); + + // (3) when the files are from an nonexistent bucket + // (3a) no exception on SKIP_IF_MISSING + gcsUtil.removeV2(errPaths); + gcsUtil.remove(errPaths, MissingStrategy.SKIP_IF_MISSING); + + // (3b) raise exception on FAIL_IF_MISSING + assertThrows( + FileNotFoundException.class, + () -> gcsUtil.remove(errPaths, MissingStrategy.FAIL_IF_MISSING)); + } else { + final List srcList = + srcPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); + final List errList = + errPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); + + // (1) when the files to remove exist + gcsUtil.remove(srcList); + assertNotExists(srcPaths.get(0)); + assertNotExists(srcPaths.get(1)); + + // (2) when the files to remove have been deleted, no exception + gcsUtil.remove(srcList); + + // (3) when the files are from an nonexistent bucket, no exception + gcsUtil.remove(errList); + } + } finally { + tearDownTestBucketHelper(existingBucket); + } + } + + @Test + public void testRename() throws IOException { + final String existingBucket = randomBucketName(); + final String nonExistentBucket = "my-random-test-bucket-12345"; + + try { + final List srcPaths = createTestBucketHelper(existingBucket, true); + final List tmpPaths = + srcPaths.stream() + .map(o -> GcsPath.fromComponents(existingBucket, "tmp/" + o.getObject())) + .collect(Collectors.toList()); + final List dstPaths = + srcPaths.stream() + .map(o -> GcsPath.fromComponents(existingBucket, o.getObject() + ".bak")) + .collect(Collectors.toList()); + final List errPaths = + srcPaths.stream() + .map(o -> GcsPath.fromComponents(nonExistentBucket, o.getObject())) + .collect(Collectors.toList()); + + assertNotExists(dstPaths.get(0)); + assertNotExists(dstPaths.get(1)); + if (experiment.equals("use_gcsutil_v2")) { + // Make a copy of sources + gcsUtil.copyV2(srcPaths, tmpPaths); + + // (1) when the source files exist and target files do not + gcsUtil.renameV2(tmpPaths, dstPaths); + assertNotExists(tmpPaths.get(0)); + assertNotExists(tmpPaths.get(1)); + assertExists(dstPaths.get(0)); + assertExists(dstPaths.get(1)); + + // (2) when the source files do not exist + // (2a) no exception if IGNORE_MISSING_FILES is set + gcsUtil.renameV2(errPaths, dstPaths, MoveOptions.StandardMoveOptions.IGNORE_MISSING_FILES); + + // (2b) raise exception if if IGNORE_MISSING_FILES is not set + assertThrows(FileNotFoundException.class, () -> gcsUtil.renameV2(errPaths, dstPaths)); + + // (3) when both source files and target files exist + gcsUtil.renameV2( + srcPaths, dstPaths, MoveOptions.StandardMoveOptions.SKIP_IF_DESTINATION_EXISTS); + gcsUtil.renameV2(srcPaths, dstPaths); + } else { + final List srcList = + srcPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); + final List tmpList = + tmpPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); + final List dstList = + dstPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); + final List errList = + errPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); + + // Make a copy of sources + gcsUtil.copy(srcList, tmpList); + + // (1) when the source files exist and target files do not + gcsUtil.rename(tmpList, dstList); + assertNotExists(tmpPaths.get(0)); + assertNotExists(tmpPaths.get(1)); + assertExists(dstPaths.get(0)); + assertExists(dstPaths.get(1)); + + // (2) when the source files do not exist + // (2a) no exception if IGNORE_MISSING_FILES is set + gcsUtil.rename(errList, dstList, MoveOptions.StandardMoveOptions.IGNORE_MISSING_FILES); + + // (2b) raise exception if if IGNORE_MISSING_FILES is not set + assertThrows(FileNotFoundException.class, () -> gcsUtil.rename(errList, dstList)); + + // (3) when both source files and target files exist + assertExists(srcPaths.get(0)); + assertExists(srcPaths.get(1)); + assertExists(dstPaths.get(0)); + assertExists(dstPaths.get(1)); + + // There is a bug in V1 where SKIP_IF_DESTINATION_EXISTS is not honored. + gcsUtil.rename( + srcList, dstList, MoveOptions.StandardMoveOptions.SKIP_IF_DESTINATION_EXISTS); + + assertNotExists(srcPaths.get(0)); // BUG! The renaming is supposed to be skipped + assertNotExists(srcPaths.get(1)); // BUG! The renaming is supposed to be skipped + // assertExists(srcPaths.get(0)); + // assertExists(srcPaths.get(1)); + assertExists(dstPaths.get(0)); + assertExists(dstPaths.get(1)); + } + } finally { + tearDownTestBucketHelper(existingBucket); + } + } + + private void assertExists(GcsPath path) throws IOException { + if (experiment.equals("use_gcsutil_v2")) { + gcsUtil.getBlob(path); + } else { + gcsUtil.getObject(path); + } + } + + private void assertNotExists(GcsPath path) throws IOException { + if (experiment.equals("use_gcsutil_v2")) { + assertThrows(FileNotFoundException.class, () -> gcsUtil.getBlob(path)); + } else { + assertThrows(FileNotFoundException.class, () -> gcsUtil.getObject(path)); + } + } + + String computeHash(ByteBuffer buffer) throws NoSuchAlgorithmException { + MessageDigest digest = MessageDigest.getInstance("SHA-256"); + digest.update(buffer); + byte[] hashBytes = digest.digest(); + + // Convert bytes to Hex String + StringBuilder sb = new StringBuilder(); + for (byte b : hashBytes) { + sb.append(String.format("%02x", b)); + } + return sb.toString(); + } + + @Test + public void testRead() throws IOException, NoSuchAlgorithmException { + final GcsPath gcsPath = GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kinglear.txt"); + final String expectedHash = "674a2725884307c96398440497c889ad8cecccedf5689df85e6b0faabe4e0fe8"; + final long expectedSize = 157283L; + + try (SeekableByteChannel channel = gcsUtil.open(gcsPath)) { + // Verify Size + assertEquals(expectedSize, channel.size()); + assertEquals(0, channel.position()); + + // Read content into ByteBuffer. + // Allocate a larger buffer to ensure we receive the EOF at the expected place. + ByteBuffer buffer = ByteBuffer.allocate((int) expectedSize + 1024); + int bytesRead = StorageChannelUtils.blockingFillFrom(buffer, channel); + + // Verify total bytes read and position + assertEquals(expectedSize, bytesRead); + assertEquals(expectedSize, channel.position()); + + // Flip the buffer to prepare it for reading (sets limit to current position, position to 0) + buffer.flip(); + + // Verify hash + String actualHash = computeHash(buffer); + assertEquals("Content hash should match", expectedHash, actualHash); + } + } + + @Test + public void testWriteAndRead() throws IOException { + final String bucketName = randomBucketName(); + final GcsPath targetPath = + GcsPath.fromComponents(bucketName, "test-object-" + UUID.randomUUID() + ".txt"); + final byte[] content = "Hello, GCS!".getBytes(StandardCharsets.UTF_8); + + try { + createTestBucketHelper(bucketName, false); + + // Write content to a GCS file + CreateOptions options = CreateOptions.builder().setExpectFileToNotExist(true).build(); + try (WritableByteChannel writer = gcsUtil.create(targetPath, options)) { + writer.write(ByteBuffer.wrap(content)); + } + + // Read content into a buffer + ByteArrayOutputStream readContent = new ByteArrayOutputStream(); + try (ReadableByteChannel reader = gcsUtil.open(targetPath)) { + ByteBuffer buffer = ByteBuffer.allocate(1024); + while (reader.read(buffer) != -1) { + buffer.flip(); + readContent.write(buffer.array(), 0, buffer.limit()); + buffer.clear(); + } + } + + // Verify content + assertArrayEquals(content, readContent.toByteArray()); + } finally { + tearDownTestBucketHelper(bucketName); + } + } + + // --------------------------------------------------------------------------------------------- + // Metrics parity: the same assertions must hold for both GcsUtilV1 and GcsUtilV2. + // --------------------------------------------------------------------------------------------- + + /** Returns a {@link GcsUtil} with every GCS metric flag turned on. */ + private GcsUtil gcsUtilWithAllMetrics() { + GcsOptions gcsOptions = options.as(GcsOptions.class); + gcsOptions.setGcsPerformanceMetrics(true); + gcsOptions.setEnableBucketReadMetricCounter(true); + gcsOptions.setEnableBucketWriteMetricCounter(true); + gcsOptions.setGcsReadCounterPrefix(READ_COUNTER_PREFIX); + gcsOptions.setGcsWriteCounterPrefix(WRITE_COUNTER_PREFIX); + // Built directly, as getGcsUtil() returns the instance cached in setUp(). + return new GcsUtil(gcsOptions); + } + + private static long counter(MetricsContainerImpl container, MetricName name) { + CounterCell cell = container.tryGetCounter(name); + assertNotNull("counter " + name + " was not reported", cell); + return cell.getCumulative(); + } + + private static long gcsCounter(MetricsContainerImpl container, String name) { + return counter(container, MetricName.named(GcsUtil.METRIC_NAMESPACE, name)); + } + + private static long bucketCounter(MetricsContainerImpl container, String prefix, String bucket) { + return counter(container, MetricName.named(GcsUtil.class, prefix + "_" + bucket)); + } + + /** + * Sums the API request counter for {@code method} and {@code status} on {@code bucket}. + * + *

The {@code GCS_PROJECT_ID} label is deliberately not matched: V1 reports the project of its + * gcsio options (which is never set, so it reports {@code "null"}), while V2 reports the + * pipeline's project. + */ + private static long apiRequestCount( + MetricsContainerImpl container, String method, String status, String bucket) { + long total = 0; + for (MetricUpdate update : container.getCumulative().counterUpdates()) { + MetricName name = update.getKey().metricName(); + if (!(name instanceof MonitoringInfoMetricName)) { + continue; + } + MonitoringInfoMetricName miName = (MonitoringInfoMetricName) name; + Map labels = miName.getLabels(); + if (MonitoringInfoConstants.Urns.API_REQUEST_COUNT.equals(miName.getUrn()) + && "Storage".equals(labels.get(MonitoringInfoConstants.Labels.SERVICE)) + && method.equals(labels.get(MonitoringInfoConstants.Labels.METHOD)) + && status.equals(labels.get(MonitoringInfoConstants.Labels.STATUS)) + && bucket.equals(labels.get(MonitoringInfoConstants.Labels.GCS_BUCKET)) + && GcpResourceIdentifiers.cloudStorageBucket(bucket) + .equals(labels.get(MonitoringInfoConstants.Labels.RESOURCE))) { + total += update.getUpdate(); + } + } + return total; + } + + @Test + public void testReadMetrics() throws IOException { + final String bucket = "apache-beam-samples"; + final GcsPath gcsPath = GcsPath.fromComponents(bucket, "shakespeare/kinglear.txt"); + final long expectedSize = 157283L; + GcsUtil metricsGcsUtil = gcsUtilWithAllMetrics(); + + MetricsContainerImpl container = new MetricsContainerImpl("step"); + MetricsContainerImpl processWide = new MetricsContainerImpl(null); + MetricsEnvironment.setCurrentContainer(container); + MetricsEnvironment.setProcessWideContainer(processWide); + try { + try (SeekableByteChannel channel = metricsGcsUtil.open(gcsPath)) { + ByteBuffer buffer = ByteBuffer.allocate((int) expectedSize + 1024); + assertEquals(expectedSize, StorageChannelUtils.blockingFillFrom(buffer, channel)); + } + } finally { + MetricsEnvironment.setCurrentContainer(null); + MetricsEnvironment.setProcessWideContainer(null); + } + + // --enableBucketReadMetricCounter + assertEquals(expectedSize, bucketCounter(container, READ_COUNTER_PREFIX, bucket)); + // --gcsPerformanceMetrics: wire bytes and HTTP counters + assertEquals(expectedSize, gcsCounter(container, "gcs_http_read_wire_bytes_received")); + assertTrue(gcsCounter(container, "gcs_http_read_request_count") >= 1); + assertTrue(gcsCounter(container, "gcs_http_read_status_2xx") >= 1); + assertEquals( + gcsCounter(container, "gcs_http_read_request_count"), + gcsCounter(container, "gcs_http_read_request_count_ranged") + + gcsCounter(container, "gcs_http_read_request_count_unbounded") + + gcsCounter(container, "gcs_http_read_request_count_other")); + // A read must not produce write-side counters. + assertNull( + container.tryGetCounter( + MetricName.named(GcsUtil.METRIC_NAMESPACE, "gcs_http_write_request_count"))); + // API request metric + assertEquals(1, apiRequestCount(processWide, "GcsGet", "ok", bucket)); + } + + @Test + public void testWriteMetrics() throws IOException { + final String bucket = + "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(); + + MetricsContainerImpl container = new MetricsContainerImpl("step"); + MetricsContainerImpl processWide = new MetricsContainerImpl(null); + try { + createTestBucketHelper(bucket, false); + + MetricsEnvironment.setCurrentContainer(container); + MetricsEnvironment.setProcessWideContainer(processWide); + try (WritableByteChannel writer = + metricsGcsUtil.create( + targetPath, CreateOptions.builder().setExpectFileToNotExist(true).build())) { + writer.write(ByteBuffer.wrap(content)); + } finally { + MetricsEnvironment.setCurrentContainer(null); + MetricsEnvironment.setProcessWideContainer(null); + } + } finally { + tearDownTestBucketHelper(bucket); + } + + // --enableBucketWriteMetricCounter + assertEquals(content.length, bucketCounter(container, WRITE_COUNTER_PREFIX, bucket)); + // --gcsPerformanceMetrics: wire bytes and HTTP counters + assertEquals(content.length, gcsCounter(container, "gcs_http_write_wire_bytes_sent")); + assertTrue(gcsCounter(container, "gcs_http_write_request_count") >= 1); + assertTrue(gcsCounter(container, "gcs_http_write_status_2xx") >= 1); + // A write must not produce read-side counters. + assertNull( + container.tryGetCounter( + MetricName.named(GcsUtil.METRIC_NAMESPACE, "gcs_http_read_request_count"))); + // API request metric + assertEquals(1, apiRequestCount(processWide, "GcsInsert", "ok", bucket)); + } + + // --------------------------------------------------------------------------------------------- + // V1-only tests. + // --------------------------------------------------------------------------------------------- + /** Tests a rewrite operation that requires multiple API calls (using a continuation token). */ @Test public void testRewriteMultiPart() throws IOException { + // V2 copies each file with a single call, without rewrite tokens. + assumeTrue(experiment.equals("use_gcsutil_v1")); + TestPipelineOptions options = TestPipeline.testingPipelineOptions().as(TestPipelineOptions.class); // Using a KMS key is necessary to trigger multi-part rewrites (bucket is created @@ -93,6 +899,9 @@ public void testRewriteMultiPart() throws IOException { // the failed case. The interface of GcsGrpcOptions can also be removed. @Test public void testWriteAndReadGcsWithGrpc() throws IOException { + // GcsUtilV2 does not support gRPC yet. + assumeTrue(experiment.equals("use_gcsutil_v1")); + final String outputPattern = "%s/GcsUtilIT-%tF-%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(Parameterized.class) -@Category(UsesKms.class) -public class GcsUtilParameterizedIT { - - private static final String READ_COUNTER_PREFIX = "it_read_bytes"; - private static final String WRITE_COUNTER_PREFIX = "it_write_bytes"; - - @Parameters(name = "{0}") - public static Iterable data() { - return Arrays.asList("use_gcsutil_v1", "use_gcsutil_v2"); - } - - @Parameter public String experiment; - - private TestPipelineOptions options; - private GcsUtil gcsUtil; - - @Before - public void setUp() { - options = TestPipeline.testingPipelineOptions().as(TestPipelineOptions.class); - - // set the experimental flag. - ExperimentalOptions experimentalOptions = options.as(ExperimentalOptions.class); - experimentalOptions.setExperiments(Collections.singletonList(experiment)); - - GcsOptions gcsOptions = options.as(GcsOptions.class); - gcsUtil = gcsOptions.getGcsUtil(); - } - - @Test - public void testFileSize() throws IOException { - final GcsPath gcsPath = GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kinglear.txt"); - final long expectedSize = 157283L; - - assertEquals(expectedSize, gcsUtil.fileSize(gcsPath)); - } - - @Test - public void testGetObjectOrGetBlob() throws IOException { - final GcsPath existingPath = - GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kinglear.txt"); - final String expectedCRC = "s0a3Tg=="; - - String crc; - if (experiment.equals("use_gcsutil_v2")) { - Blob blob = gcsUtil.getBlob(existingPath); - crc = blob.getCrc32c(); - } else { - StorageObject obj = gcsUtil.getObject(existingPath); - crc = obj.getCrc32c(); - } - assertEquals(expectedCRC, crc); - - final GcsPath nonExistentPath = - GcsPath.fromUri("gs://my-random-test-bucket-12345/unknown-12345.txt"); - final GcsPath forbiddenPath = GcsPath.fromUri("gs://test-bucket/unknown-12345.txt"); - - if (experiment.equals("use_gcsutil_v2")) { - assertThrows(FileNotFoundException.class, () -> gcsUtil.getBlob(nonExistentPath)); - // For V2, we are returning AccessDeniedException (a subclass of IOException) for forbidden - // paths. - assertThrows(AccessDeniedException.class, () -> gcsUtil.getBlob(forbiddenPath)); - } else { - assertThrows(FileNotFoundException.class, () -> gcsUtil.getObject(nonExistentPath)); - assertThrows(IOException.class, () -> gcsUtil.getObject(forbiddenPath)); - } - } - - @Test - public void testGetObjectsOrGetBlobs() throws IOException { - final GcsPath existingPath = - GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kinglear.txt"); - final GcsPath nonExistentPath = - GcsPath.fromUri("gs://my-random-test-bucket-12345/unknown-12345.txt"); - final List paths = Arrays.asList(existingPath, nonExistentPath); - - if (experiment.equals("use_gcsutil_v2")) { - List results = gcsUtil.getBlobs(paths); - assertEquals(2, results.size()); - assertTrue(results.get(0).blob() != null); - assertTrue(results.get(0).ioException() == null); - assertTrue(results.get(1).blob() == null); - assertTrue(results.get(1).ioException() != null); - } else { - List results = gcsUtil.getObjects(paths); - assertEquals(2, results.size()); - assertTrue(results.get(0).storageObject() != null); - assertTrue(results.get(0).ioException() == null); - assertTrue(results.get(1).storageObject() == null); - assertTrue(results.get(1).ioException() != null); - } - } - - @Test - public void testListObjectsOrListBlobs() throws IOException { - final String bucket = "apache-beam-samples"; - final String prefix = "shakespeare/kingrichard"; - - List names; - if (experiment.equals("use_gcsutil_v2")) { - Page blobs = gcsUtil.listBlobs(bucket, prefix, null); - names = blobs.streamAll().map(blob -> blob.getName()).collect(Collectors.toList()); - } else { - Objects objs = gcsUtil.listObjects(bucket, prefix, null); - names = objs.getItems().stream().map(obj -> obj.getName()).collect(Collectors.toList()); - } - assertEquals( - Arrays.asList("shakespeare/kingrichardii.txt", "shakespeare/kingrichardiii.txt"), names); - - final String randomPrefix = "my-random-prefix/random"; - if (experiment.equals("use_gcsutil_v2")) { - Page blobs = gcsUtil.listBlobs(bucket, randomPrefix, null); - assertEquals(0, blobs.streamAll().count()); - } else { - Objects objs = gcsUtil.listObjects(bucket, randomPrefix, null); - assertEquals(null, objs.getItems()); - } - } - - @Test - public void testExpand() throws IOException { - final GcsPath existingPattern = - GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kingrichardii*.txt"); - List paths = gcsUtil.expand(existingPattern); - - assertEquals( - Arrays.asList( - GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kingrichardii.txt"), - GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kingrichardiii.txt")), - paths); - - final GcsPath nonExistentPattern1 = - GcsPath.fromUri("gs://apache-beam-samples/my_random_folder/random*.txt"); - assertTrue(gcsUtil.expand(nonExistentPattern1).isEmpty()); - - final GcsPath nonExistentPattern2 = - GcsPath.fromUri("gs://apache-beam-samples/shakespeare/king*.csv"); - assertTrue(gcsUtil.expand(nonExistentPattern2).isEmpty()); - } - - @Test - public void testGetBucketOrGetBucketWithOptions() throws IOException { - final GcsPath existingPath = GcsPath.fromUri("gs://apache-beam-samples"); - - String bucket; - if (experiment.equals("use_gcsutil_v2")) { - bucket = gcsUtil.getBucketWithOptions(existingPath).getName(); - } else { - bucket = gcsUtil.getBucket(existingPath).getName(); - } - assertEquals("apache-beam-samples", bucket); - - final GcsPath nonExistentPath = GcsPath.fromUri("gs://my-random-test-bucket-12345"); - final GcsPath forbiddenPath = GcsPath.fromUri("gs://test-bucket"); - - if (experiment.equals("use_gcsutil_v2")) { - assertThrows( - FileNotFoundException.class, () -> gcsUtil.getBucketWithOptions(nonExistentPath)); - assertThrows(AccessDeniedException.class, () -> gcsUtil.getBucketWithOptions(forbiddenPath)); - } else { - assertThrows(FileNotFoundException.class, () -> gcsUtil.getBucket(nonExistentPath)); - assertThrows(AccessDeniedException.class, () -> gcsUtil.getBucket(forbiddenPath)); - } - } - - @Test - public void testBucketAccessible() throws IOException { - final GcsPath existingPath = GcsPath.fromUri("gs://apache-beam-samples"); - final GcsPath nonExistentPath = GcsPath.fromUri("gs://my-random-test-bucket-12345"); - final GcsPath forbiddenPath = GcsPath.fromUri("gs://test-bucket"); - - assertEquals(true, gcsUtil.bucketAccessible(existingPath)); - assertEquals(false, gcsUtil.bucketAccessible(nonExistentPath)); - assertEquals(false, gcsUtil.bucketAccessible(forbiddenPath)); - } - - @Test - public void testBucketOwner() throws IOException { - final GcsPath existingPath = GcsPath.fromUri("gs://apache-beam-samples"); - final long expectedProjectNumber = 844138762903L; // apache-beam-testing - assertEquals(expectedProjectNumber, gcsUtil.bucketOwner(existingPath)); - - final GcsPath nonExistentPath = GcsPath.fromUri("gs://my-random-test-bucket-12345"); - final GcsPath forbiddenPath = GcsPath.fromUri("gs://test-bucket"); - assertThrows(FileNotFoundException.class, () -> gcsUtil.bucketOwner(nonExistentPath)); - assertThrows(AccessDeniedException.class, () -> gcsUtil.bucketOwner(forbiddenPath)); - } - - @Test - public void testCreateAndRemoveBucket() throws IOException { - final GcsPath gcsPath = GcsPath.fromUri("gs://apache-beam-test-bucket-12345"); - - if (experiment.equals("use_gcsutil_v2")) { - BucketInfo bucketInfo = BucketInfo.of(gcsPath.getBucket()); - try { - assertFalse(gcsUtil.bucketAccessible(gcsPath)); - gcsUtil.createBucket(bucketInfo); - assertTrue(gcsUtil.bucketAccessible(gcsPath)); - - // raise exception when the bucket already exists during creation - assertThrows(FileAlreadyExistsException.class, () -> gcsUtil.createBucket(bucketInfo)); - - assertTrue(gcsUtil.bucketAccessible(gcsPath)); - gcsUtil.removeBucket(bucketInfo); - assertFalse(gcsUtil.bucketAccessible(gcsPath)); - - // raise exception when the bucket does not exist during removal - assertThrows(FileNotFoundException.class, () -> gcsUtil.removeBucket(bucketInfo)); - } finally { - // clean up and ignore errors no matter what - try { - gcsUtil.removeBucket(bucketInfo); - } catch (IOException e) { - } - } - } else { - Bucket bucket = new Bucket().setName(gcsPath.getBucket()); - GcsOptions gcsOptions = options.as(GcsOptions.class); - String projectId = gcsOptions.getProject(); - try { - assertFalse(gcsUtil.bucketAccessible(gcsPath)); - gcsUtil.createBucket(projectId, bucket); - assertTrue(gcsUtil.bucketAccessible(gcsPath)); - - // raise exception when the bucket already exists during creation - assertThrows( - FileAlreadyExistsException.class, () -> gcsUtil.createBucket(projectId, bucket)); - - assertTrue(gcsUtil.bucketAccessible(gcsPath)); - gcsUtil.removeBucket(bucket); - assertFalse(gcsUtil.bucketAccessible(gcsPath)); - - // raise exception when the bucket does not exist during removal - assertThrows(FileNotFoundException.class, () -> gcsUtil.removeBucket(bucket)); - } finally { - // clean up and ignore errors no matter what - try { - gcsUtil.removeBucket(bucket); - } catch (IOException e) { - } - } - } - } - - private List createTestBucketHelper(String bucketName, boolean copyData) - throws IOException { - final List originPaths = - Arrays.asList( - GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kingrichardii.txt"), - GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kingrichardiii.txt")); - - final List testPaths = - originPaths.stream() - .map(o -> GcsPath.fromComponents(bucketName, o.getObject())) - .collect(Collectors.toList()); - - // create bucket and copy some initial files into there - if (experiment.equals("use_gcsutil_v2")) { - gcsUtil.createBucket(BucketInfo.of(bucketName)); - - if (copyData) { - gcsUtil.copyV2(originPaths, testPaths); - } else { - return Collections.emptyList(); - } - } else { - GcsOptions gcsOptions = options.as(GcsOptions.class); - gcsUtil.createBucket(gcsOptions.getProject(), new Bucket().setName(bucketName)); - - if (copyData) { - final List originList = - originPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); - final List testList = - testPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); - gcsUtil.copy(originList, testList); - } else { - return Collections.emptyList(); - } - } - - return testPaths; - } - - private void tearDownTestBucketHelper(String bucketName) { - try { - // use "**" in the pattern to match any characters including "/". - final List paths = - gcsUtil.expand(GcsPath.fromUri(String.format("gs://%s/**", bucketName))); - if (experiment.equals("use_gcsutil_v2")) { - gcsUtil.remove(paths, MissingStrategy.SKIP_IF_MISSING); - gcsUtil.removeBucket(BucketInfo.of(bucketName)); - } else { - gcsUtil.remove(paths.stream().map(GcsPath::toString).collect(Collectors.toList())); - gcsUtil.removeBucket(new Bucket().setName(bucketName)); - } - } catch (IOException e) { - System.err.println( - "Error during tear down of test bucket " + bucketName + ": " + e.getMessage()); - } - } - - @Test - public void testCopy() throws IOException { - final String existingBucket = "apache-beam-temp-bucket-12345"; - final String nonExistentBucket = "my-random-test-bucket-12345"; - - try { - final List srcPaths = createTestBucketHelper(existingBucket, true); - final List dstPaths = - srcPaths.stream() - .map(o -> GcsPath.fromComponents(existingBucket, o.getObject() + ".bak")) - .collect(Collectors.toList()); - final List errPaths = - srcPaths.stream() - .map(o -> GcsPath.fromComponents(nonExistentBucket, o.getObject())) - .collect(Collectors.toList()); - - assertNotExists(dstPaths.get(0)); - assertNotExists(dstPaths.get(1)); - - if (experiment.equals("use_gcsutil_v2")) { - // (1) when the target files do not exist - gcsUtil.copyV2(srcPaths, dstPaths); - assertExists(dstPaths.get(0)); - assertExists(dstPaths.get(1)); - - // (2) when the target files exist - // (2a) no exception on SAFE_OVERWRITE, ALWAYS_OVERWRITE, SKIP_IF_EXISTS - gcsUtil.copyV2(srcPaths, dstPaths); - gcsUtil.copy(srcPaths, dstPaths, OverwriteStrategy.ALWAYS_OVERWRITE); - gcsUtil.copy(srcPaths, dstPaths, OverwriteStrategy.SKIP_IF_EXISTS); - - // (2b) raise exception on FAIL_IF_EXISTS - assertThrows( - FileAlreadyExistsException.class, - () -> gcsUtil.copy(srcPaths, dstPaths, OverwriteStrategy.FAIL_IF_EXISTS)); - - // (3) raise exception when the target bucket is nonexistent. - assertThrows(FileNotFoundException.class, () -> gcsUtil.copyV2(srcPaths, errPaths)); - - // (4) raise exception when the source files are nonexistent. - assertThrows(FileNotFoundException.class, () -> gcsUtil.copyV2(errPaths, dstPaths)); - } else { - final List srcList = - srcPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); - final List dstList = - dstPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); - final List errList = - errPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); - - // (1) when the target files do not exist - gcsUtil.copy(srcList, dstList); - assertExists(dstPaths.get(0)); - assertExists(dstPaths.get(1)); - - // (2) when the target files exist, no exception - gcsUtil.copy(srcList, dstList); - - // (3) raise exception when the target bucket is nonexistent. - assertThrows(FileNotFoundException.class, () -> gcsUtil.copy(srcList, errList)); - - // (4) raise exception when the source files are nonexistent. - assertThrows(FileNotFoundException.class, () -> gcsUtil.copy(errList, dstList)); - } - } finally { - tearDownTestBucketHelper(existingBucket); - } - } - - @Test - public void testRemove() throws IOException { - final String existingBucket = "apache-beam-temp-bucket-12345"; - final String nonExistentBucket = "my-random-test-bucket-12345"; - - try { - final List srcPaths = createTestBucketHelper(existingBucket, true); - final List errPaths = - srcPaths.stream() - .map(o -> GcsPath.fromComponents(nonExistentBucket, o.getObject())) - .collect(Collectors.toList()); - - assertExists(srcPaths.get(0)); - assertExists(srcPaths.get(1)); - - if (experiment.equals("use_gcsutil_v2")) { - // (1) when the files to remove exist - gcsUtil.removeV2(srcPaths); - assertNotExists(srcPaths.get(0)); - assertNotExists(srcPaths.get(1)); - - // (2) when the files to remove have been deleted - // (2a) no exception on SKIP_IF_MISSING - gcsUtil.removeV2(srcPaths); - gcsUtil.remove(srcPaths, MissingStrategy.SKIP_IF_MISSING); - - // (2b) raise exception on FAIL_IF_MISSING - assertThrows( - FileNotFoundException.class, - () -> gcsUtil.remove(srcPaths, MissingStrategy.FAIL_IF_MISSING)); - - // (3) when the files are from an nonexistent bucket - // (3a) no exception on SKIP_IF_MISSING - gcsUtil.removeV2(errPaths); - gcsUtil.remove(errPaths, MissingStrategy.SKIP_IF_MISSING); - - // (3b) raise exception on FAIL_IF_MISSING - assertThrows( - FileNotFoundException.class, - () -> gcsUtil.remove(errPaths, MissingStrategy.FAIL_IF_MISSING)); - } else { - final List srcList = - srcPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); - final List errList = - errPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); - - // (1) when the files to remove exist - gcsUtil.remove(srcList); - assertNotExists(srcPaths.get(0)); - assertNotExists(srcPaths.get(1)); - - // (2) when the files to remove have been deleted, no exception - gcsUtil.remove(srcList); - - // (3) when the files are from an nonexistent bucket, no exception - gcsUtil.remove(errList); - } - } finally { - tearDownTestBucketHelper(existingBucket); - } - } - - @Test - public void testRename() throws IOException { - final String existingBucket = "apache-beam-temp-bucket-12345"; - final String nonExistentBucket = "my-random-test-bucket-12345"; - - try { - final List srcPaths = createTestBucketHelper(existingBucket, true); - final List tmpPaths = - srcPaths.stream() - .map(o -> GcsPath.fromComponents(existingBucket, "tmp/" + o.getObject())) - .collect(Collectors.toList()); - final List dstPaths = - srcPaths.stream() - .map(o -> GcsPath.fromComponents(existingBucket, o.getObject() + ".bak")) - .collect(Collectors.toList()); - final List errPaths = - srcPaths.stream() - .map(o -> GcsPath.fromComponents(nonExistentBucket, o.getObject())) - .collect(Collectors.toList()); - - assertNotExists(dstPaths.get(0)); - assertNotExists(dstPaths.get(1)); - if (experiment.equals("use_gcsutil_v2")) { - // Make a copy of sources - gcsUtil.copyV2(srcPaths, tmpPaths); - - // (1) when the source files exist and target files do not - gcsUtil.renameV2(tmpPaths, dstPaths); - assertNotExists(tmpPaths.get(0)); - assertNotExists(tmpPaths.get(1)); - assertExists(dstPaths.get(0)); - assertExists(dstPaths.get(1)); - - // (2) when the source files do not exist - // (2a) no exception if IGNORE_MISSING_FILES is set - gcsUtil.renameV2(errPaths, dstPaths, MoveOptions.StandardMoveOptions.IGNORE_MISSING_FILES); - - // (2b) raise exception if if IGNORE_MISSING_FILES is not set - assertThrows(FileNotFoundException.class, () -> gcsUtil.renameV2(errPaths, dstPaths)); - - // (3) when both source files and target files exist - gcsUtil.renameV2( - srcPaths, dstPaths, MoveOptions.StandardMoveOptions.SKIP_IF_DESTINATION_EXISTS); - gcsUtil.renameV2(srcPaths, dstPaths); - } else { - final List srcList = - srcPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); - final List tmpList = - tmpPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); - final List dstList = - dstPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); - final List errList = - errPaths.stream().map(o -> o.toString()).collect(Collectors.toList()); - - // Make a copy of sources - gcsUtil.copy(srcList, tmpList); - - // (1) when the source files exist and target files do not - gcsUtil.rename(tmpList, dstList); - assertNotExists(tmpPaths.get(0)); - assertNotExists(tmpPaths.get(1)); - assertExists(dstPaths.get(0)); - assertExists(dstPaths.get(1)); - - // (2) when the source files do not exist - // (2a) no exception if IGNORE_MISSING_FILES is set - gcsUtil.rename(errList, dstList, MoveOptions.StandardMoveOptions.IGNORE_MISSING_FILES); - - // (2b) raise exception if if IGNORE_MISSING_FILES is not set - assertThrows(FileNotFoundException.class, () -> gcsUtil.rename(errList, dstList)); - - // (3) when both source files and target files exist - assertExists(srcPaths.get(0)); - assertExists(srcPaths.get(1)); - assertExists(dstPaths.get(0)); - assertExists(dstPaths.get(1)); - - // There is a bug in V1 where SKIP_IF_DESTINATION_EXISTS is not honored. - gcsUtil.rename( - srcList, dstList, MoveOptions.StandardMoveOptions.SKIP_IF_DESTINATION_EXISTS); - - assertNotExists(srcPaths.get(0)); // BUG! The renaming is supposed to be skipped - assertNotExists(srcPaths.get(1)); // BUG! The renaming is supposed to be skipped - // assertExists(srcPaths.get(0)); - // assertExists(srcPaths.get(1)); - assertExists(dstPaths.get(0)); - assertExists(dstPaths.get(1)); - } - } finally { - tearDownTestBucketHelper(existingBucket); - } - } - - private void assertExists(GcsPath path) throws IOException { - if (experiment.equals("use_gcsutil_v2")) { - gcsUtil.getBlob(path); - } else { - gcsUtil.getObject(path); - } - } - - private void assertNotExists(GcsPath path) throws IOException { - if (experiment.equals("use_gcsutil_v2")) { - assertThrows(FileNotFoundException.class, () -> gcsUtil.getBlob(path)); - } else { - assertThrows(FileNotFoundException.class, () -> gcsUtil.getObject(path)); - } - } - - String computeHash(ByteBuffer buffer) throws NoSuchAlgorithmException { - MessageDigest digest = MessageDigest.getInstance("SHA-256"); - digest.update(buffer); - byte[] hashBytes = digest.digest(); - - // Convert bytes to Hex String - StringBuilder sb = new StringBuilder(); - for (byte b : hashBytes) { - sb.append(String.format("%02x", b)); - } - return sb.toString(); - } - - @Test - public void testRead() throws IOException, NoSuchAlgorithmException { - final GcsPath gcsPath = GcsPath.fromUri("gs://apache-beam-samples/shakespeare/kinglear.txt"); - final String expectedHash = "674a2725884307c96398440497c889ad8cecccedf5689df85e6b0faabe4e0fe8"; - final long expectedSize = 157283L; - - try (SeekableByteChannel channel = gcsUtil.open(gcsPath)) { - // Verify Size - assertEquals(expectedSize, channel.size()); - assertEquals(0, channel.position()); - - // Read content into ByteBuffer. - // Allocate a larger buffer to ensure we receive the EOF at the expected place. - ByteBuffer buffer = ByteBuffer.allocate((int) expectedSize + 1024); - int bytesRead = StorageChannelUtils.blockingFillFrom(buffer, channel); - - // Verify total bytes read and position - assertEquals(expectedSize, bytesRead); - assertEquals(expectedSize, channel.position()); - - // Flip the buffer to prepare it for reading (sets limit to current position, position to 0) - buffer.flip(); - - // Verify hash - String actualHash = computeHash(buffer); - assertEquals("Content hash should match", expectedHash, actualHash); - } - } - - @Test - public void testWriteAndRead() throws IOException { - final String bucketName = "apache-beam-temp-bucket-12345"; - final GcsPath targetPath = - GcsPath.fromComponents(bucketName, "test-object-" + java.util.UUID.randomUUID() + ".txt"); - final byte[] content = "Hello, GCS!".getBytes(StandardCharsets.UTF_8); - - try { - createTestBucketHelper(bucketName, false); - - // Write content to a GCS file - CreateOptions options = CreateOptions.builder().setExpectFileToNotExist(true).build(); - try (WritableByteChannel writer = gcsUtil.create(targetPath, options)) { - writer.write(ByteBuffer.wrap(content)); - } - - // Read content into a buffer - ByteArrayOutputStream readContent = new ByteArrayOutputStream(); - try (ReadableByteChannel reader = gcsUtil.open(targetPath)) { - ByteBuffer buffer = ByteBuffer.allocate(1024); - while (reader.read(buffer) != -1) { - buffer.flip(); - readContent.write(buffer.array(), 0, buffer.limit()); - buffer.clear(); - } - } - - // Verify content - assertArrayEquals(content, readContent.toByteArray()); - } finally { - tearDownTestBucketHelper(bucketName); - } - } - - // --------------------------------------------------------------------------------------------- - // Metrics parity: the same assertions must hold for both GcsUtilV1 and GcsUtilV2. - // --------------------------------------------------------------------------------------------- - - /** Returns a {@link GcsUtil} with every GCS metric flag turned on. */ - private GcsUtil gcsUtilWithAllMetrics() { - GcsOptions gcsOptions = options.as(GcsOptions.class); - gcsOptions.setGcsPerformanceMetrics(true); - gcsOptions.setEnableBucketReadMetricCounter(true); - gcsOptions.setEnableBucketWriteMetricCounter(true); - gcsOptions.setGcsReadCounterPrefix(READ_COUNTER_PREFIX); - gcsOptions.setGcsWriteCounterPrefix(WRITE_COUNTER_PREFIX); - // Built directly, as getGcsUtil() returns the instance cached in setUp(). - return new GcsUtil(gcsOptions); - } - - private static long counter(MetricsContainerImpl container, MetricName name) { - CounterCell cell = container.tryGetCounter(name); - assertNotNull("counter " + name + " was not reported", cell); - return cell.getCumulative(); - } - - private static long gcsCounter(MetricsContainerImpl container, String name) { - return counter(container, MetricName.named(GcsUtil.METRIC_NAMESPACE, name)); - } - - private static long bucketCounter(MetricsContainerImpl container, String prefix, String bucket) { - return counter(container, MetricName.named(GcsUtil.class, prefix + "_" + bucket)); - } - - /** - * Sums the API request counter for {@code method} and {@code status} on {@code bucket}. - * - *

The {@code GCS_PROJECT_ID} label is deliberately not matched: V1 reports the project of its - * gcsio options (which is never set, so it reports {@code "null"}), while V2 reports the - * pipeline's project. - */ - private static long apiRequestCount( - MetricsContainerImpl container, String method, String status, String bucket) { - long total = 0; - for (MetricUpdate update : container.getCumulative().counterUpdates()) { - MetricName name = update.getKey().metricName(); - if (!(name instanceof MonitoringInfoMetricName)) { - continue; - } - MonitoringInfoMetricName miName = (MonitoringInfoMetricName) name; - Map labels = miName.getLabels(); - if (MonitoringInfoConstants.Urns.API_REQUEST_COUNT.equals(miName.getUrn()) - && "Storage".equals(labels.get(MonitoringInfoConstants.Labels.SERVICE)) - && method.equals(labels.get(MonitoringInfoConstants.Labels.METHOD)) - && status.equals(labels.get(MonitoringInfoConstants.Labels.STATUS)) - && bucket.equals(labels.get(MonitoringInfoConstants.Labels.GCS_BUCKET)) - && GcpResourceIdentifiers.cloudStorageBucket(bucket) - .equals(labels.get(MonitoringInfoConstants.Labels.RESOURCE))) { - total += update.getUpdate(); - } - } - return total; - } - - @Test - public void testReadMetrics() throws IOException { - final String bucket = "apache-beam-samples"; - final GcsPath gcsPath = GcsPath.fromComponents(bucket, "shakespeare/kinglear.txt"); - final long expectedSize = 157283L; - GcsUtil metricsGcsUtil = gcsUtilWithAllMetrics(); - - MetricsContainerImpl container = new MetricsContainerImpl("step"); - MetricsContainerImpl processWide = new MetricsContainerImpl(null); - MetricsEnvironment.setCurrentContainer(container); - MetricsEnvironment.setProcessWideContainer(processWide); - try { - try (SeekableByteChannel channel = metricsGcsUtil.open(gcsPath)) { - ByteBuffer buffer = ByteBuffer.allocate((int) expectedSize + 1024); - assertEquals(expectedSize, StorageChannelUtils.blockingFillFrom(buffer, channel)); - } - } finally { - MetricsEnvironment.setCurrentContainer(null); - MetricsEnvironment.setProcessWideContainer(null); - } - - // --enableBucketReadMetricCounter - assertEquals(expectedSize, bucketCounter(container, READ_COUNTER_PREFIX, bucket)); - // --gcsPerformanceMetrics: wire bytes and HTTP counters - assertEquals(expectedSize, gcsCounter(container, "gcs_http_read_wire_bytes_received")); - assertTrue(gcsCounter(container, "gcs_http_read_request_count") >= 1); - assertTrue(gcsCounter(container, "gcs_http_read_status_2xx") >= 1); - assertEquals( - gcsCounter(container, "gcs_http_read_request_count"), - gcsCounter(container, "gcs_http_read_request_count_ranged") - + gcsCounter(container, "gcs_http_read_request_count_unbounded") - + gcsCounter(container, "gcs_http_read_request_count_other")); - // A read must not produce write-side counters. - assertNull( - container.tryGetCounter( - MetricName.named(GcsUtil.METRIC_NAMESPACE, "gcs_http_write_request_count"))); - // API request metric - assertEquals(1, apiRequestCount(processWide, "GcsGet", "ok", bucket)); - } - - @Test - public void testWriteMetrics() throws IOException { - final String bucket = - "apache-beam-temp-metrics-" + java.util.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(); - - MetricsContainerImpl container = new MetricsContainerImpl("step"); - MetricsContainerImpl processWide = new MetricsContainerImpl(null); - try { - createTestBucketHelper(bucket, false); - - MetricsEnvironment.setCurrentContainer(container); - MetricsEnvironment.setProcessWideContainer(processWide); - try (WritableByteChannel writer = - metricsGcsUtil.create( - targetPath, CreateOptions.builder().setExpectFileToNotExist(true).build())) { - writer.write(ByteBuffer.wrap(content)); - } finally { - MetricsEnvironment.setCurrentContainer(null); - MetricsEnvironment.setProcessWideContainer(null); - } - } finally { - tearDownTestBucketHelper(bucket); - } - - // --enableBucketWriteMetricCounter - assertEquals(content.length, bucketCounter(container, WRITE_COUNTER_PREFIX, bucket)); - // --gcsPerformanceMetrics: wire bytes and HTTP counters - assertEquals(content.length, gcsCounter(container, "gcs_http_write_wire_bytes_sent")); - assertTrue(gcsCounter(container, "gcs_http_write_request_count") >= 1); - assertTrue(gcsCounter(container, "gcs_http_write_status_2xx") >= 1); - // A write must not produce read-side counters. - assertNull( - container.tryGetCounter( - MetricName.named(GcsUtil.METRIC_NAMESPACE, "gcs_http_read_request_count"))); - // API request metric - assertEquals(1, apiRequestCount(processWide, "GcsInsert", "ok", bucket)); - } -} 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..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 @@ -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; @@ -29,105 +24,42 @@ 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.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.Blob; 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.time.Duration; +import java.time.Instant; +import java.time.ZoneOffset; 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; -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.function.ThrowingRunnable; -import org.junit.rules.ExpectedException; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; import org.mockito.Mockito; @@ -135,7 +67,6 @@ /** Test case for {@link GcsUtil}. */ @RunWith(JUnit4.class) public class GcsUtilTest { - @Rule public ExpectedException thrown = ExpectedException.none(); MetricsContainerImpl testMetricsContainer; @Before @@ -165,30 +96,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 +104,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 +136,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. @@ -1921,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; @@ -1997,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")); @@ -2021,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"); @@ -2044,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. @@ -2085,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), @@ -2117,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); @@ -2127,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); @@ -2154,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); @@ -2200,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); @@ -2212,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); } @@ -2300,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); @@ -2346,104 +577,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; - } - - /** - * 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 - * 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 #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 #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/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..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 @@ -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,33 @@ 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.BlobInfo; 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 +68,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: * *

* - *

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 { @@ -81,6 +97,7 @@ public class GcsUtilV2Test { @After public void tearDown() { + MetricsEnvironment.setProcessWideContainer(null); MetricsEnvironment.setCurrentContainer(null); } @@ -152,7 +169,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 +290,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) { @@ -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 GcsUtilIT. + // --------------------------------------------------------------------------------------------- + + /** + * 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(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); + Storage storage = Mockito.mock(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 { + 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(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 { + Storage storage = Mockito.mock(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 { + Storage storage = Mockito.mock(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() { + 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); + + IOException thrown = + assertThrows( + IOException.class, + () -> gcsUtil.bucketAccessible(GcsPath.fromComponents("testbucket", "testobject"))); + assertSame(serverError, thrown.getCause()); + } }