Fix test class location search for bytecode enhanced test classes

Author: aloubyansky

test-framework/common/src/main/java/io/quarkus/test/common/PathTestHelper.java

@@ -18,10 +18,12 @@
 import io.quarkus.bootstrap.app.CuratedApplication;
 import io.quarkus.bootstrap.classloading.ClassPathElement;
 import io.quarkus.bootstrap.classloading.QuarkusClassLoader;
+import io.quarkus.bootstrap.model.ApplicationModel;
 import io.quarkus.bootstrap.workspace.ArtifactSources;
 import io.quarkus.bootstrap.workspace.SourceDir;
 import io.quarkus.bootstrap.workspace.WorkspaceModule;
 import io.quarkus.commons.classloading.ClassLoaderHelper;
+import io.quarkus.maven.dependency.DependencyFlags;
 import io.quarkus.paths.PathTree;
 import io.quarkus.paths.PathVisit;
 import io.quarkus.runtime.util.ClassPathUtils;
@@ -32,6 +34,7 @@
 public final class PathTestHelper {
     private static final String TARGET = "target";
     private static final Map<String, String> TEST_TO_MAIN_DIR_FRAGMENTS = new HashMap<>();
+    private static final List<String> TEST_DIRS;
 
     static {
         //region Eclipse
@@ -131,6 +134,7 @@ public final class PathTestHelper {
                         }
                     });
         }
