diff --git a/pom.xml b/pom.xml
index 9e80263..8e979df 100644
--- a/pom.xml
+++ b/pom.xml
@@ -277,6 +277,13 @@
6.12
test
+
+
+ org.apache.commons
+ commons-collections4
+ 4.4
+ test
+
diff --git a/src/test/java/com/amazonaws/services/neptune/propertygraph/LazyQueriesRangeFactoryProviderTest.java b/src/test/java/com/amazonaws/services/neptune/propertygraph/LazyQueriesRangeFactoryProviderTest.java
new file mode 100644
index 0000000..7810a6b
--- /dev/null
+++ b/src/test/java/com/amazonaws/services/neptune/propertygraph/LazyQueriesRangeFactoryProviderTest.java
@@ -0,0 +1,110 @@
+/*
+Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License").
+You may not use this file except in compliance with the License.
+A copy of the License is located at
+ http://www.apache.org/licenses/LICENSE-2.0
+or in the "license" file accompanying this file. This file 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.amazonaws.services.neptune.propertygraph;
+
+import com.amazonaws.services.neptune.cluster.ConcurrencyConfig;
+import com.amazonaws.services.neptune.export.FeatureToggles;
+import com.amazonaws.services.neptune.propertygraph.schema.ExportSpecification;
+import com.amazonaws.services.neptune.propertygraph.schema.GraphElementType;
+import org.apache.tinkerpop.gremlin.process.traversal.dsl.graph.GraphTraversalSource;
+import org.apache.tinkerpop.gremlin.tinkergraph.structure.TinkerFactory;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.times;
+
+public class LazyQueriesRangeFactoryProviderTest {
+
+ private NeptuneGremlinClient mockClient;
+ private LazyQueriesRangeFactoryProvider rangeFactoryProvider;
+ private GraphTraversalSource g;
+
+ @Before
+ public void setup() {
+ mockClient = mock(NeptuneGremlinClient.class);
+ g = spy(TinkerFactory.createModern().traversal());
+ when(mockClient.newTraversalSource()).thenReturn(g);
+
+ rangeFactoryProvider = new LazyQueriesRangeFactoryProvider(
+ new RangeConfig(4, 0, Long.MAX_VALUE, -1, -1),
+ new ConcurrencyConfig(4),
+ mockClient,
+ createExportSpecifications(),
+ new FeatureToggles(Collections.EMPTY_SET)
+ );
+ }
+
+ @Test
+ public void shouldOnlyCountNodesWhenFirstRequested() {
+ // Validate lazy initialization (count and creation of g is only done when requested)
+ verify(mockClient, times(0)).newTraversalSource();
+ verify(g, times(0)).V();
+
+ // Validate traversal source is created and queried after first invocation
+ rangeFactoryProvider.getNodesRangeFactory();
+ verify(mockClient, times(1)).newTraversalSource();
+ verify(g, times(1)).V();
+
+ // Validate traversal source is not recreated or re-queried on second invocation
+ rangeFactoryProvider.getNodesRangeFactory();
+ verify(mockClient, times(1)).newTraversalSource();
+ verify(g, times(1)).V();
+ }
+
+ @Test
+ public void shouldOnlyCountEdgesWhenFirstRequested() {
+ // Validate lazy initialization (count and creation of g is only done when requested)
+ verify(mockClient, times(0)).newTraversalSource();
+ verify(g, times(0)).E();
+
+ // Validate traversal source is created and queried after first invocation
+ rangeFactoryProvider.getEdgesRangeFactory();
+ verify(mockClient, times(1)).newTraversalSource();
+ verify(g, times(1)).E();
+
+ // Validate traversal source is not recreated or re-queried on second invocation
+ rangeFactoryProvider.getEdgesRangeFactory();
+ verify(mockClient, times(1)).newTraversalSource();
+ verify(g, times(1)).E();
+ }
+
+ private Collection createExportSpecifications() {
+ Collection exportSpecifications = new ArrayList<>();
+ exportSpecifications.add(new ExportSpecification(
+ GraphElementType.nodes,
+ new AllLabels(NodeLabelStrategy.nodeLabelsOnly),
+ GremlinFilters.EMPTY,
+ null,
+ false,
+ new FeatureToggles(Collections.EMPTY_SET)
+ ));
+ exportSpecifications.add(new ExportSpecification(
+ GraphElementType.edges,
+ new AllLabels(NodeLabelStrategy.nodeLabelsOnly),
+ GremlinFilters.EMPTY,
+ null,
+ false,
+ new FeatureToggles(Collections.EMPTY_SET)
+ ));
+ return exportSpecifications;
+ }
+
+}
diff --git a/src/test/java/com/amazonaws/services/neptune/propertygraph/NamedQueriesCollectionTest.java b/src/test/java/com/amazonaws/services/neptune/propertygraph/NamedQueriesCollectionTest.java
new file mode 100644
index 0000000..1fe30cc
--- /dev/null
+++ b/src/test/java/com/amazonaws/services/neptune/propertygraph/NamedQueriesCollectionTest.java
@@ -0,0 +1,50 @@
+/*
+Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License").
+You may not use this file except in compliance with the License.
+A copy of the License is located at
+ http://www.apache.org/licenses/LICENSE-2.0
+or in the "license" file accompanying this file. This file 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.amazonaws.services.neptune.propertygraph;
+
+import com.amazonaws.services.neptune.cluster.ConcurrencyConfig;
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.Collection;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
+
+public class NamedQueriesCollectionTest {
+ @Test
+ public void splitQueriesShouldSplitAllNamedQueries() {
+ Collection namedQueries = new ArrayList<>();
+ namedQueries.add(mock(NamedQueries.class));
+ namedQueries.add(mock(NamedQueries.class));
+ namedQueries.add(mock(NamedQueries.class));
+ NamedQueriesCollection namedQueriesCollection = new NamedQueriesCollection(namedQueries);
+
+ LazyQueriesRangeFactoryProvider rangeFactoryProvider = NamedQueriesTest.initRangeFactoryProvider(mock(RangeConfig.class), mock(ConcurrencyConfig.class));
+
+ for(NamedQueries nq: namedQueries) {
+ verify(nq, times(0)).split(any());
+ }
+
+ namedQueriesCollection.splitQueries(rangeFactoryProvider);
+
+ for(NamedQueries nq: namedQueries) {
+ verify(nq, times(1)).split(rangeFactoryProvider);
+ verify(nq, times(1)).split(any());
+ verifyNoMoreInteractions(nq);
+ }
+ }
+}
diff --git a/src/test/java/com/amazonaws/services/neptune/propertygraph/NamedQueriesTest.java b/src/test/java/com/amazonaws/services/neptune/propertygraph/NamedQueriesTest.java
new file mode 100644
index 0000000..0341da6
--- /dev/null
+++ b/src/test/java/com/amazonaws/services/neptune/propertygraph/NamedQueriesTest.java
@@ -0,0 +1,175 @@
+/*
+Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License").
+You may not use this file except in compliance with the License.
+A copy of the License is located at
+ http://www.apache.org/licenses/LICENSE-2.0
+or in the "license" file accompanying this file. This file 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.amazonaws.services.neptune.propertygraph;
+
+import com.amazonaws.services.neptune.cluster.ConcurrencyConfig;
+import com.amazonaws.services.neptune.export.FeatureToggles;
+import com.amazonaws.services.neptune.propertygraph.schema.ExportSpecification;
+import com.amazonaws.services.neptune.propertygraph.schema.GraphElementType;
+import org.apache.commons.collections.CollectionUtils;
+import org.apache.commons.collections.collection.UnmodifiableCollection;
+import org.apache.tinkerpop.gremlin.tinkergraph.structure.TinkerFactory;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class NamedQueriesTest {
+
+ Collection testQueries;
+ final Collection sampleQueries;
+
+ public NamedQueriesTest() {
+ Collection sampleData = new ArrayList<>();
+ sampleData.add("g.V().hasLabel(\"person\")");
+ sampleData.add("g.V().out()");
+ sampleData.add("g.E().outV()");
+ sampleQueries = UnmodifiableCollection.decorate(sampleData);
+ }
+
+ @Before
+ public void setupClass() {
+ testQueries = new ArrayList<>();
+ testQueries.addAll(sampleQueries);
+ }
+
+ @Test
+ public void shouldNotModifyQueriesWithEffectiveConcurrency1() {
+ NamedQueries namedQueries = new NamedQueries("name", testQueries);
+ namedQueries.split(initRangeFactoryProvider(
+ new RangeConfig(-1, 0, Long.MAX_VALUE, -1, -1),
+ new ConcurrencyConfig(1)
+ ));
+
+ assertTrue(CollectionUtils.isEqualCollection(sampleQueries, namedQueries.queries()));
+ }
+
+ @Test
+ public void shouldSplitQueriesAccordingToRangeConfig() {
+ NamedQueries namedQueries = new NamedQueries("name", testQueries);
+ namedQueries.split(initRangeFactoryProvider(
+ new RangeConfig(5, 0, Long.MAX_VALUE, -1, -1),
+ new ConcurrencyConfig(4)
+ ));
+
+ Collection expectedQueries = new ArrayList<>();
+ expectedQueries.add("g.V().range(0, 5).hasLabel(\"person\")");
+ expectedQueries.add("g.V().range(5, -1).hasLabel(\"person\")");
+ expectedQueries.add("g.V().range(0, 5).out()");
+ expectedQueries.add("g.V().range(5, -1).out()");
+ expectedQueries.add("g.E().range(0, 5).outV()");
+ expectedQueries.add("g.E().range(5, 10).outV()");
+ expectedQueries.add("g.E().range(10, -1).outV()");
+
+ assertTrue(CollectionUtils.isEqualCollection(expectedQueries, namedQueries.queries()));
+ }
+
+ @Test
+ public void shouldOnlySplitEdgeQueries() {
+ NamedQueries namedQueries = new NamedQueries("name", testQueries);
+ namedQueries.split(initRangeFactoryProvider(
+ new RangeConfig(10, 0, Long.MAX_VALUE, -1, -1),
+ new ConcurrencyConfig(4)
+ ));
+
+ Collection expectedQueries = new ArrayList<>();
+ expectedQueries.add("g.V().hasLabel(\"person\")");
+ expectedQueries.add("g.V().out()");
+ expectedQueries.add("g.E().range(0, 10).outV()");
+ expectedQueries.add("g.E().range(10, -1).outV()");
+
+ assertTrue(CollectionUtils.isEqualCollection(expectedQueries, namedQueries.queries()));
+ }
+
+ @Test
+ public void shouldOnlySplitVertexQueries() {
+ NamedQueries namedQueries = new NamedQueries("name", testQueries);
+ namedQueries.split(initRangeFactoryProvider(
+ new RangeConfig(15, 0, Long.MAX_VALUE, 30, -1),
+ new ConcurrencyConfig(4)
+ ));
+
+ Collection expectedQueries = new ArrayList<>();
+ expectedQueries.add("g.V().range(0, 15).hasLabel(\"person\")");
+ expectedQueries.add("g.V().range(15, -1).hasLabel(\"person\")");
+ expectedQueries.add("g.V().range(0, 15).out()");
+ expectedQueries.add("g.V().range(15, -1).out()");
+ expectedQueries.add("g.E().outV()");
+
+ assertTrue(CollectionUtils.isEqualCollection(expectedQueries, namedQueries.queries()));
+ }
+
+ @Test
+ public void shouldSplitBasedOnApproxCountsAndConcurrency() {
+ NamedQueries namedQueries = new NamedQueries("name", testQueries);
+ namedQueries.split(initRangeFactoryProvider(
+ new RangeConfig(-1, 0, Long.MAX_VALUE, 1200, 1500),
+ new ConcurrencyConfig(3)
+ ));
+
+ Collection expectedQueries = new ArrayList<>();
+ expectedQueries.add("g.V().range(0, 401).hasLabel(\"person\")");
+ expectedQueries.add("g.V().range(401, 802).hasLabel(\"person\")");
+ expectedQueries.add("g.V().range(802, -1).hasLabel(\"person\")");
+ expectedQueries.add("g.V().range(0, 401).out()");
+ expectedQueries.add("g.V().range(401, 802).out()");
+ expectedQueries.add("g.V().range(802, -1).out()");
+ expectedQueries.add("g.E().range(0, 501).outV()");
+ expectedQueries.add("g.E().range(501, 1002).outV()");
+ expectedQueries.add("g.E().range(1002, -1).outV()");
+
+ assertTrue(CollectionUtils.isEqualCollection(expectedQueries, namedQueries.queries()));
+ }
+
+ static LazyQueriesRangeFactoryProvider initRangeFactoryProvider(RangeConfig rangeConfig, ConcurrencyConfig concurrencyConfig) {
+ NeptuneGremlinClient mockClient = mock(NeptuneGremlinClient.class);
+ when(mockClient.newTraversalSource()).thenReturn(TinkerFactory.createTheCrew().traversal());
+
+ return new LazyQueriesRangeFactoryProvider(
+ rangeConfig,
+ concurrencyConfig,
+ mockClient,
+ createExportSpecifications(),
+ new FeatureToggles(Collections.EMPTY_SET)
+ );
+ }
+
+ private static Collection createExportSpecifications() {
+ Collection exportSpecifications = new ArrayList<>();
+ exportSpecifications.add(new ExportSpecification(
+ GraphElementType.nodes,
+ new AllLabels(NodeLabelStrategy.nodeLabelsOnly),
+ GremlinFilters.EMPTY,
+ null,
+ false,
+ new FeatureToggles(Collections.EMPTY_SET)
+ ));
+ exportSpecifications.add(new ExportSpecification(
+ GraphElementType.edges,
+ new AllLabels(NodeLabelStrategy.nodeLabelsOnly),
+ GremlinFilters.EMPTY,
+ null,
+ false,
+ new FeatureToggles(Collections.EMPTY_SET)
+ ));
+ return exportSpecifications;
+ }
+
+}