Added unit test for DefineClassInstrumentation (#12483)

Co-authored-by: Lauri Tulmin <[email protected]>

Author: JonasKunz

instrumentation/internal/internal-class-loader/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/DefineClassInstrumentationTest.java

@@ -0,0 +1,94 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.internal.classloader;
+
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.same;
+import static org.mockito.Mockito.verify;
+
+import io.opentelemetry.javaagent.bootstrap.DefineClassHelper;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.security.ProtectionDomain;
+import org.apache.commons.compress.utils.IOUtils;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
+import org.mockito.Mockito;
+
+class DefineClassInstrumentationTest {
+
+  public static class DefiningClassLoader extends ClassLoader {
+
+    public DefiningClassLoader() {
+      super(null);
+    }
+
+    // Suppressing warnings to force testing of deprecated method
+    @SuppressWarnings("deprecation")
+    public Class<?> doDefineClass(byte[] b, int off, int len) {
+      return defineClass(b, off, len);
+    }
+
+    public Class<?> doDefineClass(String name, byte[] b, int off, int len) {
+      return defineClass(name, b, off, len);
+    }
+
+    public Class<?> doDefineClass(
+        String name, byte[] b, int off, int len, ProtectionDomain protectionDomain) {
+      return defineClass(name, b, off, len, protectionDomain);
+    }
+  }
+
+  @ParameterizedTest
+  @ValueSource(ints = {3, 4, 5})
+  @SuppressWarnings("DirectInvocationOnMock")
+  void ensureDefineClassInstrumented(int argCount) throws IOException {
+    String className = Dummy.class.getName();
+    String resource = className.replace('.', '/') + ".class";
+    ByteArrayOutputStream byteCodeStream = new ByteArrayOutputStream();
+    try (InputStream classfile = getClass().getResourceAsStream("/" + resource)) {
+      IOUtils.copy(classfile, byteCodeStream, 1024);
+    }
+    byte[] bytecode = byteCodeStream.toByteArray();
+
+    DefiningClassLoader cl = new DefiningClassLoader();
+
+    DefineClassHelper.Handler mockHandler = Mockito.mock(DefineClassHelper.Handler.class);
+    // We need to initialize mockHandler by invoking it early, otherwise this leads to problems
+    mockHandler.beforeDefineClass(cl, className, bytecode, 0, bytecode.length);
+    Mockito.reset(mockHandler);
+
+    DefineClassHelper.Handler originalHandler =
+        DefineClassHelper.internalSetHandlerForTests(mockHandler);
+    String expectedClassName;
+    try {
+      switch (argCount) {
+        case 3:
+          expectedClassName = null;
+          cl.doDefineClass(bytecode, 0, bytecode.length);
+          break;
+        case 4:
+          expectedClassName = className;
+          cl.doDefineClass(className, bytecode, 0, bytecode.length);
+          break;
+        case 5:
+          expectedClassName = className;
+          cl.doDefineClass(className, bytecode, 0, bytecode.length, null);
+          break;
+        default:
+          throw new IllegalStateException();
+      }
+    } finally {
+      DefineClassHelper.internalSetHandlerForTests(originalHandler);
+    }
+
+    verify(mockHandler)
+        .beforeDefineClass(
+            same(cl), eq(expectedClassName), eq(bytecode), eq(0), eq(bytecode.length));
+    verify(mockHandler).afterDefineClass(eq(null));
+  }
+}

instrumentation/internal/internal-class-loader/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/internal/classloader/Dummy.java

@@ -0,0 +1,8 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.instrumentation.internal.classloader;
+
+class Dummy {}

javaagent-bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/DefineClassHelper.java

@@ -70,5 +70,18 @@ public static void internalSetHandler(Handler handler) {
     DefineClassHelper.handler = handler;
   }
 
+  /**
+   * Only for testing. In contrast to {@link #internalSetHandler(Handler)} allows replacing the
+   * handler if it already has been set.
+   *
+   * @param handler the handler to set
+   * @return the previously active handler
+   */
+  public static Handler internalSetHandlerForTests(Handler handler) {
+    Handler oldHandler = DefineClassHelper.handler;
+    DefineClassHelper.handler = handler;
+    return oldHandler;
+  }
+
   private DefineClassHelper() {}
 }