/*
* 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.util;
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.assertThat;
import static org.junit.Assert.assertTrue;
import static org.mockito.Matchers.any;
import static org.mockito.Mockito.when;
import com.google.api.client.googleapis.batch.BatchRequest;
import com.google.api.client.googleapis.json.GoogleJsonError.ErrorInfo;
import com.google.api.client.googleapis.json.GoogleJsonResponseException;
import com.google.api.client.http.HttpRequest;
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.jackson2.JacksonFactory;
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.StorageObject;
import com.google.cloud.hadoop.gcsio.GoogleCloudStorageReadChannel;
import com.google.cloud.hadoop.util.ClientRequestHelper;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import java.io.ByteArrayInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.math.BigInteger;
import java.net.SocketTimeoutException;
import java.nio.channels.SeekableByteChannel;
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.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import org.apache.beam.sdk.extensions.gcp.auth.TestCredential;
import org.apache.beam.sdk.extensions.gcp.options.GcsOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.util.GcsUtil.StorageObjectOrIOException;
import org.apache.beam.sdk.util.gcsfs.GcsPath;
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;
/** Test case for {@link GcsUtil}. */
@RunWith(JUnit4.class)
public class GcsUtilTest {
@Rule public ExpectedException thrown = ExpectedException.none();
@Test
public void testGlobTranslation() {
assertEquals("foo", GcsUtil.wildcardToRegexp("foo"));
assertEquals("fo[^/]*o", GcsUtil.wildcardToRegexp("fo*o"));
assertEquals("f[^/]*o\\.[^/]", GcsUtil.wildcardToRegexp("f*o.?"));
assertEquals("foo-[0-9][^/]*", GcsUtil.wildcardToRegexp("foo-[0-9]*"));
assertEquals("foo-[0-9].*", GcsUtil.wildcardToRegexp("foo-[0-9]**"));
assertEquals(".*foo", GcsUtil.wildcardToRegexp("**/*foo"));
assertEquals(".*foo", GcsUtil.wildcardToRegexp("**foo"));
assertEquals("foo/[^/]*", GcsUtil.wildcardToRegexp("foo/*"));
assertEquals("foo[^/]*", GcsUtil.wildcardToRegexp("foo*"));
assertEquals("foo/[^/]*/[^/]*/[^/]*", GcsUtil.wildcardToRegexp("foo/*/*/*"));
assertEquals("foo/[^/]*/.*", GcsUtil.wildcardToRegexp("foo/*/**"));
assertEquals("foo.*baz", GcsUtil.wildcardToRegexp("foo**baz"));
}
private static GcsOptions gcsOptionsWithTestCredential() {
GcsOptions pipelineOptions = PipelineOptionsFactory.as(GcsOptions.class);
pipelineOptions.setGcpCredential(new TestCredential());
return pipelineOptions;
}
@Test
public void testCreationWithDefaultOptions() {
GcsOptions pipelineOptions = gcsOptionsWithTestCredential();
assertNotNull(pipelineOptions.getGcpCredential());
}
@Test
public void testUploadBufferSizeDefault() {
GcsOptions pipelineOptions = gcsOptionsWithTestCredential();
GcsUtil util = pipelineOptions.getGcsUtil();
assertNull(util.getUploadBufferSizeBytes());
}
@Test
public void testUploadBufferSizeUserSpecified() {
GcsOptions pipelineOptions = gcsOptionsWithTestCredential();
pipelineOptions.setGcsUploadBufferSizeBytes(12345);
GcsUtil util = pipelineOptions.getGcsUtil();
assertEquals((Integer) 12345, util.getUploadBufferSizeBytes());
}
@Test
public void testCreationWithExecutorServiceProvided() {
GcsOptions pipelineOptions = gcsOptionsWithTestCredential();
pipelineOptions.setExecutorService(Executors.newCachedThreadPool());
assertSame(pipelineOptions.getExecutorService(), pipelineOptions.getGcsUtil().executorService);
}
@Test
public void testCreationWithGcsUtilProvided() {
GcsOptions pipelineOptions = PipelineOptionsFactory.as(GcsOptions.class);
GcsUtil gcsUtil = Mockito.mock(GcsUtil.class);
pipelineOptions.setGcsUtil(gcsUtil);
assertSame(gcsUtil, pipelineOptions.getGcsUtil());
}
@Test
public void testMultipleThreadsCanCompleteOutOfOrderWithDefaultThreadPool() throws Exception {
GcsOptions pipelineOptions = PipelineOptionsFactory.as(GcsOptions.class);
ExecutorService executorService = pipelineOptions.getExecutorService();
int numThreads = 100;
final CountDownLatch[] countDownLatches = new CountDownLatch[numThreads];
for (int i = 0; i < numThreads; i++) {
final int currentLatch = i;
countDownLatches[i] = new CountDownLatch(1);
executorService.execute(
new Runnable() {
@Override
public void run() {
// Wait for latch N and then release latch N - 1
try {
countDownLatches[currentLatch].await();
if (currentLatch > 0) {
countDownLatches[currentLatch - 1].countDown();
}
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new RuntimeException(e);
}
}
});
}
// Release the last latch starting the chain reaction.
countDownLatches[countDownLatches.length - 1].countDown();
executorService.shutdown();
assertTrue("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.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<StorageObject> 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<GcsPath> 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<GcsPath> 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<GcsPath> 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<GcsPath> 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<GcsPath> 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.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<StorageObject> 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<GcsPath> 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.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.EMPTY_LIST, 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.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.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.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.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.getObject(
GcsPath.fromComponents("testbucket", "testobject"),
mockBackOff,
new FastNanoClockAndSleeper()).getSize().longValue());
assertEquals(BackOff.STOP, mockBackOff.nextBackOffMillis());
}
@Test
public void testGetSizeBytesWhenFileNotFoundBatch() throws Exception {
JsonFactory jsonFactory = new JacksonFactory();
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";
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.setStorageClient(new Storage(mockTransport, Transport.getJsonFactory(), null));
gcsUtil.fileSizes(ImmutableList.of(GcsPath.fromComponents("testbucket", "testobject")));
}
@Test
public void testGetSizeBytesWhenFileNotFoundBatchRetry() throws Exception {
JsonFactory jsonFactory = new JacksonFactory();
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";
thrown.expect(FileNotFoundException.class);
final LowLevelHttpResponse mockResponse = Mockito.mock(LowLevelHttpResponse.class);
when(mockResponse.getContentType()).thenReturn("multipart/mixed; boundary=" + contentBoundary);
// 429: Too many requests, then 200: OK.
when(mockResponse.getStatusCode()).thenReturn(429, 200);
when(mockResponse.getContent()).thenReturn(toStream("error"), toStream(content));
// A mock transport that lets us mock the API responses.
MockHttpTransport mockTransport =
new MockHttpTransport.Builder()
.setLowLevelHttpRequest(
new MockLowLevelHttpRequest() {
@Override
public LowLevelHttpResponse execute() throws IOException {
return mockResponse;
}
})
.build();
GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil();
gcsUtil.setStorageClient(
new Storage(mockTransport, Transport.getJsonFactory(), new RetryHttpRequestInitializer()));
gcsUtil.fileSizes(ImmutableList.of(GcsPath.fromComponents("testbucket", "testobject")));
}
@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.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.createBucket("a", new Bucket(), mockBackOff, new FastNanoClockAndSleeper());
}
@Test
public void testCreateBucketAccessErrors() throws IOException {
GcsOptions pipelineOptions = gcsOptionsWithTestCredential();
GcsUtil gcsUtil = pipelineOptions.getGcsUtil();
Storage mockStorage = Mockito.mock(Storage.class);
gcsUtil.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.createBucket("a", new Bucket(), mockBackOff, new FastNanoClockAndSleeper());
}
@Test
public void testBucketAccessible() throws IOException {
GcsOptions pipelineOptions = gcsOptionsWithTestCredential();
GcsUtil gcsUtil = pipelineOptions.getGcsUtil();
Storage mockStorage = Mockito.mock(Storage.class);
gcsUtil.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.bucketAccessible(GcsPath.fromComponents("testbucket", "testobject"),
mockBackOff, new FastNanoClockAndSleeper()));
}
@Test
public void testBucketDoesNotExistBecauseOfAccessError() throws IOException {
GcsOptions pipelineOptions = gcsOptionsWithTestCredential();
GcsUtil gcsUtil = pipelineOptions.getGcsUtil();
Storage mockStorage = Mockito.mock(Storage.class);
gcsUtil.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.bucketAccessible(GcsPath.fromComponents("testbucket", "testobject"),
mockBackOff, new FastNanoClockAndSleeper()));
}
@Test
public void testBucketDoesNotExist() throws IOException {
GcsOptions pipelineOptions = gcsOptionsWithTestCredential();
GcsUtil gcsUtil = pipelineOptions.getGcsUtil();
Storage mockStorage = Mockito.mock(Storage.class);
gcsUtil.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.bucketAccessible(GcsPath.fromComponents("testbucket", "testobject"),
mockBackOff, new FastNanoClockAndSleeper()));
}
@Test
public void testGetBucket() throws IOException {
GcsOptions pipelineOptions = gcsOptionsWithTestCredential();
GcsUtil gcsUtil = pipelineOptions.getGcsUtil();
Storage mockStorage = Mockito.mock(Storage.class);
gcsUtil.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.getBucket(GcsPath.fromComponents("testbucket", "testobject"),
mockBackOff, new FastNanoClockAndSleeper()));
}
@Test
public void testGetBucketNotExists() throws IOException {
GcsOptions pipelineOptions = gcsOptionsWithTestCredential();
GcsUtil gcsUtil = pipelineOptions.getGcsUtil();
Storage mockStorage = Mockito.mock(Storage.class);
gcsUtil.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.getBucket(GcsPath.fromComponents("testbucket", "testobject"),
mockBackOff, new FastNanoClockAndSleeper());
}
@Test
public void testGCSChannelCloseIdempotent() throws IOException {
SeekableByteChannel channel =
new GoogleCloudStorageReadChannel(null, "dummybucket", "dummyobject", null,
new ClientRequestHelper<StorageObject>());
channel.close();
channel.close();
}
/**
* 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 JacksonFactory();
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<String> makeStrings(String s, int n) {
ImmutableList.Builder<String> ret = ImmutableList.builder();
for (int i = 0; i < n; ++i) {
ret.add(String.format("gs://bucket/%s%d", s, i));
}
return ret.build();
}
private static List<GcsPath> makeGcsPaths(String s, int n) {
ImmutableList.Builder<GcsPath> 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<BatchRequest> batches) {
int ret = 0;
for (BatchRequest b : batches) {
ret += b.size();
assertThat(b.size(), greaterThan(0));
}
return ret;
}
@Test
public void testMakeCopyBatches() throws IOException {
GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil();
// Small number of files fits in 1 batch
List<BatchRequest> batches = gcsUtil.makeCopyBatches(makeStrings("s", 3), makeStrings("d", 3));
assertThat(batches.size(), equalTo(1));
assertThat(sumBatchSizes(batches), equalTo(3));
// 1 batch of files fits in 1 batch
batches = gcsUtil.makeCopyBatches(makeStrings("s", 100), makeStrings("d", 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.makeCopyBatches(makeStrings("s", 501), makeStrings("d", 501));
assertThat(batches.size(), equalTo(6));
assertThat(sumBatchSizes(batches), equalTo(501));
}
@Test
public void testInvalidCopyBatches() throws IOException {
GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil();
thrown.expect(IllegalArgumentException.class);
thrown.expectMessage("Number of source files 3");
gcsUtil.makeCopyBatches(makeStrings("s", 3), makeStrings("d", 1));
}
@Test
public void testMakeRemoveBatches() throws IOException {
GcsUtil gcsUtil = gcsOptionsWithTestCredential().getGcsUtil();
// Small number of files fits in 1 batch
List<BatchRequest> batches = gcsUtil.makeRemoveBatches(makeStrings("s", 3));
assertThat(batches.size(), equalTo(1));
assertThat(sumBatchSizes(batches), equalTo(3));
// 1 batch of files fits in 1 batch
batches = gcsUtil.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.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<StorageObjectOrIOException[]> results = Lists.newArrayList();
List<BatchRequest> batches = gcsUtil.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.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.makeGetBatches(makeGcsPaths("s", 501), results);
assertThat(batches.size(), equalTo(6));
assertThat(sumBatchSizes(batches), equalTo(501));
assertEquals(501, results.size());
}
/**
* 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));
}
}