Order test descriptor children in place

documentation/src/docs/asciidoc/release-notes/release-notes-5.12.0-RC1.adoc

@@ -26,7 +26,8 @@ JUnit repository on GitHub.
 [[release-notes-5.12.0-RC1-junit-platform-new-features-and-improvements]]
 ==== New Features and Improvements
 
-* ❓
+* New `TestDescriptor.orderChildren(UnaryOperator<List<TestDescriptor>> orderer)`
+  method to order children in place
 
 
 [[release-notes-5.12.0-RC1-junit-jupiter]]

junit-jupiter-engine/src/main/java/org/junit/jupiter/engine/discovery/AbstractOrderingVisitor.java

@@ -10,15 +10,16 @@
 
 package org.junit.jupiter.engine.discovery;
 
+import static java.util.Comparator.comparing;
 import static java.util.stream.Collectors.toCollection;
+import static java.util.stream.Collectors.toList;
 
 import java.util.ArrayList;
-import java.util.LinkedHashSet;
+import java.util.HashMap;
 import java.util.List;
-import java.util.Set;
+import java.util.Map;
 import java.util.function.Consumer;
 import java.util.function.Function;
-import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import org.junit.jupiter.engine.descriptor.ClassBasedTestDescriptor;
@@ -60,9 +61,8 @@ protected void doWithMatchingDescriptor(Class<PARENT> parentTestDescriptorType,
 	protected void orderChildrenTestDescriptors(TestDescriptor parentTestDescriptor, Class<CHILD> matchingChildrenType,
 			Function<CHILD, WRAPPER> descriptorWrapperFactory, DescriptorWrapperOrderer descriptorWrapperOrderer) {
 
-		Set<? extends TestDescriptor> children = parentTestDescriptor.getChildren();
-
-		List<WRAPPER> matchingDescriptorWrappers = children.stream()//
+		List<WRAPPER> matchingDescriptorWrappers = parentTestDescriptor.getChildren()//
+				.stream()//
 				.filter(matchingChildrenType::isInstance)//
 				.map(matchingChildrenType::cast)//
 				.map(descriptorWrapperFactory)//
@@ -74,50 +74,33 @@ protected void orderChildrenTestDescriptors(TestDescriptor parentTestDescriptor,
 		}
 
 		if (descriptorWrapperOrderer.canOrderWrappers()) {
-			List<TestDescriptor> nonMatchingTestDescriptors = children.stream()//
-					.filter(childTestDescriptor -> !matchingChildrenType.isInstance(childTestDescriptor))//
-					.collect(Collectors.toList());
-
-			// Make a local copy for later validation
-			Set<WRAPPER> originalWrappers = new LinkedHashSet<>(matchingDescriptorWrappers);
-
-			descriptorWrapperOrderer.orderWrappers(matchingDescriptorWrappers);
-
-			int difference = matchingDescriptorWrappers.size() - originalWrappers.size();
-			if (difference > 0) {
-				descriptorWrapperOrderer.logDescriptorsAddedWarning(difference);
-			}
-			else if (difference < 0) {
-				descriptorWrapperOrderer.logDescriptorsRemovedWarning(difference);
-			}
-
-			Set<TestDescriptor> orderedTestDescriptors = matchingDescriptorWrappers.stream()//
-					.filter(originalWrappers::contains)//
-					.map(AbstractAnnotatedDescriptorWrapper::getTestDescriptor)//
-					.collect(toCollection(LinkedHashSet::new));
-
-			// There is currently no way to removeAll or addAll children at once, so we
-			// first remove them all and then add them all back.
-			Stream.concat(orderedTestDescriptors.stream(), nonMatchingTestDescriptors.stream())//
-					.forEach(parentTestDescriptor::removeChild);
-
-			// If we are ordering children of type ClassBasedTestDescriptor, that means we
-			// are ordering top-level classes or @Nested test classes. Thus, the
-			// nonMatchingTestDescriptors list is either empty (for top-level classes) or
-			// contains only local test methods (for @Nested classes) which must be executed
-			// before tests in @Nested test classes. So we add the test methods before adding
-			// the @Nested test classes.
-			if (matchingChildrenType == ClassBasedTestDescriptor.class) {
-				Stream.concat(nonMatchingTestDescriptors.stream(), orderedTestDescriptors.stream())//
-						.forEach(parentTestDescriptor::addChild);
-			}
-			// Otherwise, we add the ordered descriptors before the non-matching descriptors,
-			// which is the case when we are ordering test methods. In other words, local
-			// test methods always get added before @Nested test classes.
-			else {
-				Stream.concat(orderedTestDescriptors.stream(), nonMatchingTestDescriptors.stream())//
-						.forEach(parentTestDescriptor::addChild);
-			}
+			parentTestDescriptor.orderChildren(children -> {
+				Stream<TestDescriptor> nonMatchingTestDescriptors = children.stream()//
+						.filter(childTestDescriptor -> !matchingChildrenType.isInstance(childTestDescriptor));
+
+				descriptorWrapperOrderer.orderWrappers(matchingDescriptorWrappers);
+
+				Stream<TestDescriptor> orderedTestDescriptors = matchingDescriptorWrappers.stream()//
+						.map(AbstractAnnotatedDescriptorWrapper::getTestDescriptor);
+
+				// If we are ordering children of type ClassBasedTestDescriptor, that means we
+				// are ordering top-level classes or @Nested test classes. Thus, the
+				// nonMatchingTestDescriptors list is either empty (for top-level classes) or
+				// contains only local test methods (for @Nested classes) which must be executed
+				// before tests in @Nested test classes. So we add the test methods before adding
+				// the @Nested test classes.
+				if (matchingChildrenType == ClassBasedTestDescriptor.class) {
+					return Stream.concat(nonMatchingTestDescriptors, orderedTestDescriptors)//
+							.collect(toList());
+				}
+				// Otherwise, we add the ordered descriptors before the non-matching descriptors,
+				// which is the case when we are ordering test methods. In other words, local
+				// test methods always get added before @Nested test classes.
+				else {
+					return Stream.concat(orderedTestDescriptors, nonMatchingTestDescriptors)//
+							.collect(toList());
+				}
+			});
 		}
 
 		// Recurse through the children in order to support ordering for @Nested test classes.
@@ -167,7 +150,32 @@ private boolean canOrderWrappers() {
 		}
 
 		private void orderWrappers(List<WRAPPER> wrappers) {
-			this.orderingAction.accept(wrappers);
+			List<WRAPPER> orderedWrappers = new ArrayList<>(wrappers);
+			this.orderingAction.accept(orderedWrappers);
+			Map<Object, Integer> distinctWrappersToIndex = distinctWrappersToIndex(orderedWrappers);
+
+			int difference = orderedWrappers.size() - wrappers.size();
+			int distinctDifference = distinctWrappersToIndex.size() - wrappers.size();
+			if (difference > 0) { // difference >= distinctDifference
+				logDescriptorsAddedWarning(difference);
+			}
+			if (distinctDifference < 0) { // distinctDifference <= difference
+				logDescriptorsRemovedWarning(distinctDifference);
+			}
+
+			wrappers.sort(comparing(wrapper -> distinctWrappersToIndex.getOrDefault(wrapper, -1)));
+		}
+
+		private Map<Object, Integer> distinctWrappersToIndex(List<?> wrappers) {
+			Map<Object, Integer> toIndex = new HashMap<>();
+			for (int i = 0; i < wrappers.size(); i++) {
+				// Avoid ClassCastException if a misbehaving ordering action added a non-WRAPPER
+				Object wrapper = wrappers.get(i);
+				if (!toIndex.containsKey(wrapper)) {
+					toIndex.put(wrapper, i);
+				}
+			}
+			return toIndex;
 		}
 
 		private void logDescriptorsAddedWarning(int number) {

junit-platform-engine/src/main/java/org/junit/platform/engine/TestDescriptor.java

@@ -10,12 +10,16 @@
 
 package org.junit.platform.engine;
 
+import static org.apiguardian.api.API.Status.EXPERIMENTAL;
 import static org.apiguardian.api.API.Status.STABLE;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.LinkedHashSet;
+import java.util.List;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.UnaryOperator;
 
 import org.apiguardian.api.API;
 import org.junit.platform.commons.util.Preconditions;
@@ -172,6 +176,39 @@ default Set<? extends TestDescriptor> getDescendants() {
 	 */
 	void removeFromHierarchy();
 
+	/**
+	 * Order the children from this descriptor.
+	 *
+	 * <p>The {@code orderer} is provided a modifiable list of child test
+	 * descriptors in this test descriptor; never {@code null}. The
+	 * {@code orderer} must return a list containing the same descriptors in any
+	 * order; potentially the same list, but never {@code null}. If descriptors
+	 * were added or removed, an exception is thrown.
+	 *
+	 * @param orderer a unary operator to order the children of this test
+	 * descriptor.
+	 */
+	@API(since = "5.12", status = EXPERIMENTAL)
+	default void orderChildren(UnaryOperator<List<TestDescriptor>> orderer) {
+		Preconditions.notNull(orderer, "orderer must not be null");
+		Set<? extends TestDescriptor> originalChildren = getChildren();
+		List<TestDescriptor> suggestedOrder = orderer.apply(new ArrayList<>(originalChildren));
+		Preconditions.notNull(suggestedOrder, "orderer may not return null");
+
+		Set<? extends TestDescriptor> orderedChildren = new LinkedHashSet<>(suggestedOrder);
+		boolean unmodified = originalChildren.equals(orderedChildren);
+		Preconditions.condition(unmodified && originalChildren.size() == suggestedOrder.size(),
+			"orderer may not add or remove test descriptors");
+
+		suggestedOrder.stream() //
+				.distinct() //
+				.filter(originalChildren::contains)//
+				.forEach(testDescriptor -> {
+					removeChild(testDescriptor);
+					addChild(testDescriptor);
+				});
+	}
+
 	/**
 	 * Determine if this descriptor is a <em>root</em> descriptor.
 	 *

junit-platform-engine/src/main/java/org/junit/platform/engine/support/descriptor/AbstractTestDescriptor.java

@@ -13,10 +13,13 @@
 import static java.util.Collections.emptySet;
 import static org.apiguardian.api.API.Status.STABLE;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.LinkedHashSet;
+import java.util.List;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.UnaryOperator;
 
 import org.apiguardian.api.API;
 import org.junit.platform.commons.util.Preconditions;
@@ -147,6 +150,24 @@ public void removeFromHierarchy() {
 		this.children.clear();
 	}
 
+	/**
+	 * {@inheritDoc}
+	 */
+	@Override
+	public void orderChildren(UnaryOperator<List<TestDescriptor>> orderer) {
+		Preconditions.notNull(orderer, "orderer must not be null");
+		List<TestDescriptor> suggestedOrder = orderer.apply(new ArrayList<>(this.children));
+		Preconditions.notNull(suggestedOrder, "orderer may not return null");
+
+		Set<? extends TestDescriptor> orderedChildren = new LinkedHashSet<>(suggestedOrder);
+		boolean unmodified = this.children.equals(orderedChildren);
+		Preconditions.condition(unmodified && this.children.size() == suggestedOrder.size(),
+			"orderer may not add or remove test descriptors");
+
+		this.children.clear();
+		this.children.addAll(orderedChildren);
+	}
+
 	@Override
 	public Optional<? extends TestDescriptor> findByUniqueId(UniqueId uniqueId) {
 		Preconditions.notNull(uniqueId, "UniqueId must not be null");