From 7cdc9f7fcf8ddc63a8e86c8317c438ccd89a71cc Mon Sep 17 00:00:00 2001
From: david-leifker <114954101+david-leifker@users.noreply.github.com>
Date: Mon, 24 Oct 2022 10:49:06 -0500
Subject: [PATCH] test(neo4j): Improve test performance (#6142)

Reuse neo4j graph service in test for performance improvements.
---
 .../metadata/graph/neo4j/Neo4jGraphService.java     |  6 ++++++
 .../metadata/graph/neo4j/Neo4jGraphServiceTest.java | 13 ++++++++++---
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/metadata-io/src/main/java/com/linkedin/metadata/graph/neo4j/Neo4jGraphService.java b/metadata-io/src/main/java/com/linkedin/metadata/graph/neo4j/Neo4jGraphService.java
index 85c17a46b9..6a834352b9 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/graph/neo4j/Neo4jGraphService.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/graph/neo4j/Neo4jGraphService.java
@@ -3,6 +3,7 @@ package com.linkedin.metadata.graph.neo4j;
 import com.codahale.metrics.Timer;
 import com.datahub.util.Statement;
 import com.datahub.util.exception.RetryLimitReached;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableMap;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.metadata.graph.Edge;
@@ -254,6 +255,11 @@ public class Neo4jGraphService implements GraphService {
     removeNodesMatchingLabel(".*");
   }
 
+  @VisibleForTesting
+  public void wipe() {
+    runQuery(new Statement("MATCH (n) DETACH DELETE n", Map.of())).consume();
+  }
+
   // visible for testing
   @Nonnull
   Statement buildStatement(@Nonnull String queryTemplate, @Nonnull Map<String, Object> params) {
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/graph/neo4j/Neo4jGraphServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/graph/neo4j/Neo4jGraphServiceTest.java
index 409c6d7168..9f9ddef8a3 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/graph/neo4j/Neo4jGraphServiceTest.java
+++ b/metadata-io/src/test/java/com/linkedin/metadata/graph/neo4j/Neo4jGraphServiceTest.java
@@ -10,7 +10,8 @@ import com.linkedin.metadata.query.filter.RelationshipFilter;
 import org.neo4j.driver.Driver;
 import org.neo4j.driver.GraphDatabase;
 import org.testng.SkipException;
-import org.testng.annotations.AfterMethod;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
@@ -28,15 +29,21 @@ public class Neo4jGraphServiceTest extends GraphServiceTestBase {
   private Driver _driver;
   private Neo4jGraphService _client;
 
-  @BeforeMethod
+  @BeforeClass
   public void init() {
     _serverBuilder = new Neo4jTestServerBuilder();
     _serverBuilder.newServer();
     _driver = GraphDatabase.driver(_serverBuilder.boltURI());
     _client = new Neo4jGraphService(new LineageRegistry(SnapshotEntityRegistry.getInstance()), _driver);
+    _client.clear();
   }
 
-  @AfterMethod
+  @BeforeMethod
+  public void wipe() {
+    _client.wipe();
+  }
+
+  @AfterClass
   public void tearDown() {
     _serverBuilder.shutdown();
   }