added tests

pom.xml

@@ -277,6 +277,13 @@
             <version>6.12</version>
             <scope>test</scope>
         </dependency>
+
+        <dependency>
+            <groupId>org.apache.commons</groupId>
+            <artifactId>commons-collections4</artifactId>
+            <version>4.4</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>

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<ExportSpecification> createExportSpecifications() {
+        Collection<ExportSpecification> 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;
+    }
+
+}

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> 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);
+        }
+    }
+}

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<String> testQueries;
+    final Collection<String> sampleQueries;
+
+    public NamedQueriesTest() {
+        Collection<String> 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<String> 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<String> 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<String> 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<String> 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<ExportSpecification> createExportSpecifications() {
+        Collection<ExportSpecification> 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;
+    }
+
+}