diff --git a/docs/dev/IntegrationTests.md b/docs/dev/IntegrationTests.md index e62381da..3383c58e 100644 --- a/docs/dev/IntegrationTests.md +++ b/docs/dev/IntegrationTests.md @@ -2,16 +2,20 @@ ## Overview -The integration tests for Neptune Export run a set of full export jobs against a Neptune instance and validate that the output is correct. A user must provide their own Neptune instance with the correct sample data pre-loaded in order to run these tests. +The integration tests for Neptune Export run a set of full export jobs against a Neptune instance and validate that the output is correct. A user must provide their own Neptune instance with the correct sample data pre-loaded in order to run these tests. The tests +assume that IAM authentication is enabled on the provided Neptune Cluster. ## Setup - Setup a clean Neptune instance to run integration tests against. - Using the Neptune bulk loader, fill the database with Kelvin Lawrence's air-routes-small graph data [found here](https://github.com/krlawrence/graph/blob/master/sample-data/air-routes-small.graphml). GraphML files can be run through [graphml2csv](https://github.com/awslabs/amazon-neptune-tools/tree/master/graphml2csv) for property graph loading and csv files can go through [csv to rdf](https://github.com/aws/amazon-neptune-csv-to-rdf-converter) for RDF loading. -- Ensure your Neptune instance is accessible to your testing environment (either run tests from within the -same VPC as your instance or connect through a bastion host) +- Ensure your Neptune instance is accessible to your testing environment (use a public endpoint,run tests from within the +same VPC as your instance, or connect through a bastion host) +- Create an S3 bucket which can be read and written to by `ExportServiceIntegrationTest`. ## Running Tests - Set a `NEPTUNE_ENDPOINT` environment variable to your Neptune cluster endpoint +- Set a `SERVICE_REGION` environment variable for IAM auth +- Set a `S3_PATH` environment variable for ExportServiceIntegrationTest - Run Junit integration tests through your IDE \ No newline at end of file diff --git a/src/main/java/com/amazonaws/services/neptune/export/NeptuneExportLambda.java b/src/main/java/com/amazonaws/services/neptune/export/NeptuneExportLambda.java index 8596ac2e..4a0fabd3 100644 --- a/src/main/java/com/amazonaws/services/neptune/export/NeptuneExportLambda.java +++ b/src/main/java/com/amazonaws/services/neptune/export/NeptuneExportLambda.java @@ -181,9 +181,7 @@ public void handleRequest(InputStream inputStream, OutputStream outputStream, Co } if (outputS3ObjectInfo != null) { - try (Writer writer = new BufferedWriter(new OutputStreamWriter(outputStream, UTF_8))) { - writer.write(outputS3ObjectInfo.toString()); - } + outputStream.write(outputS3ObjectInfo.toString().getBytes(UTF_8)); } else { System.exit(-1); } diff --git a/src/test/java/com/amazonaws/services/neptune/CreatePgConfigIntegrationTest.java b/src/test/java/com/amazonaws/services/neptune/CreatePgConfigIntegrationTest.java index a05b9fa4..addf8bbc 100644 --- a/src/test/java/com/amazonaws/services/neptune/CreatePgConfigIntegrationTest.java +++ b/src/test/java/com/amazonaws/services/neptune/CreatePgConfigIntegrationTest.java @@ -21,7 +21,7 @@ public class CreatePgConfigIntegrationTest extends AbstractExportIntegrationTest @Test public void testCreatePgConfig() { - final String[] command = {"create-pg-config", "-e", neptuneEndpoint, "-d", outputDir.getPath()}; + final String[] command = {"create-pg-config", "-e", neptuneEndpoint, "-d", outputDir.getPath(), "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -33,7 +33,7 @@ public void testCreatePgConfig() { @Test public void testCreatePgConfigWithGremlinFilter() { final String[] command = {"create-pg-config", "-e", neptuneEndpoint, "-d", outputDir.getPath(), - "--gremlin-filter", "has(\"runways\", 2)"}; + "--gremlin-filter", "has(\"runways\", 2)", "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -45,7 +45,7 @@ public void testCreatePgConfigWithGremlinFilter() { @Test public void testCreatePgConfigWithEdgeGremlinFilter() { final String[] command = {"create-pg-config", "-e", neptuneEndpoint, "-d", outputDir.getPath(), - "--gremlin-filter", "hasLabel(\"route\")"}; + "--gremlin-filter", "hasLabel(\"route\")", "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -57,7 +57,7 @@ public void testCreatePgConfigWithEdgeGremlinFilter() { @Test public void testCreatePgConfigWithEdgeGremlinFilterAndEarlyGremlinFilter() { final String[] command = {"create-pg-config", "-e", neptuneEndpoint, "-d", outputDir.getPath(), - "--gremlin-filter", "hasLabel(\"route\")", "--filter-edges-early"}; + "--gremlin-filter", "hasLabel(\"route\")", "--filter-edges-early", "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); diff --git a/src/test/java/com/amazonaws/services/neptune/ExportPgFromConfigIntegrationTest.java b/src/test/java/com/amazonaws/services/neptune/ExportPgFromConfigIntegrationTest.java index ca07c21f..425f9f6c 100644 --- a/src/test/java/com/amazonaws/services/neptune/ExportPgFromConfigIntegrationTest.java +++ b/src/test/java/com/amazonaws/services/neptune/ExportPgFromConfigIntegrationTest.java @@ -23,7 +23,7 @@ public class ExportPgFromConfigIntegrationTest extends AbstractExportIntegration public void testExportPgFromConfig() { final String[] command = {"export-pg-from-config", "-e", neptuneEndpoint, "-c", "src/test/resources/IntegrationTest/ExportPgFromConfigIntegrationTest/input/config.json", - "-d", outputDir.getPath()}; + "-d", outputDir.getPath(), "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -36,7 +36,7 @@ public void testExportPgFromConfig() { public void testExportPgFromConfigWithGremlinFilter() { final String[] command = {"export-pg-from-config", "-e", neptuneEndpoint, "-d", outputDir.getPath(), "-c", "src/test/resources/IntegrationTest/ExportPgFromConfigIntegrationTest/input/config.json", - "--gremlin-filter", "has(\"runways\", 2)"}; + "--gremlin-filter", "has(\"runways\", 2)", "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -49,7 +49,7 @@ public void testExportPgFromConfigWithGremlinFilter() { public void testExportEdgesFromConfigWithGremlinFilter() { final String[] command = {"export-pg-from-config", "-e", neptuneEndpoint, "-d", outputDir.getPath(), "-c", "src/test/resources/IntegrationTest/ExportPgFromConfigIntegrationTest/input/config.json", - "--gremlin-filter", "hasLabel(\"route\")"}; + "--gremlin-filter", "hasLabel(\"route\")", "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -62,7 +62,7 @@ public void testExportEdgesFromConfigWithGremlinFilter() { public void testExportEdgesFromConfigWithGremlinFilterWithEarlyGremlinFilter() { final String[] command = {"export-pg-from-config", "-e", neptuneEndpoint, "-d", outputDir.getPath(), "-c", "src/test/resources/IntegrationTest/ExportPgFromConfigIntegrationTest/input/config.json", - "--gremlin-filter", "hasLabel(\"route\")", "--filter-edges-early"}; + "--gremlin-filter", "hasLabel(\"route\")", "--filter-edges-early", "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); diff --git a/src/test/java/com/amazonaws/services/neptune/ExportPgFromQueriesIntegrationTest.java b/src/test/java/com/amazonaws/services/neptune/ExportPgFromQueriesIntegrationTest.java index e8ba9bf4..62e85978 100644 --- a/src/test/java/com/amazonaws/services/neptune/ExportPgFromQueriesIntegrationTest.java +++ b/src/test/java/com/amazonaws/services/neptune/ExportPgFromQueriesIntegrationTest.java @@ -25,7 +25,8 @@ public class ExportPgFromQueriesIntegrationTest extends AbstractExportIntegratio public void testExportPgFromQueries() { final String[] command = {"export-pg-from-queries", "-e", neptuneEndpoint, "-d", outputDir.getPath(), - "-q", "airport=g.V().hasLabel('airport').has('runways', gt(2)).project('code', 'runways', 'city', 'country').by('code').by('runways').by('city').by('country')" + "-q", "airport=g.V().hasLabel('airport').has('runways', gt(2)).project('code', 'runways', 'city', 'country').by('code').by('runways').by('city').by('country')", + "--use-iam-auth" }; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -39,7 +40,8 @@ public void testExportPgFromQueries() { public void testExportPgFromQueriesNoHeaders() { final String[] command = {"export-pg-from-queries", "-e", neptuneEndpoint, "-d", outputDir.getPath(), "--format", "csvNoHeaders", - "-q", "airport=g.V().hasLabel('airport').has('runways', gt(2)).project('code', 'runways', 'city', 'country').by('code').by('runways').by('city').by('country')" + "-q", "airport=g.V().hasLabel('airport').has('runways', gt(2)).project('code', 'runways', 'city', 'country').by('code').by('runways').by('city').by('country')", + "--use-iam-auth" }; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -53,7 +55,8 @@ public void testExportPgFromQueriesNoHeaders() { public void testExportPgFromQueriesWithStaggeredResults() { final String[] command = {"export-pg-from-queries", "-e", neptuneEndpoint, "-d", outputDir.getPath(), - "-q", "airport=g.inject(['code': 'YYC'], ['city': 'Vancouver', 'code': 'YVR'], ['code':'SEA', 'city':'Seattle', 'runways': 3])" + "-q", "airport=g.inject(['code': 'YYC'], ['city': 'Vancouver', 'code': 'YVR'], ['code':'SEA', 'city':'Seattle', 'runways': 3])", + "--use-iam-auth" }; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -67,7 +70,8 @@ public void testExportPgFromQueriesWithStaggeredResults() { public void testExportPgFromQueriesWithStaggeredResultsNoHeaders() { final String[] command = {"export-pg-from-queries", "-e", neptuneEndpoint, "-d", outputDir.getPath(), "--format", "csvNoHeaders", - "-q", "airport=g.inject(['code':'SEA', 'city':'Seattle', 'runways': 3], ['city': 'Vancouver', 'code': 'YVR'], ['code': 'YYC'])" + "-q", "airport=g.inject(['code':'SEA', 'city':'Seattle', 'runways': 3], ['city': 'Vancouver', 'code': 'YVR'], ['code': 'YYC'])", + "--use-iam-auth" }; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -82,7 +86,7 @@ public void testExportPgFromQueriesSplitQueries() { final String[] command = {"export-pg-from-queries", "-e", neptuneEndpoint, "-d", outputDir.getPath(), "-q", "airport=g.V().hasLabel('airport').has('runways', gt(2)).project('code', 'runways', 'city', 'country').by('code').by('runways').by('city').by('country')", - "--split-queries" + "--split-queries", "--use-iam-auth" }; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -97,7 +101,7 @@ public void testExportPgFromQueriesSplitQueriesAndRange() { final String[] command = {"export-pg-from-queries", "-e", neptuneEndpoint, "-d", outputDir.getPath(), "-q", "airport=g.V().hasLabel('airport').has('runways', gt(2)).project('code', 'runways', 'city', 'country').by('code').by('runways').by('city').by('country')", - "--split-queries", "--range", "25" + "--split-queries", "--range", "25", "--use-iam-auth" }; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -113,7 +117,7 @@ public void testExportPgFromQueriesWithStructuredOutput() { "-d", outputDir.getPath(), "-q", "airport=g.V().union(hasLabel('airport'), outE()).elementMap()", "--include-type-definitions", - "--structured-output" + "--structured-output", "--use-iam-auth" }; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -130,7 +134,7 @@ public void testExportPgFromQueriesWithStructuredOutputSplitQueries() { "-q", "airport=g.V().union(hasLabel('airport'), outE()).elementMap()", "--include-type-definitions", "--split-queries", "--range", "25", - "--structured-output" + "--structured-output", "--use-iam-auth" }; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -145,7 +149,7 @@ public void testExportPgFromQueriesWithStructuredOutputWithEdgeAndVertexLabels() final String[] command = {"export-pg-from-queries", "-e", neptuneEndpoint, "-d", outputDir.getPath(), "-q", "all=g.V().union(elementMap(), outE().elementMap())", - "--edge-label-strategy", "edgeAndVertexLabels", "--structured-output" + "--edge-label-strategy", "edgeAndVertexLabels", "--structured-output", "--use-iam-auth" }; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -161,7 +165,7 @@ public void testExportPgFromQueriesWithStructuredOutputWithEdgeAndVertexLabelsIn "-d", outputDir.getPath(), "-q", "airport=g.V().union(hasLabel('airport'), outE()).elementMap()", "--include-type-definitions", "--edge-label-strategy", "edgeAndVertexLabels", - "--structured-output" + "--structured-output", "--use-iam-auth" }; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); diff --git a/src/test/java/com/amazonaws/services/neptune/ExportPgIntegrationTest.java b/src/test/java/com/amazonaws/services/neptune/ExportPgIntegrationTest.java index 6e0714dc..76ae9a54 100644 --- a/src/test/java/com/amazonaws/services/neptune/ExportPgIntegrationTest.java +++ b/src/test/java/com/amazonaws/services/neptune/ExportPgIntegrationTest.java @@ -22,7 +22,7 @@ public class ExportPgIntegrationTest extends AbstractExportIntegrationTest{ @Test public void testExportPgToCsv() { - final String[] command = {"export-pg", "-e", neptuneEndpoint, "-d", outputDir.getPath()}; + final String[] command = {"export-pg", "-e", neptuneEndpoint, "-d", outputDir.getPath(), "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -34,7 +34,7 @@ public void testExportPgToCsv() { @Test public void testExportPgWithEdgeAndVertexLabels() { final String[] command = {"export-pg", "-e", neptuneEndpoint, "-d", outputDir.getPath(), - "--edge-label-strategy", "edgeAndVertexLabels"}; + "--edge-label-strategy", "edgeAndVertexLabels", "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -45,7 +45,7 @@ public void testExportPgWithEdgeAndVertexLabels() { @Test public void testExportPgToCsvWithJanus() { - final String[] command = {"export-pg", "-e", neptuneEndpoint, "-d", outputDir.getPath(), "--janus"}; + final String[] command = {"export-pg", "-e", neptuneEndpoint, "-d", outputDir.getPath(), "--janus", "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -57,7 +57,7 @@ public void testExportPgToCsvWithJanus() { @Test public void testExportPgToCsvWithGremlinFilter() { final String[] command = {"export-pg", "-e", neptuneEndpoint, "-d", outputDir.getPath(), - "--gremlin-filter", "has(\"runways\", 2)"}; + "--gremlin-filter", "has(\"runways\", 2)", "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -69,7 +69,7 @@ public void testExportPgToCsvWithGremlinFilter() { @Test public void testExportEdgesToCsvWithGremlinFilter() { final String[] command = {"export-pg", "-e", neptuneEndpoint, "-d", outputDir.getPath(), - "--gremlin-filter", "hasLabel(\"route\")"}; + "--gremlin-filter", "hasLabel(\"route\")", "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -81,7 +81,7 @@ public void testExportEdgesToCsvWithGremlinFilter() { @Test public void testExportEdgesToCsvWithGremlinFilterWithEarlyGremlinFilter() { final String[] command = {"export-pg", "-e", neptuneEndpoint, "-d", outputDir.getPath(), - "--gremlin-filter", "hasLabel(\"route\")", "--filter-edges-early"}; + "--gremlin-filter", "hasLabel(\"route\")", "--filter-edges-early", "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); diff --git a/src/test/java/com/amazonaws/services/neptune/ExportRdfIntegrationTest.java b/src/test/java/com/amazonaws/services/neptune/ExportRdfIntegrationTest.java index be218352..b0045c33 100644 --- a/src/test/java/com/amazonaws/services/neptune/ExportRdfIntegrationTest.java +++ b/src/test/java/com/amazonaws/services/neptune/ExportRdfIntegrationTest.java @@ -29,7 +29,7 @@ public class ExportRdfIntegrationTest extends AbstractExportIntegrationTest{ @Test public void testExportRdf() { - final String[] command = {"export-rdf", "-e", neptuneEndpoint, "-d", outputDir.getPath()}; + final String[] command = {"export-rdf", "-e", neptuneEndpoint, "-d", outputDir.getPath(), "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -41,7 +41,7 @@ public void testExportRdf() { @Test public void testExportRdfSingleNamedGraphVersion() { final String[] command = {"export-rdf", "-e", neptuneEndpoint, "-d", outputDir.getPath(), - "--named-graph", "http://aws.amazon.com/neptune/csv2rdf/graph/version"}; + "--named-graph", "http://aws.amazon.com/neptune/csv2rdf/graph/version", "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -53,7 +53,7 @@ public void testExportRdfSingleNamedGraphVersion() { @Test public void testExportRdfSingleNamedGraphDefault() { final String[] command = {"export-rdf", "-e", neptuneEndpoint, "-d", outputDir.getPath(), - "--named-graph", "http://aws.amazon.com/neptune/vocab/v01/DefaultNamedGraph"}; + "--named-graph", "http://aws.amazon.com/neptune/vocab/v01/DefaultNamedGraph", "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -65,7 +65,7 @@ public void testExportRdfSingleNamedGraphDefault() { @Test public void testExportRdfNoGSP() { final String[] command = {"export-rdf", "-e", neptuneEndpoint, "-d", outputDir.getPath(), - "--feature-toggle", "No_GSP"}; + "--feature-toggle", "No_GSP", "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -78,7 +78,7 @@ public void testExportRdfNoGSP() { public void testExportRdfSingleNamedGraphVersionNoGSP() { final String[] command = {"export-rdf", "-e", neptuneEndpoint, "-d", outputDir.getPath(), "--named-graph", "http://aws.amazon.com/neptune/csv2rdf/graph/version", - "--feature-toggle", "No_GSP"}; + "--feature-toggle", "No_GSP", "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); @@ -91,7 +91,7 @@ public void testExportRdfSingleNamedGraphVersionNoGSP() { public void testExportRdfSingleNamedGraphDefaultNoGSP() { final String[] command = {"export-rdf", "-e", neptuneEndpoint, "-d", outputDir.getPath(), "--named-graph", "http://aws.amazon.com/neptune/vocab/v01/DefaultNamedGraph", - "--feature-toggle", "No_GSP"}; + "--feature-toggle", "No_GSP", "--use-iam-auth"}; final NeptuneExportRunner runner = new NeptuneExportRunner(command); runner.run(); diff --git a/src/test/java/com/amazonaws/services/neptune/ExportServiceIntegrationTest.java b/src/test/java/com/amazonaws/services/neptune/ExportServiceIntegrationTest.java index de776b56..43573a4f 100644 --- a/src/test/java/com/amazonaws/services/neptune/ExportServiceIntegrationTest.java +++ b/src/test/java/com/amazonaws/services/neptune/ExportServiceIntegrationTest.java @@ -13,22 +13,47 @@ package com.amazonaws.services.neptune; import com.amazonaws.services.neptune.export.NeptuneExportRunner; +import com.amazonaws.services.neptune.util.S3ObjectInfo; +import com.amazonaws.services.neptune.util.TransferManagerWrapper; +import org.junit.BeforeClass; import org.junit.Rule; import org.junit.Test; import org.junit.contrib.java.lang.system.Assertion; import org.junit.contrib.java.lang.system.ExpectedSystemExit; +import software.amazon.awssdk.services.s3.S3Client; +import software.amazon.awssdk.services.s3.model.*; +import software.amazon.awssdk.transfer.s3.model.DownloadFileRequest; +import software.amazon.awssdk.transfer.s3.model.FileDownload; import java.io.File; +import java.util.List; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CompletionException; +import java.util.stream.Collectors; + +import static org.junit.Assert.assertNotNull; public class ExportServiceIntegrationTest extends AbstractExportIntegrationTest{ + private static String s3Path; + @Rule public final ExpectedSystemExit exit = ExpectedSystemExit.none(); + @BeforeClass + public static void setupS3() { + s3Path = System.getenv("S3_PATH"); + assertNotNull("S3 path must be provided through \"S3_PATH\" environment variable", s3Path); + if (s3Path.endsWith("/")) { + s3Path = s3Path.substring(0, s3Path.length() - 1); + } + } + @Test public void testExportPgToCsv() { exit.expectSystemExitWithStatus(0); exit.checkAssertionAfterwards(new EquivalentResultsAssertion("src/test/resources/IntegrationTest/testExportPgToCsv")); + exit.checkAssertionAfterwards(new S3EquivalentResultsAssertion(s3Path, "src/test/resources/IntegrationTest/testExportPgToCsv", "testExportPgToCsv")); final String[] command = { "nesvc", @@ -36,8 +61,10 @@ public void testExportPgToCsv() { "--json", "{"+ "\"command\": \"export-pg\",\n" + " \"params\": {\n" + - " \"endpoint\": \""+neptuneEndpoint+"\"\n" + - " }\n" + + " \"endpoint\": \""+neptuneEndpoint+"\",\n" + + " \"useIamAuth\": true\n" + + " },\n" + + " \"outputS3Path\": \""+s3Path+"/testExportPgToCsv\"\n" + "}" }; final NeptuneExportRunner runner = new NeptuneExportRunner(command); @@ -49,6 +76,7 @@ public void testExportPgML() { exit.expectSystemExitWithStatus(0); exit.checkAssertionAfterwards(new EquivalentResultsAssertion("src/test/resources/IntegrationTest/testExportPgWithEdgeAndVertexLabelsWithoutTypes")); exit.checkAssertionAfterwards(new EquivalentTrainingConfigAssertion("src/test/resources/IntegrationTest/ml-training-data-configs/v2.json")); + exit.checkAssertionAfterwards(new S3EquivalentResultsAssertion(s3Path, "src/test/resources/IntegrationTest/testExportPgWithEdgeAndVertexLabelsWithoutTypes", "testExportPgML")); final String[] command = { "nesvc", @@ -57,8 +85,10 @@ public void testExportPgML() { " \"command\": \"export-pg\",\n" + " \"params\": {\n" + " \"endpoint\": \""+neptuneEndpoint+"\",\n" + - " \"profile\": \"neptune_ml\"\n" + + " \"profile\": \"neptune_ml\",\n" + + " \"useIamAuth\": true\n" + " },\n" + + " \"outputS3Path\": \""+s3Path+"/testExportPgML\",\n" + " \"additionalParams\": {\n" + " \"neptune_ml\": {\n" + " \"version\": \"v2.0\",\n" + @@ -83,6 +113,7 @@ public void testExportPgFromQueriesML() { exit.expectSystemExitWithStatus(0); exit.checkAssertionAfterwards(new EquivalentResultsAssertion("src/test/resources/IntegrationTest/testExportPgWithEdgeAndVertexLabelsWithoutTypes")); exit.checkAssertionAfterwards(new EquivalentTrainingConfigAssertion("src/test/resources/IntegrationTest/ml-training-data-configs/v2.json")); + exit.checkAssertionAfterwards(new S3EquivalentResultsAssertion(s3Path, "src/test/resources/IntegrationTest/testExportPgWithEdgeAndVertexLabelsWithoutTypes", "testExportPgFromQueriesML")); final String[] command = { "nesvc", @@ -93,8 +124,10 @@ public void testExportPgFromQueriesML() { " \"endpoint\": \""+neptuneEndpoint+"\",\n" + " \"profile\": \"neptune_ml\",\n" + " \"query\" : \"query=g.V().union(elementMap(), outE().elementMap())\",\n" + - " \"structuredOutput\" : true \n" + + " \"structuredOutput\" : true,\n" + + " \"useIamAuth\": true\n" + " },\n" + + " \"outputS3Path\": \""+s3Path+"/testExportPgFromQueriesML\",\n" + " \"additionalParams\": {\n" + " \"neptune_ml\": {\n" + " \"version\": \"v2.0\",\n" + @@ -147,4 +180,109 @@ public void checkAssertion() throws Exception { } } + private class S3EquivalentResultsAssertion implements Assertion { + private final String s3BasePath; + private final String expectedResultsPath; + private final String testName; + + public S3EquivalentResultsAssertion(String s3BasePath, String expectedResultsPath, String testName) { + this.s3BasePath = s3BasePath; + this.expectedResultsPath = expectedResultsPath; + this.testName = testName; + } + + @Override + public void checkAssertion() throws Exception { + String fullS3Path = s3BasePath + "/" + testName; + S3ObjectInfo s3ObjectInfo = new S3ObjectInfo(fullS3Path); + + File downloadDir = tempFolder.newFolder("s3-download-" + testName); + + try { + downloadDirectoryFromS3(s3ObjectInfo, downloadDir); + + File resultDir = downloadDir.listFiles()[0]; + assertEquivalentResults(new File(expectedResultsPath), resultDir); + } finally { + deleteS3Directory(s3ObjectInfo); + } + } + } + + private void downloadDirectoryFromS3(S3ObjectInfo s3ObjectInfo, File targetDir) { + try (TransferManagerWrapper transferManager = new TransferManagerWrapper(null)) { + S3Client s3Client = S3Client.create(); + + ListObjectsV2Request listRequest = ListObjectsV2Request.builder() + .bucket(s3ObjectInfo.bucket()) + .prefix(s3ObjectInfo.key()) + .build(); + + ListObjectsV2Response listResponse = s3Client.listObjectsV2(listRequest); + + for (S3Object s3Object : listResponse.contents()) { + String key = s3Object.key(); + String relativePath = key.substring(s3ObjectInfo.key().length()); + if (relativePath.startsWith("/")) { + relativePath = relativePath.substring(1); + } + + File targetFile = new File(targetDir, relativePath); + targetFile.getParentFile().mkdirs(); + + DownloadFileRequest downloadRequest = DownloadFileRequest.builder() + .getObjectRequest(GetObjectRequest.builder() + .bucket(s3ObjectInfo.bucket()) + .key(key) + .build()) + .destination(targetFile.toPath()) + .build(); + + FileDownload download = transferManager.get().downloadFile(downloadRequest); + download.completionFuture().join(); + } + + s3Client.close(); + } catch (CancellationException | CompletionException e) { + throw new RuntimeException("Failed to download directory from S3", e); + } + } + + private void deleteS3Directory(S3ObjectInfo s3ObjectInfo) { + S3Client s3Client = S3Client.create(); + + try { + ListObjectsV2Request listRequest = ListObjectsV2Request.builder() + .bucket(s3ObjectInfo.bucket()) + .prefix(s3ObjectInfo.key()) + .build(); + + ListObjectsV2Response listResponse; + do { + listResponse = s3Client.listObjectsV2(listRequest); + + if (listResponse.hasContents()) { + List objectsToDelete = listResponse.contents().stream() + .map(s3Object -> ObjectIdentifier.builder().key(s3Object.key()).build()) + .collect(Collectors.toList()); + + DeleteObjectsRequest deleteRequest = DeleteObjectsRequest.builder() + .bucket(s3ObjectInfo.bucket()) + .delete(Delete.builder().objects(objectsToDelete).build()) + .build(); + + s3Client.deleteObjects(deleteRequest); + } + + listRequest = listRequest.toBuilder() + .continuationToken(listResponse.nextContinuationToken()) + .build(); + + } while (listResponse.isTruncated()); + + } finally { + s3Client.close(); + } + } + }