/** * Copyright (C) 2016-2017 Expedia Inc. * * Licensed 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 com.hotels.bdp.circustrain.s3s3copier; import static org.hamcrest.CoreMatchers.is; import static org.hamcrest.CoreMatchers.startsWith; import static org.junit.Assert.assertThat; import static org.junit.Assert.fail; import static org.mockito.Matchers.any; import static org.mockito.Matchers.eq; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import java.io.File; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import org.apache.hadoop.fs.Path; import org.apache.hadoop.fs.s3a.BasicAWSCredentialsProvider; import org.junit.Before; import org.junit.Rule; import org.junit.Test; import org.junit.rules.TemporaryFolder; import org.junit.runner.RunWith; import org.mockito.Mock; import org.mockito.Mockito; import org.mockito.runners.MockitoJUnitRunner; import com.amazonaws.AmazonClientException; import com.amazonaws.AmazonServiceException; import com.amazonaws.client.builder.AwsClientBuilder.EndpointConfiguration; import com.amazonaws.regions.Regions; import com.amazonaws.services.s3.AmazonS3; import com.amazonaws.services.s3.AmazonS3ClientBuilder; import com.amazonaws.services.s3.AmazonS3URI; import com.amazonaws.services.s3.model.CopyObjectRequest; import com.amazonaws.services.s3.model.ListObjectsRequest; import com.amazonaws.services.s3.model.S3Object; import com.amazonaws.services.s3.transfer.Copy; import com.amazonaws.services.s3.transfer.TransferManager; import com.amazonaws.services.s3.transfer.TransferProgress; import com.amazonaws.services.s3.transfer.internal.TransferStateChangeListener; import com.amazonaws.util.IOUtils; import com.codahale.metrics.MetricRegistry; import com.google.common.base.Charsets; import com.google.common.collect.Lists; import com.google.common.io.Files; import com.hotels.bdp.circustrain.api.CircusTrainException; import com.hotels.bdp.circustrain.api.metrics.Metrics; import com.hotels.bdp.circustrain.common.test.junit.rules.S3ProxyRule; import com.hotels.bdp.circustrain.s3s3copier.aws.AmazonS3ClientFactory; import com.hotels.bdp.circustrain.s3s3copier.aws.ListObjectsRequestFactory; import com.hotels.bdp.circustrain.s3s3copier.aws.TransferManagerFactory; @RunWith(MockitoJUnitRunner.class) public class S3S3CopierTest { private static final String AWS_ACCESS_KEY = "access"; private static final String AWS_SECRET_KEY = "secret"; public @Rule TemporaryFolder temp = new TemporaryFolder(); public @Rule S3ProxyRule s3Proxy = S3ProxyRule.builder().withCredentials(AWS_ACCESS_KEY, AWS_SECRET_KEY).build(); private final ListObjectsRequestFactory listObjectsRequestFactory = new ListObjectsRequestFactory(); private final TransferManagerFactory transferManagerFactory = new TransferManagerFactory(); private final MetricRegistry registry = new MetricRegistry(); private final S3S3CopierOptions s3S3CopierOptions = new S3S3CopierOptions(new HashMap<String, Object>()); private @Mock AmazonS3ClientFactory s3ClientFactory; private File inputData; private AmazonS3 client; @Before public void setUp() throws Exception { inputData = temp.newFile("data"); Files.write("bar foo", inputData, Charsets.UTF_8); client = newClient(); client.createBucket("source"); client.createBucket("target"); when(s3ClientFactory.newInstance(any(AmazonS3URI.class), any(S3S3CopierOptions.class))).thenReturn(newClient()); } private AmazonS3 newClient() { EndpointConfiguration endpointConfiguration = new EndpointConfiguration(s3Proxy.getProxyUrl(), Regions.DEFAULT_REGION.getName()); AmazonS3 newClient = AmazonS3ClientBuilder .standard() .withCredentials(new BasicAWSCredentialsProvider(AWS_ACCESS_KEY, AWS_SECRET_KEY)) .withEndpointConfiguration(endpointConfiguration) .build(); return newClient; } @Test public void copyOneObject() throws Exception { client.putObject("source", "data", inputData); Path sourceBaseLocation = new Path("s3://source/"); Path replicaLocation = new Path("s3://target/"); List<Path> sourceSubLocations = new ArrayList<>(); S3S3Copier s3s3Copier = newS3S3Copier(sourceBaseLocation, sourceSubLocations, replicaLocation); Metrics metrics = s3s3Copier.copy(); assertThat(metrics.getBytesReplicated(), is(7L)); assertThat(metrics.getMetrics().get(S3S3CopierMetrics.Metrics.TOTAL_BYTES_TO_REPLICATE.name()), is(7L)); S3Object object = client.getObject("target", "data"); String data = IOUtils.toString(object.getObjectContent()); assertThat(data, is("bar foo")); assertThat(registry.getGauges().containsKey(RunningMetrics.S3S3_CP_BYTES_REPLICATED.name()), is(true)); } private S3S3Copier newS3S3Copier(Path sourceBaseLocation, List<Path> sourceSubLocations, Path replicaLocation) { return new S3S3Copier(sourceBaseLocation, sourceSubLocations, replicaLocation, s3ClientFactory, transferManagerFactory, listObjectsRequestFactory, registry, s3S3CopierOptions); } @Test public void copyOneObjectUsingKeys() throws Exception { client.putObject("source", "bar/data", inputData); Path sourceBaseLocation = new Path("s3://source/bar/"); Path replicaLocation = new Path("s3://target/foo/"); List<Path> sourceSubLocations = new ArrayList<>(); S3S3Copier s3s3Copier = newS3S3Copier(sourceBaseLocation, sourceSubLocations, replicaLocation); s3s3Copier.copy(); S3Object object = client.getObject("target", "foo/data"); String data = IOUtils.toString(object.getObjectContent()); assertThat(data, is("bar foo")); } @Test public void copyOneObjectPartitioned() throws Exception { client.putObject("source", "year=2016/data", inputData); Path sourceBaseLocation = new Path("s3://source/"); Path replicaLocation = new Path("s3://target/foo/"); List<Path> sourceSubLocations = Lists.newArrayList(new Path(sourceBaseLocation, "year=2016")); S3S3Copier s3s3Copier = newS3S3Copier(sourceBaseLocation, sourceSubLocations, replicaLocation); Metrics metrics = s3s3Copier.copy(); assertThat(metrics.getBytesReplicated(), is(7L)); S3Object object = client.getObject("target", "foo/year=2016/data"); String data = IOUtils.toString(object.getObjectContent()); assertThat(data, is("bar foo")); } @Test public void copyOneObjectPartitionedSourceBaseNested() throws Exception { client.putObject("source", "nested/year=2016/data", inputData); Path sourceBaseLocation = new Path("s3://source/nested");// no slash at the end Path replicaLocation = new Path("s3://target/foo/"); List<Path> sourceSubLocations = Lists.newArrayList(new Path(sourceBaseLocation, "year=2016")); S3S3Copier s3s3Copier = newS3S3Copier(sourceBaseLocation, sourceSubLocations, replicaLocation); s3s3Copier.copy(); S3Object object = client.getObject("target", "foo/year=2016/data"); String data = IOUtils.toString(object.getObjectContent()); assertThat(data, is("bar foo")); } @Test public void copyOneObjectPartitionedHandlingS3ASchemes() throws Exception { client.putObject("source", "year=2016/data", inputData); Path sourceBaseLocation = new Path("s3a://source/"); Path replicaLocation = new Path("s3a://target/foo/"); List<Path> sourceSubLocations = Lists.newArrayList(new Path(sourceBaseLocation, "year=2016")); S3S3Copier s3s3Copier = newS3S3Copier(sourceBaseLocation, sourceSubLocations, replicaLocation); s3s3Copier.copy(); S3Object object = client.getObject("target", "foo/year=2016/data"); String data = IOUtils.toString(object.getObjectContent()); assertThat(data, is("bar foo")); } @Test public void copyMultipleObjects() throws Exception { // Making sure we only request 1 file at the time so we need to loop ListObjectsRequestFactory mockListObjectRequestFactory = Mockito.mock(ListObjectsRequestFactory.class); when(mockListObjectRequestFactory.newInstance()).thenReturn(new ListObjectsRequest().withMaxKeys(1)); client.putObject("source", "bar/data1", inputData); client.putObject("source", "bar/data2", inputData); Path sourceBaseLocation = new Path("s3://source/bar/"); Path replicaLocation = new Path("s3://target/foo/"); List<Path> sourceSubLocations = new ArrayList<>(); S3S3Copier s3s3Copier = new S3S3Copier(sourceBaseLocation, sourceSubLocations, replicaLocation, s3ClientFactory, transferManagerFactory, mockListObjectRequestFactory, registry, s3S3CopierOptions); Metrics metrics = s3s3Copier.copy(); assertThat(metrics.getBytesReplicated(), is(14L)); S3Object object1 = client.getObject("target", "foo/data1"); String data1 = IOUtils.toString(object1.getObjectContent()); assertThat(data1, is("bar foo")); S3Object object2 = client.getObject("target", "foo/data2"); String data2 = IOUtils.toString(object2.getObjectContent()); assertThat(data2, is("bar foo")); } @Test public void copyCheckTransferManagerIsShutdown() throws Exception { client.putObject("source", "data", inputData); Path sourceBaseLocation = new Path("s3://source/"); Path replicaLocation = new Path("s3://target/"); List<Path> sourceSubLocations = new ArrayList<>(); TransferManagerFactory mockedTransferManagerFactory = Mockito.mock(TransferManagerFactory.class); TransferManager mockedTransferManager = Mockito.mock(TransferManager.class); when(mockedTransferManagerFactory.newInstance(any(AmazonS3.class), eq(s3S3CopierOptions))) .thenReturn(mockedTransferManager); Copy copy = Mockito.mock(Copy.class); when(mockedTransferManager.copy(any(CopyObjectRequest.class), any(AmazonS3.class), any(TransferStateChangeListener.class))).thenReturn(copy); TransferProgress transferProgress = new TransferProgress(); when(copy.getProgress()).thenReturn(transferProgress); S3S3Copier s3s3Copier = new S3S3Copier(sourceBaseLocation, sourceSubLocations, replicaLocation, s3ClientFactory, mockedTransferManagerFactory, listObjectsRequestFactory, registry, s3S3CopierOptions); s3s3Copier.copy(); verify(mockedTransferManager).shutdownNow(); } @Test public void copySafelyShutDownTransferManagerWhenNotInitialised() throws Exception { Path sourceBaseLocation = new Path("s3://source/"); Path replicaLocation = new Path("s3://target/"); List<Path> sourceSubLocations = new ArrayList<>(); TransferManagerFactory mockedTransferManagerFactory = Mockito.mock(TransferManagerFactory.class); when(mockedTransferManagerFactory.newInstance(any(AmazonS3.class), eq(s3S3CopierOptions))) .thenThrow(new RuntimeException("error in instance")); S3S3Copier s3s3Copier = new S3S3Copier(sourceBaseLocation, sourceSubLocations, replicaLocation, s3ClientFactory, mockedTransferManagerFactory, listObjectsRequestFactory, registry, s3S3CopierOptions); try { s3s3Copier.copy(); } catch (RuntimeException e) { assertThat(e.getMessage(), is("error in instance")); } } @Test public void copyCheckTransferManagerIsShutdownWhenSubmittingJobExceptionsAreThrown() throws Exception { client.putObject("source", "data", inputData); Path sourceBaseLocation = new Path("s3://source/"); Path replicaLocation = new Path("s3://target/"); List<Path> sourceSubLocations = new ArrayList<>(); TransferManagerFactory mockedTransferManagerFactory = Mockito.mock(TransferManagerFactory.class); TransferManager mockedTransferManager = Mockito.mock(TransferManager.class); when(mockedTransferManagerFactory.newInstance(any(AmazonS3.class), eq(s3S3CopierOptions))) .thenReturn(mockedTransferManager); when(mockedTransferManager.copy(any(CopyObjectRequest.class), any(AmazonS3.class), any(TransferStateChangeListener.class))).thenThrow(new AmazonServiceException("MyCause")); S3S3Copier s3s3Copier = new S3S3Copier(sourceBaseLocation, sourceSubLocations, replicaLocation, s3ClientFactory, mockedTransferManagerFactory, listObjectsRequestFactory, registry, s3S3CopierOptions); try { s3s3Copier.copy(); fail("exception should have been thrown"); } catch (CircusTrainException e) { verify(mockedTransferManager).shutdownNow(); assertThat(e.getCause().getMessage(), startsWith("MyCause")); } } @Test public void copyCheckTransferManagerIsShutdownWhenCopyExceptionsAreThrown() throws Exception { client.putObject("source", "data", inputData); Path sourceBaseLocation = new Path("s3://source/"); Path replicaLocation = new Path("s3://target/"); List<Path> sourceSubLocations = new ArrayList<>(); TransferManagerFactory mockedTransferManagerFactory = Mockito.mock(TransferManagerFactory.class); TransferManager mockedTransferManager = Mockito.mock(TransferManager.class); when(mockedTransferManagerFactory.newInstance(any(AmazonS3.class), eq(s3S3CopierOptions))) .thenReturn(mockedTransferManager); Copy copy = Mockito.mock(Copy.class); when(copy.getProgress()).thenReturn(new TransferProgress()); when(mockedTransferManager.copy(any(CopyObjectRequest.class), any(AmazonS3.class), any(TransferStateChangeListener.class))).thenReturn(copy); doThrow(new AmazonClientException("cause")).when(copy).waitForCompletion(); S3S3Copier s3s3Copier = new S3S3Copier(sourceBaseLocation, sourceSubLocations, replicaLocation, s3ClientFactory, mockedTransferManagerFactory, listObjectsRequestFactory, registry, s3S3CopierOptions); try { s3s3Copier.copy(); fail("exception should have been thrown"); } catch (CircusTrainException e) { verify(mockedTransferManager).shutdownNow(); assertThat(e.getCause().getMessage(), is("cause")); } } }