apoc.util.s3.S3TestUtil Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of apoc-test-utils Show documentation
Test utils for Neo4j Procedures
There is a newer version: 5.24.0
/*
 * Copyright (c) "Neo4j"
 * Neo4j Sweden AB [http://neo4j.com]
 *
 * This file is part of Neo4j.
 *
 * 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 apoc.util.s3;

import static org.junit.Assert.assertEquals;

import com.amazonaws.AmazonClientException;
import com.amazonaws.services.s3.AmazonS3;
import com.amazonaws.services.s3.model.S3Object;
import com.amazonaws.services.s3.model.S3ObjectInputStream;
import com.amazonaws.util.IOUtils;
import java.io.IOException;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.concurrent.TimeUnit;
import org.neo4j.test.assertion.Assert;

/**
 * Utility class for testing Amazon S3 related functionality.
 */
public class S3TestUtil {

    /**
     * Read file object as a string from S3 bucket. This code expects valid AWS credentials are set up.
     * @param s3Url String containing url to S3 bucket.
     * @return the s3 string object
     */
    public static String readS3FileToString(String s3Url) throws AmazonClientException {
        try {
            S3Object s3object = getS3Object(s3Url);
            S3ObjectInputStream inputStream = s3object.getObjectContent();
            return IOUtils.toString(inputStream);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    public static S3Object getS3Object(String s3Url) throws MalformedURLException, AmazonClientException {
        S3Params s3Params = S3ParamsExtractor.extract(new URL(s3Url));
        S3Aws s3Aws = new S3Aws(s3Params, s3Params.getRegion());
        AmazonS3 s3Client = s3Aws.getClient();

        return s3Client.getObject(s3Params.getBucket(), s3Params.getKey());
    }

    public static void assertStringFileEquals(String expected, String s3Url) {
        assertS3KeyEventually(() -> {
            final String actual = readS3FileToString(s3Url);
            assertEquals(expected, actual);
        });
    }

    public static void assertS3KeyEventually(Runnable runnable) {
        Assert.assertEventually(
                () -> {
                    try {
                        runnable.run();
                        return true;
                    } catch (AmazonClientException e) {
                        if (e.getMessage().contains("The specified key does not exist")) {
                            return false;
                        }
                        throw e;
                    }
                },
                v -> v,
                30L,
                TimeUnit.SECONDS);
    }
}