include functions on algo.list()

neo4j-contrib · Mar 5, 2019 · fac8cdc · fac8cdc
1 parent e7cbde8
commit fac8cdc
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 8 deletions.
diff --git a/algo/src/main/java/org/neo4j/graphalgo/ListProc.java b/algo/src/main/java/org/neo4j/graphalgo/ListProc.java
@@ -31,10 +31,17 @@
 
 public class ListProc {
 
-    private static final String QUERY = "CALL dbms.procedures() " +
+    private static final String QUERY =
+            " CALL dbms.procedures() " +
             " YIELD name, signature, description " +
             " WHERE name starts with 'algo.' AND name <> 'algo.list' AND ($name IS NULL OR name CONTAINS $name) " +
-            " RETURN name, signature, description ORDER BY name";
+            " RETURN name, signature, description, 'procedure' AS type " +
+            " ORDER BY name UNION " +
+            " CALL dbms.functions() " +
+            " YIELD name, signature, description " +
+            " WHERE name starts with 'algo.' AND ($name IS NULL OR name CONTAINS $name) " +
+            " RETURN name, signature, description, 'function' AS type " +
+            " ORDER BY name";
 
     @Context
     public GraphDatabaseService db;
@@ -49,11 +56,13 @@ public static class ListResult {
         public String name;
         public String description;
         public String signature;
+        public String type;
 
         public ListResult(Map<String, Object> row) {
             this.name = (String) row.get("name");
             this.description = (String) row.get("description");
             this.signature = (String) row.get("signature");
+            this.type = (String) row.get("type");
         }
     }
 }
diff --git a/tests/src/test/java/org/neo4j/graphalgo/algo/ListProcTest.java b/tests/src/test/java/org/neo4j/graphalgo/algo/ListProcTest.java
@@ -18,16 +18,25 @@
  */
 package org.neo4j.graphalgo.algo;
 
+import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.ClassRule;
 import org.junit.Test;
 import org.neo4j.graphalgo.ListProc;
 import org.neo4j.graphalgo.PageRankProc;
+import org.neo4j.graphalgo.linkprediction.LinkPrediction;
+import org.neo4j.graphdb.GraphDatabaseService;
+import org.neo4j.graphdb.factory.GraphDatabaseSettings;
 import org.neo4j.kernel.impl.proc.Procedures;
+import org.neo4j.kernel.internal.GraphDatabaseAPI;
+import org.neo4j.test.TestGraphDatabaseFactory;
 import org.neo4j.test.rule.ImpermanentDatabaseRule;
 
+import java.util.Arrays;
+import java.util.Collection;
 import java.util.List;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static java.util.Arrays.asList;
 import static java.util.Collections.*;
@@ -40,28 +49,44 @@
 public class ListProcTest {
     @ClassRule
     public static ImpermanentDatabaseRule DB = new ImpermanentDatabaseRule();
+    public static final List<String> PROCEDURES = asList("algo.pageRank", "algo.pageRank.stream");
+    public static final List<String> FUNCTIONS = Arrays.asList("algo.linkprediction.adamicAdar", "algo.linkprediction.commonNeighbors", "algo.linkprediction.preferentialAttachment", "algo.linkprediction.resourceAllocation", "algo.linkprediction.sameCommunity",
+            "algo.linkprediction.totalNeighbors");
+    public static final List<String> ALL = Stream.of(PROCEDURES, FUNCTIONS).flatMap(Collection::stream).collect(Collectors.toList());
 
     @BeforeClass
     public static void setUp() throws Exception {
         Procedures procedures = DB.getDependencyResolver().resolveDependency(Procedures.class);
         procedures.registerProcedure(ListProc.class);
         procedures.registerProcedure(PageRankProc.class);
+        procedures.registerFunction(LinkPrediction.class);
+    }
+
+    @AfterClass
+    public static void tearDown() {
+        DB.shutdown();
     }
 
     @Test
-    public void list() throws Exception {
-        assertEquals(asList("algo.pageRank","algo.pageRank.stream"), listProcs(null));
-        assertEquals(asList("algo.pageRank","algo.pageRank.stream"), listProcs("page"));
+    public void listProcedures() throws Exception {
+        assertEquals(ALL, listProcs(null));
+        assertEquals(PROCEDURES, listProcs("page"));
         assertEquals(singletonList("algo.pageRank.stream"), listProcs("stream"));
         assertEquals(emptyList(), listProcs("foo"));
     }
 
-    private List<String> listProcs(Object name) {
-        return DB.execute("CALL algo.list($name)", singletonMap("name", name)).<String>columnAs("name").stream().collect(Collectors.toList());
+    @Test
+    public void listFunctions() throws Exception {
+        assertEquals(FUNCTIONS, listProcs("linkprediction"));
     }
+
     @Test
     public void listEmpty() throws Exception {
-        assertEquals(asList("algo.pageRank","algo.pageRank.stream"),
+        assertEquals(ALL,
                 DB.execute("CALL algo.list()").<String>columnAs("name").stream().collect(Collectors.toList()));
     }
+
+    private List<String> listProcs(Object name) {
+        return DB.execute("CALL algo.list($name)", singletonMap("name", name)).<String>columnAs("name").stream().collect(Collectors.toList());
+    }
 }