add clear graph data method in unit test (#99)

* add clear graph data method
apache · Sep 12, 2021 · c4c71a2 · c4c71a2
1 parent b254f7c
commit c4c71a2
Showing 1 changed file with 49 additions and 0 deletions.
diff --git a/computer-test/src/main/java/com/baidu/hugegraph/computer/suite/unit/UnitTestBase.java b/computer-test/src/main/java/com/baidu/hugegraph/computer/suite/unit/UnitTestBase.java
@@ -21,6 +21,7 @@
 
 import java.io.IOException;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Random;
 
@@ -50,6 +51,14 @@
 import com.baidu.hugegraph.computer.core.util.ComputerContextUtil;
 import com.baidu.hugegraph.computer.core.worker.MockComputationParams;
 import com.baidu.hugegraph.config.TypedOption;
+import com.baidu.hugegraph.driver.HugeClient;
+import com.baidu.hugegraph.driver.SchemaManager;
+import com.baidu.hugegraph.structure.graph.Edge;
+import com.baidu.hugegraph.structure.graph.Vertex;
+import com.baidu.hugegraph.structure.schema.EdgeLabel;
+import com.baidu.hugegraph.structure.schema.IndexLabel;
+import com.baidu.hugegraph.structure.schema.PropertyKey;
+import com.baidu.hugegraph.structure.schema.VertexLabel;
 import com.baidu.hugegraph.testutil.Assert;
 import com.baidu.hugegraph.util.E;
 import com.baidu.hugegraph.util.Log;
@@ -62,6 +71,42 @@ public class UnitTestBase {
                                         "0123456789" +
                                         "abcdefghijklmnopqrstuvxyz";
 
+    private static final String URL = ComputerOptions.HUGEGRAPH_URL
+                                                     .defaultValue();
+    private static final String GRAPH = ComputerOptions.HUGEGRAPH_GRAPH_NAME
+                                                       .defaultValue();
+    private static final HugeClient CLIENT = HugeClient.builder(URL, GRAPH)
+                                                       .build();
+
+    protected static void clearAll() {
+        clearData();
+        clearSchema();
+    }
+
+    protected static void clearData() {
+        List<Edge> edges = CLIENT.graph().listEdges();
+        edges.forEach(edge -> CLIENT.graph().removeEdge(edge.id()));
+
+        List<Vertex> vertices = CLIENT.graph().listVertices();
+        vertices.forEach(vertex -> CLIENT.graph().removeVertex(vertex.id()));
+    }
+
+    private static void clearSchema() {
+        SchemaManager schema = CLIENT.schema();
+
+        List<IndexLabel> indexLabels = schema.getIndexLabels();
+        indexLabels.forEach(label -> schema.removeIndexLabel(label.name()));
+
+        List<EdgeLabel> edgeLabels = schema.getEdgeLabels();
+        edgeLabels.forEach(label -> schema.removeEdgeLabel(label.name()));
+
+        List<VertexLabel> vertexLabels = schema.getVertexLabels();
+        vertexLabels.forEach(label -> schema.removeVertexLabel(label.name()));
+
+        List<PropertyKey> propertyKeys = schema.getPropertyKeys();
+        propertyKeys.forEach(label -> schema.removePropertyKey(label.name()));
+    }
+
     public static void assertIdEqualAfterWriteAndRead(Id oldId)
                                                       throws IOException {
         byte[] bytes;
@@ -196,6 +241,10 @@ protected static StreamGraphOutput newStreamGraphOutput(
                                                                output);
     }
 
+    protected static HugeClient client() {
+        return CLIENT;
+    }
+
     public static boolean existError(Throwable[] exceptions) {
         boolean error = false;
         for (Throwable e : exceptions) {