Search in sources :

Example 11 with BackOff

use of com.google.api.client.util.BackOff in project beam by apache.

the class GcsUtilTest method testBucketAccessible.

@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()));
}
Also used : Storage(com.google.api.services.storage.Storage) SocketTimeoutException(java.net.SocketTimeoutException) Bucket(com.google.api.services.storage.model.Bucket) GcsOptions(org.apache.beam.sdk.extensions.gcp.options.GcsOptions) BackOff(com.google.api.client.util.BackOff) Test(org.junit.Test)

Example 12 with BackOff

use of com.google.api.client.util.BackOff in project beam by apache.

the class GcsUtilTest method testRetryFileSizeNonBatch.

@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());
}
Also used : Storage(com.google.api.services.storage.Storage) SocketTimeoutException(java.net.SocketTimeoutException) StorageObject(com.google.api.services.storage.model.StorageObject) Objects(com.google.api.services.storage.model.Objects) GcsOptions(org.apache.beam.sdk.extensions.gcp.options.GcsOptions) BackOff(com.google.api.client.util.BackOff) Test(org.junit.Test)

Aggregations

BackOff (com.google.api.client.util.BackOff)12 Storage (com.google.api.services.storage.Storage)8 GcsOptions (org.apache.beam.sdk.extensions.gcp.options.GcsOptions)8 Test (org.junit.Test)8 SocketTimeoutException (java.net.SocketTimeoutException)5 Bucket (com.google.api.services.storage.model.Bucket)4 GoogleJsonResponseException (com.google.api.client.googleapis.json.GoogleJsonResponseException)3 IOException (java.io.IOException)3 Sleeper (com.google.api.client.util.Sleeper)2 Job (com.google.api.services.bigquery.model.Job)1 JobStatus (com.google.api.services.bigquery.model.JobStatus)1 JobMessage (com.google.api.services.dataflow.model.JobMessage)1 Objects (com.google.api.services.storage.model.Objects)1 StorageObject (com.google.api.services.storage.model.StorageObject)1 VisibleForTesting (com.google.common.annotations.VisibleForTesting)1 FileNotFoundException (java.io.FileNotFoundException)1 WritableByteChannel (java.nio.channels.WritableByteChannel)1 ExecutionException (java.util.concurrent.ExecutionException)1 Nullable (javax.annotation.Nullable)1 Duration (org.joda.time.Duration)1