visitor offers method to delete current element

junit-team · Nov 9, 2015 · 1d6a7d5 · 1d6a7d5
1 parent e79a619
commit 1d6a7d5
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 14 deletions.
diff --git a/junit-engine-api/src/main/java/org/junit/gen5/engine/AbstractTestDescriptor.java b/junit-engine-api/src/main/java/org/junit/gen5/engine/AbstractTestDescriptor.java
@@ -51,12 +51,13 @@ protected void removeChild(AbstractTestDescriptor abstractTestDescriptor) {
 		children.remove(abstractTestDescriptor);
 	}
 
-	@Override
+	protected void removeFromHierarchy() {
-	public void remove() {
+		if (isRoot())
+			return;
 		if (parent instanceof AbstractTestDescriptor)
 			((AbstractTestDescriptor) parent).removeChild(this);
 		setParent(null);
-		children.clear(); //to prevent visiting
+		children.clear();
 	}
 
 	@Override
@@ -74,7 +75,8 @@ public final Set<TestDescriptor> getChildren() {
 
 	@Override
 	public void accept(Visitor visitor) {
-		visitor.visit(this);
+		Runnable delete = () -> removeFromHierarchy();
+		visitor.visit(this, delete);
 		new LinkedHashSet<>(getChildren()).forEach(child -> child.accept(visitor));
 	}
 

diff --git a/junit-engine-api/src/main/java/org/junit/gen5/engine/TestDescriptor.java b/junit-engine-api/src/main/java/org/junit/gen5/engine/TestDescriptor.java
@@ -32,17 +32,19 @@ public interface TestDescriptor {
 
 	boolean isTest();
 
+	default boolean isRoot() {
+		return getParent() == null;
+	}
+
 	Set<TestTag> getTags();
 
 	void addChild(TestDescriptor descriptor);
 
-	void remove();
-
 	Set<TestDescriptor> getChildren();
 
 	interface Visitor {
 
-		void visit(TestDescriptor descriptor);
+		void visit(TestDescriptor descriptor, Runnable deleteCurrent);
 	}
 
 	void accept(Visitor visitor);

diff --git a/junit-engine-api/src/test/java/org/junit/gen5/engine/AbstractTestDescriptorTest.java b/junit-engine-api/src/test/java/org/junit/gen5/engine/AbstractTestDescriptorTest.java
@@ -61,21 +61,21 @@ public String getId() {
 	@Test
 	public void visitAllNodes() {
 		List<TestDescriptor> visited = new ArrayList<>();
-		engineDescriptor.accept(visited::add);
+		engineDescriptor.accept((descriptor, delete) -> visited.add(descriptor));
 
 		Assert.assertEquals(8, visited.size());
 	}
 
 	@Test
 	public void pruneLeaf() {
-		TestDescriptor.Visitor visitor = (TestDescriptor descriptor) -> {
+		TestDescriptor.Visitor visitor = (TestDescriptor descriptor, Runnable delete) -> {
 			if (descriptor.getUniqueId().equals("leaf1-1"))
-				descriptor.remove();
+				delete.run();
 		};
 		engineDescriptor.accept(visitor);
 
 		List<String> visited = new ArrayList<>();
-		engineDescriptor.accept((descriptor) -> visited.add(descriptor.getUniqueId()));
+		engineDescriptor.accept((descriptor, delete) -> visited.add(descriptor.getUniqueId()));
 
 		Assert.assertEquals(7, visited.size());
 		Assert.assertTrue(visited.contains("group1"));
@@ -85,17 +85,17 @@ public void pruneLeaf() {
 	@Test
 	public void pruneGroup() {
 		final AtomicInteger countVisited = new AtomicInteger();
-		TestDescriptor.Visitor visitor = (TestDescriptor descriptor) -> {
+		TestDescriptor.Visitor visitor = (descriptor, delete) -> {
 			if (descriptor.getUniqueId().equals("group1"))
-				descriptor.remove();
+				delete.run();
 			countVisited.incrementAndGet();
 		};
 		engineDescriptor.accept(visitor);
 
 		Assert.assertEquals("Children of pruned element are not visited", 4, countVisited.get());
 
 		List<String> visited = new ArrayList<>();
-		engineDescriptor.accept((descriptor) -> visited.add(descriptor.getUniqueId()));
+		engineDescriptor.accept((descriptor, delete) -> visited.add(descriptor.getUniqueId()));
 
 		Assert.assertEquals(3, visited.size());
 		Assert.assertFalse(visited.contains("group1"));