address comments

Author: jpbempel

buildSrc/call-site-instrumentation-plugin/src/test/java/datadog/trace/plugin/csi/impl/assertion/AssertBuilder.java

@@ -18,6 +18,7 @@
 import java.io.FileNotFoundException;
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
+import java.util.Collections;
 import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -39,10 +40,8 @@ public CallSiteAssert build() {
     Method enabled = null;
     Set<String> enabledArgs = null;
     Object[] enabledDeclaration = getEnabledDeclaration(targetType, interfaces);
-    if (enabledDeclaration != null) {
-      enabled = (Method) enabledDeclaration[0];
-      enabledArgs = (Set<String>) enabledDeclaration[1];
-    }
+    enabled = (Method) enabledDeclaration[0];
+    enabledArgs = (Set<String>) enabledDeclaration[1];
     return new CallSiteAssert(
         interfaces,
         getSpi(targetType),
@@ -53,8 +52,8 @@ public CallSiteAssert build() {
   }
 
   protected Set<Class<?>> getSpi(ClassOrInterfaceDeclaration type) {
-    return type.getAnnotationByName("AutoService").stream()
-        .flatMap(
+    return type.getAnnotationByName("AutoService")
+        .<Set<Class<?>>>map(
             annotation ->
                 annotation.asNormalAnnotationExpr().getPairs().stream()
                     .filter(pair -> pair.getNameAsString().equals("value"))
@@ -70,9 +69,10 @@ protected Set<Class<?>> getSpi(ClassOrInterfaceDeclaration type) {
                                             .asReferenceType()
                                             .getTypeDeclaration()
                                             .get()
-                                            .getQualifiedName())))
-        .map(AssertBuilder::loadClass)
-        .collect(Collectors.toSet());
+                                            .getQualifiedName()))
+                    .map(AssertBuilder::loadClass)
+                    .collect(Collectors.toSet()))
+        .orElse(Collections.emptySet());
   }
 
   protected Set<Class<?>> getInterfaces(ClassOrInterfaceDeclaration type) {
@@ -93,31 +93,26 @@ protected Set<Class<?>> getInterfaces(ClassOrInterfaceDeclaration type) {
   }
 
   private static Class<?> loadClass(String qualifiedName) {
-    // Try the name as-is first
-    try {
-      return Class.forName(qualifiedName);
-    } catch (ClassNotFoundException e) {
-      // Try progressively replacing dots with $ from right to left for inner classes
-      // We need to try all possible combinations
-      String current = qualifiedName;
-      int lastDot = current.lastIndexOf('.');
-      while (lastDot > 0) {
-        current = current.substring(0, lastDot) + "$" + current.substring(lastDot + 1);
-        try {
-          return Class.forName(current);
-        } catch (ClassNotFoundException ex) {
-          // Continue trying with the next dot
-          lastDot = current.lastIndexOf('.', lastDot - 1);
+    // Try progressively replacing dots with $ from right to left for inner classes
+    String current = qualifiedName;
+    int lastDot = current.lastIndexOf('.');
+    do {
+      try {
+        return Class.forName(current);
+      } catch (ClassNotFoundException e) {
+        if (lastDot <= 0) {
+          throw new RuntimeException(new ClassNotFoundException(qualifiedName));
         }
+        current = current.substring(0, lastDot) + "$" + current.substring(lastDot + 1);
+        lastDot = current.lastIndexOf('.', lastDot - 1);
       }
-      throw new RuntimeException(new ClassNotFoundException(qualifiedName));
-    }
+    } while (true);
   }
 
   protected Object[] getEnabledDeclaration(
       ClassOrInterfaceDeclaration type, Set<Class<?>> interfaces) {
     if (!interfaces.contains(CallSites.HasEnabledProperty.class)) {
-      return null;
+      return new Object[] {null, null};
     }
     MethodDeclaration isEnabled = type.getMethodsByName("isEnabled").get(0);
     MethodCallExpr enabledMethodCall =

buildSrc/call-site-instrumentation-plugin/src/test/java/datadog/trace/plugin/csi/impl/ext/IastExtensionTest.java

@@ -4,6 +4,7 @@
 import static datadog.trace.plugin.csi.util.CallSiteUtils.classNameToType;
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -30,13 +31,11 @@
 import java.util.Set;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
-import java.util.stream.Stream;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
 import org.junit.jupiter.params.ParameterizedTest;
-import org.junit.jupiter.params.provider.Arguments;
-import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.CsvSource;
 import org.objectweb.asm.Type;
 
 class IastExtensionTest extends BaseCsiPluginTest {
@@ -56,14 +55,14 @@ void setup() throws Exception {
     Files.createDirectories(srcFolder);
   }
 
-  static Stream<Arguments> testThatExtensionOnlyAppliesToIastAdvicesProvider() {
-    return Stream.of(
-        Arguments.of(CallSites.class.getName(), false),
-        Arguments.of(IastExtension.IAST_CALL_SITES_FQCN, true));
-  }
-
   @ParameterizedTest
-  @MethodSource("testThatExtensionOnlyAppliesToIastAdvicesProvider")
+  @CsvSource(
+      delimiter = '|',
+      nullValues = "null",
+      value = {
+        "datadog.trace.agent.tooling.csi.CallSites      | false",
+        "datadog.trace.agent.tooling.iast.IastCallSites | true"
+      })
   void testThatExtensionOnlyAppliesToIastAdvices(String typeName, boolean expected) {
     Type type = classNameToType(typeName);
     Type[] types = new Type[] {type};
@@ -85,9 +84,7 @@ void testThatExtensionGeneratesACallSiteWithTelemetry() throws Exception {
     CallSiteSpecification spec = buildClassSpecification(IastExtensionCallSite.class);
     AdviceGenerator generator = buildAdviceGenerator(buildDir);
     CallSiteResult result = generator.generate(spec);
-    if (!result.isSuccess()) {
-      throw new IllegalArgumentException("Error with call site " + IastExtensionCallSite.class);
-    }
+    assertTrue(result.isSuccess());
     IastExtension extension = new IastExtension();
 
     extension.apply(config, result);