+        TEST_DIRS = List.of(TEST_TO_MAIN_DIR_FRAGMENTS.keySet().toArray(new String[0]));
     }
 
     private PathTestHelper() {
@@ -143,32 +147,53 @@ private PathTestHelper() {
      * @return directory or JAR containing the test class
      */
     public static Path getTestClassesLocation(Class<?> testClass) {
-        String classFileName = testClass.getName().replace('.', File.separatorChar) + ".class";
-        URL resource = testClass.getClassLoader().getResource(fromClassNameToResourceName(testClass.getName()));
-
+        final String classFileName = fromClassNameToResourceName(testClass.getName());
+        URL resource = testClass.getClassLoader().getResource(classFileName);
         if (resource == null) {
             throw new IllegalStateException(
-                    "Could not find resource: " + testClass.getName() + " using class loader " + testClass.getClassLoader());
+                    "Could not find resource " + classFileName + " using class loader " + testClass.getClassLoader());
         }
         if (resource.getProtocol().equals("jar")) {
             try {
                 resource = URI.create(resource.getFile().substring(0, resource.getFile().indexOf('!'))).toURL();
                 return toPath(resource);
             } catch (MalformedURLException e) {
-                throw new RuntimeException("Failed to resolve the location of the JAR containing " + testClass, e);
+                throw new RuntimeException("Failed to resolve the location of the JAR containing " + classFileName, e);
+            }
+        }
+        if (resource.getProtocol().equals("quarkus")) {
+            // This is a bytecode enhanced class, the original class is either in the application module or a dependency
+            final ApplicationModel appModel = ((QuarkusClassLoader) testClass.getClassLoader()).getCuratedApplication()
+                    .getApplicationModel();
+
+            Path testLocation = getTestClassesDirOrNull(classFileName, appModel.getApplicationModule());
+            if (testLocation != null) {
+                return testLocation;
+            }
+
+            // JARs containing tests will most of the time be direct test scoped dependencies (e.g. Quarkus platform testsuite).
+            // Look among the direct dependencies first to optimize for the majority of cases.
+            // Depending on the amount of dependencies and their ordering, could be ~15-20 times faster.
+            testLocation = getTestClassLocationFromDepsOrNull(classFileName, appModel,
+                    DependencyFlags.DIRECT | DependencyFlags.RUNTIME_CP);
+            if (testLocation != null) {
+                return testLocation;
             }
-        } else if (resource.getProtocol().equals("quarkus")) {
-            // This is loaded with a quarkus classloader, so we can (sort of) ask it directly
-            QuarkusClassLoader qcl = (QuarkusClassLoader) testClass.getClassLoader();
-            return getTestClassesLocation(testClass, qcl.getCuratedApplication());
+            // Look among all the runtime dependencies. We need a more efficient way to handle this case.
+            // I don't think we have a test for this case.
+            testLocation = getTestClassLocationFromDepsOrNull(classFileName, appModel, DependencyFlags.RUNTIME_CP);
+            if (testLocation != null) {
+                return testLocation;
+            }
+            throw new RuntimeException("Failed to locate " + classFileName + " among the application dependencies");
         }
         Path path = toPath(resource);
         path = path.getRoot().resolve(path.subpath(0, path.getNameCount() - Path.of(classFileName).getNameCount()));
 
-        if (!isInTestDir(resource) && !path.getParent().getFileName().toString().equals(TARGET)) {
+        if (!isInTestDir(path) && !path.getParent().getFileName().toString().equals(TARGET)) {
             final StringBuilder msg = new StringBuilder();
             msg.append("The test class ").append(testClass.getName()).append(" is not located in any of the directories ");
-            var i = TEST_TO_MAIN_DIR_FRAGMENTS.keySet().iterator();
+            var i = TEST_DIRS.iterator();
             msg.append(i.next());
             while (i.hasNext()) {
                 msg.append(", ").append(i.next());
@@ -178,12 +203,43 @@ public static Path getTestClassesLocation(Class<?> testClass) {
         return path;
     }
 
+    /**
+     * Looks for a resource among the dependencies with specific flags. The method will return the first dependency
+     * providing the resource of null, if none of the dependencies provide the resource.
+     *
+     * @param classFileName classpath resource name
+     * @param appModel application model
+     * @param depFlags dependency flags
+     * @return the first dependency containing the resource or null, if none of the matching dependencies provide the resource
+     */
+    private static Path getTestClassLocationFromDepsOrNull(String classFileName, ApplicationModel appModel, int depFlags) {
+        for (var d : appModel.getDependencies(depFlags)) {
+            final Path root = d.getContentTree().apply(classFileName, PathTestHelper::getRootOrNull);
+            if (root != null) {
+                return root;
+            }
+        }
+        return null;
+    }
+
     public static Path getTestClassesLocation(Class<?> requiredTestClass, CuratedApplication curatedApplication) {
-        final WorkspaceModule module = curatedApplication.getApplicationModel().getAppArtifact().getWorkspaceModule();
+        final Path testClassesDir = getTestClassesDirOrNull(
+                ClassLoaderHelper.fromClassNameToResourceName(requiredTestClass.getName()),
+                curatedApplication.getApplicationModel().getApplicationModule());
+        return testClassesDir != null ? testClassesDir : getTestClassesLocation(requiredTestClass);
 
+    }
+
+    /**
+     * Looks for a resource in the output directories of a workspace module.
+     * If a directory containing the resource could not be found, the method will return null.
+     *
+     * @param testClassFileName classpath resource
+     * @param module workspace module
+     * @return output directory containing the resource or null, in case the resource could not be found
+     */
+    private static Path getTestClassesDirOrNull(String testClassFileName, WorkspaceModule module) {
         ArtifactSources testSources = module.getTestSources();
-        final String testClassFileName = ClassLoaderHelper
-                .fromClassNameToResourceName(requiredTestClass.getName());
         if (testSources != null) {
             PathTree paths = testSources.getOutputTree();
             var testClassesDir = paths.apply(testClassFileName, PathTestHelper::getRootOrNull);
@@ -202,21 +258,18 @@ public static Path getTestClassesLocation(Class<?> requiredTestClass, CuratedApp
                 }
             }
         }
-
         // If we got to this point, fall back to the filesystem search
         // This happens for maven source set scenarios
         // TODO getSourceClassifiers() should return the source sets in the maven case, but currently does not - see BuildIT.testCustomTestSourceSets test
-        return getTestClassesLocation(requiredTestClass);
-
+        return null;
     }
 
     private static Path getRootOrNull(PathVisit visit) {
         if (visit == null) {
             // this path does not exist in this path tree
             return null;
-        } else {
-            return visit.getRoot();
         }
+        return visit.getRoot();
     }
 
     public static void validateTestDir(Class<?> requiredTestClass, Path testClassesDir, WorkspaceModule module) {
@@ -261,7 +314,7 @@ public static Path getAppClassLocationForTestLocation(Path testClassLocationPath
                     // we should replace only the last occurrence of the fragment
                     final int i = testClassLocation.lastIndexOf(e.getKey());
                     final StringBuilder buf = new StringBuilder(testClassLocation.length());
-                    buf.append(testClassLocation.substring(0, i)).append(e.getValue());
+                    buf.append(testClassLocation, 0, i).append(e.getValue());
                     if (i + e.getKey().length() + 1 < testClassLocation.length()) {
                         buf.append(testClassLocation.substring(i + e.getKey().length()));
                     }
@@ -362,10 +415,9 @@ public static boolean isTestClass(String className, ClassLoader classLoader, Pat
         return testLocation.equals(Path.of(path));
     }
 
-    private static boolean isInTestDir(URL resource) {
-        String path = toPath(resource).toString();
-        return TEST_TO_MAIN_DIR_FRAGMENTS.keySet().stream()
-                .anyMatch(path::contains);
+    private static boolean isInTestDir(Path resource) {
+        final String path = resource.toString();
+        return TEST_DIRS.stream().anyMatch(path::contains);
     }
 
     private static Path toPath(URL resource) {