Fix failing tests

Author: soerenreichardt

core/src/test/java/org/neo4j/gds/core/utils/progress/BaseTaskRegistryExtensionTest.java

@@ -64,11 +64,12 @@ void test() throws Exception {
 
     public static class AlgoProc {
         @Context
-        public TaskRegistry taskRegistry;
+        public TaskRegistryFactory taskRegistryFactory;
 
         @Procedure("gds.test.algo")
         public Stream<Bar> foo() {
             var task = Tasks.leaf("foo");
+            var taskRegistry = taskRegistryFactory.newInstance();
             taskRegistry.registerTask(task);
             return Stream.empty();
         }

proc/pregel/src/test/java/org/neo4j/gds/pregel/proc/PregelProcTest.java

@@ -228,8 +228,7 @@ void cleanupTaskRegistryWhenTheAlgorithmFailsInStreamMode() {
             );
 
             assertThatThrownBy(() -> proc.stream(config, Map.of())).isNotNull();
-            // TODO: set to 1 once tracking is implemented
-            assertThat(taskStore.storeTaskCalls()).isEqualTo(0);
+            assertThat(taskStore.storeTaskCalls()).isEqualTo(1);
         }
     }
 
@@ -252,8 +251,7 @@ void cleanupTaskRegistryWhenTheAlgorithmFailsInWriteMode() {
             );
 
             assertThatThrownBy(() -> proc.write(config, Map.of())).isNotNull();
-            // TODO: set to 1 once tracking is implemented
-            assertThat(taskStore.storeTaskCalls()).isEqualTo(0);
+            assertThat(taskStore.storeTaskCalls()).isEqualTo(1);
         }
     }
 
@@ -277,8 +275,7 @@ void cleanupTaskRegistryWhenTheAlgorithmFailsInMutateMode() {
             );
 
             assertThatThrownBy(() -> proc.mutate(graphName, config)).isNotNull();
-            // TODO: set to 1 once tracking is implemented
-            assertThat(taskStore.storeTaskCalls()).isEqualTo(0);
+            assertThat(taskStore.storeTaskCalls()).isEqualTo(1);
         }
     }
 

proc/test/src/main/java/org/neo4j/gds/AlgoBaseProcTest.java

@@ -44,7 +44,7 @@
 import org.neo4j.gds.core.utils.progress.EmptyTaskRegistryFactory;
 import org.neo4j.gds.core.utils.progress.GlobalTaskStore;
 import org.neo4j.gds.core.utils.progress.JobId;
-import org.neo4j.gds.core.utils.progress.LocalTaskRegistry;
+import org.neo4j.gds.core.utils.progress.TaskRegistry;
 import org.neo4j.gds.core.utils.progress.tasks.Task;
 import org.neo4j.gds.core.write.NativeNodePropertyExporter;
 import org.neo4j.gds.core.write.NativeRelationshipExporter;
@@ -292,7 +292,7 @@ default void shouldUnregisterTaskAfterComputation() {
         String loadedGraphName = "loadedGraph";
         GraphCreateConfig graphCreateConfig = withNameAndRelationshipProjections("", loadedGraphName, relationshipProjections());
         applyOnProcedure(proc -> {
-            proc.taskRegistryFactory = () -> new LocalTaskRegistry("", taskStore);
+            proc.taskRegistryFactory = () -> new TaskRegistry("", taskStore);
 
             GraphStore graphStore = graphLoader(graphCreateConfig).graphStore();
             GraphStoreCatalog.set(