Refactor NativeVactory construction

Co-authored-by: Ioannis Panagiotas<ioannis.panagiotas@neotechnology.com>

Author: vnickolov

native-projection/src/main/java/org/neo4j/gds/projection/NativeFactory.java

@@ -41,8 +41,6 @@
 
 import java.util.Optional;
 
-import static org.neo4j.gds.projection.GraphDimensionsValidation.validate;
-
 final class NativeFactory extends CSRGraphStoreFactory<GraphProjectFromStoreConfig> {
 
     private final GraphProjectFromStoreConfig storeConfig;
@@ -64,18 +62,21 @@ static NativeFactory nativeFactory(
         return new NativeFactory(
             graphProjectFromStoreConfig,
             loadingContext,
-            dimensions
+            dimensions,
+            initProgressTracker(graphProjectFromStoreConfig, loadingContext, dimensions)
         );
     }
 
-    private NativeFactory(
+    // Package-private constructor used in tests
+    NativeFactory(
         GraphProjectFromStoreConfig graphProjectConfig,
         GraphLoaderContext loadingContext,
-        GraphDimensions graphDimensions
+        GraphDimensions graphDimensions,
+        ProgressTracker progressTracker
     ) {
         super(graphProjectConfig, ImmutableStaticCapabilities.of(WriteMode.LOCAL), loadingContext, graphDimensions);
         this.storeConfig = graphProjectConfig;
-        this.progressTracker = initProgressTracker();
+        this.progressTracker = progressTracker;
     }
 
     @Override
@@ -88,7 +89,10 @@ public MemoryEstimation estimateMemoryUsageAfterLoading() {
         return getMemoryEstimation(storeConfig.nodeProjections(), storeConfig.relationshipProjections(), false);
     }
 
-    private ProgressTracker initProgressTracker() {
+    private static ProgressTracker initProgressTracker(
+        GraphProjectFromStoreConfig graphProjectConfig,
+        GraphLoaderContext loadingContext, GraphDimensions dimensions
+    ) {
         long relationshipCount = graphProjectConfig
             .relationshipProjections()
             .projections()
@@ -137,7 +141,7 @@ public CSRGraphStore build() {
             // Start the sub-task so it will be registered in the task store
             progressTracker.beginSubTask();
             // `validate` can raise an exception which has to be handled in order to end the sub-task with failure
-            validate(dimensions, storeConfig);
+            validate();
             var concurrency = graphProjectConfig.readConcurrency();
             Nodes nodes = loadNodes(concurrency);
             RelationshipImportResult relationships = loadRelationships(nodes.idMap(), concurrency);
@@ -153,6 +157,11 @@ public CSRGraphStore build() {
         }
     }
 
+    // Allows for mocking the validation behaviour.
+    void validate() throws IllegalArgumentException {
+        GraphDimensionsValidation.validate(dimensions, storeConfig);
+    }
+
     private Nodes loadNodes(Concurrency concurrency) {
         var scanningNodesImporter = new ScanningNodesImporterBuilder()
             .concurrency(concurrency)

native-projection/src/test/java/org/neo4j/gds/projection/NativeFactoryTest.java

@@ -19,12 +19,8 @@
  */
 package org.neo4j.gds.projection;
 
-import com.carrotsearch.hppc.LongHashSet;
-import org.agrona.collections.MutableBoolean;
 import org.junit.jupiter.api.Test;
 import org.neo4j.gds.ImmutableRelationshipProjections;
-import org.neo4j.gds.NodeLabel;
-import org.neo4j.gds.NodeProjection;
 import org.neo4j.gds.NodeProjections;
 import org.neo4j.gds.Orientation;
 import org.neo4j.gds.RelationshipProjection;
@@ -35,24 +31,20 @@
 import org.neo4j.gds.core.GraphDimensions;
 import org.neo4j.gds.core.ImmutableGraphDimensions;
 import org.neo4j.gds.core.concurrency.Concurrency;
-import org.neo4j.gds.core.utils.progress.JobId;
-import org.neo4j.gds.core.utils.progress.TaskRegistry;
-import org.neo4j.gds.core.utils.progress.TaskRegistryFactory;
-import org.neo4j.gds.logging.Log;
+import org.neo4j.gds.core.utils.progress.tasks.ProgressTracker;
 import org.neo4j.gds.mem.MemoryEstimation;
 import org.neo4j.gds.mem.MemoryTree;
 
-import java.util.Map;
-import java.util.Optional;
 import java.util.concurrent.atomic.AtomicReference;
 
-import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
 
 class NativeFactoryTest {
 
@@ -161,65 +153,24 @@ void memoryEstimationForMultipleProjections() {
     }
 
     @Test
-    void shouldCleanTaskOnValidateFailure(){
+    void shouldCleanTaskOnValidateFailure() {
+        var progressTrackerMock = mock(ProgressTracker.class);
 
-        var mockRegistry = mock(TaskRegistry.class);
+        var nativeFactorySpy = spy(new NativeFactory(
+            mock(GraphProjectFromStoreConfig.class),
+            mock(GraphLoaderContext.class),
+            mock(GraphDimensions.class),
+            progressTrackerMock
+        ));
 
-        var mockRegistryFactory = mock(TaskRegistryFactory.class);
-        when(mockRegistryFactory.newInstance(any())).thenReturn(mockRegistry);
-        var mockgraphLoaderContext = mock(GraphLoaderContext.class);
-        when(mockgraphLoaderContext.taskRegistryFactory()).thenReturn(mockRegistryFactory);
-        when(mockgraphLoaderContext.log()).thenReturn(Log.noOpLog());
+        doThrow(new IllegalArgumentException("Intentionally failing validation")).when(nativeFactorySpy).validate();
 
-        var labelSet = new LongHashSet();
-        labelSet.add(GraphDimensions.NO_SUCH_LABEL);
-        var graphDimensions =ImmutableGraphDimensions.builder()
-            .nodeCount(100)
-            .nodeLabelTokens(labelSet)
-            .build();
-
-        var config = mock(GraphProjectFromStoreConfig.class);
-        when(config.relationshipProjections()).thenReturn(new RelationshipProjections() {
-            @Override
-            public Map<RelationshipType, RelationshipProjection> projections() {
-                return Map.of();
-            }
-        });
-
-        when(config.nodeProjections()).thenReturn(new NodeProjections() {
-            @Override
-            public Map<NodeLabel, NodeProjection> projections() {
-                return Map.of(
-                    NodeLabel.of("foo"), new NodeProjection() {
-                        @Override
-                        public String label() {
-                            return "bar";
-                        }
-                    }
-                );
-            }
-        });
-
-        when(config.logProgress()).thenReturn(true);
-        when(config.jobId()).thenReturn(new JobId("foo"));
-        when(config.readConcurrency()).thenReturn(new Concurrency(1));
-
-        var nativeFactory = NativeFactory.nativeFactory(
-            config,
-            mockgraphLoaderContext,
-            Optional.of(graphDimensions)
-        );
-
-        MutableBoolean failed = new MutableBoolean(false);
-        try {
-            nativeFactory.build();
-        } catch (Exception ignored){
-            failed.set(true);
-        }
-
-        assertThat(failed.get()).isTrue();
-        verify(mockRegistry,times(1)).registerTask(any());
-        verify(mockRegistry,times(1)).markCompleted();
+        assertThatIllegalArgumentException()
+            .isThrownBy(nativeFactorySpy::build)
+            .withMessageContaining("Intentionally failing validation");
 
+        verify(progressTrackerMock, times(1)).beginSubTask();
+        verify(progressTrackerMock, times(1)).endSubTaskWithFailure();
+        verifyNoMoreInteractions(progressTrackerMock);
     }
 }