Move runWith logic to testrunner for use from Groovy and replace one … (#3546)

* Move runWith logic to testrunner for use from Groovy and replace one usage.

* Comment

* Fix static init order

* Fix

* Style

instrumentation/akka-actor-fork-join-2.5/javaagent/src/test/groovy/AkkaExecutorInstrumentationTest.groovy

@@ -3,8 +3,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import static io.opentelemetry.instrumentation.test.utils.TraceUtils.runUnderTrace
-
 import akka.dispatch.forkjoin.ForkJoinPool
 import akka.dispatch.forkjoin.ForkJoinTask
 import io.opentelemetry.api.trace.SpanKind
@@ -44,7 +42,7 @@ class AkkaExecutorInstrumentationTest extends AgentInstrumentationSpecification
     new Runnable() {
       @Override
       void run() {
-        runUnderTrace("parent") {
+        runWithSpan("parent") {
           // this child will have a span
           def child1 = new AkkaAsyncChild()
           // this child won't
@@ -104,7 +102,7 @@ class AkkaExecutorInstrumentationTest extends AgentInstrumentationSpecification
     new Runnable() {
       @Override
       void run() {
-        runUnderTrace("parent") {
+        runWithSpan("parent") {
           try {
             for (int i = 0; i < 20; ++i) {
               // Our current instrumentation instrumentation does not behave very well

testing-common/src/main/groovy/io/opentelemetry/instrumentation/test/InstrumentationSpecification.groovy

@@ -13,10 +13,10 @@ import io.opentelemetry.context.ContextStorage
 import io.opentelemetry.instrumentation.test.asserts.InMemoryExporterAssert
 import io.opentelemetry.instrumentation.testing.InstrumentationTestRunner
 import io.opentelemetry.instrumentation.testing.util.TelemetryDataUtil
+import io.opentelemetry.instrumentation.testing.util.ThrowingSupplier
 import io.opentelemetry.sdk.metrics.data.MetricData
 import io.opentelemetry.sdk.trace.data.SpanData
 import spock.lang.Specification
-
 /**
  * Base class for test specifications that are shared between instrumentation libraries and agent.
  * The methods in this class are implemented by {@link AgentTestTrait} and
@@ -98,4 +98,28 @@ abstract class InstrumentationSpecification extends Specification {
     final Closure spec) {
     InMemoryExporterAssert.assertTraces({ testRunner().getExportedSpans() }, size, spec)
   }
+
+  /**
+   * Runs the provided {@code callback} inside the scope of an INTERNAL span with name {@code
+   * spanName}.
+   */
+  def runWithSpan(String spanName, Closure callback) {
+    return testRunner().runWithSpan(spanName, (ThrowingSupplier) callback)
+  }
+
+  /**
+   * Runs the provided {@code callback} inside the scope of an CLIENT span with name {@code
+   * spanName}.
+   */
+  def runWithClientSpan(String spanName, Closure callback) {
+    return testRunner().runWithClientSpan(spanName, (ThrowingSupplier) callback)
+  }
+
+  /**
+   * Runs the provided {@code callback} inside the scope of an CLIENT span with name {@code
+   * spanName}.
+   */
+  def runWithServerSpan(String spanName, Closure callback) {
+    return testRunner().runWithServerSpan(spanName, (ThrowingSupplier) callback)
+  }
 }

testing-common/src/main/groovy/io/opentelemetry/instrumentation/test/utils/TraceUtils.groovy

@@ -62,6 +62,8 @@ class TraceUtils {
     tracer.spanBuilder(spanName).startSpan().end()
   }
 
+  // Must create span within agent using annotation until
+  // https://github.com/open-telemetry/opentelemetry-java-instrumentation/issues/1726
   @WithSpan(value = "parent-client-span", kind = SpanKind.CLIENT)
   static <T> T runUnderParentClientSpan(Callable<T> r) {
     r.call()

testing-common/src/main/java/io/opentelemetry/instrumentation/testing/AgentTestRunner.java

@@ -10,6 +10,7 @@ import ch.qos.logback.classic.Logger;
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.api.OpenTelemetry;
 import io.opentelemetry.instrumentation.test.utils.LoggerUtils;
+import io.opentelemetry.instrumentation.testing.util.ThrowingSupplier;
 import io.opentelemetry.javaagent.testing.common.AgentTestingExporterAccess;
 import io.opentelemetry.javaagent.testing.common.TestAgentListenerAccess;
 import io.opentelemetry.sdk.metrics.data.MetricData;
@@ -35,6 +36,12 @@ public final class AgentTestRunner implements InstrumentationTestRunner {
     return INSTANCE;
   }
 
+  private final TestInstrumenters testInstrumenters;
+
+  private AgentTestRunner() {
+    testInstrumenters = new TestInstrumenters(getOpenTelemetry());
+  }
+
   @Override
   public void beforeTestClass() {
     TestAgentListenerAccess.reset();
@@ -81,5 +88,21 @@ public final class AgentTestRunner implements InstrumentationTestRunner {
     return AgentTestingExporterAccess.forceFlushCalled();
   }
 
-  private AgentTestRunner() {}
+  @Override
+  public <T, E extends Throwable> T runWithSpan(String spanName, ThrowingSupplier<T, E> callback)
+      throws E {
+    return testInstrumenters.runWithSpan(spanName, callback);
+  }
+
+  @Override
+  public <T, E extends Throwable> T runWithClientSpan(
+      String spanName, ThrowingSupplier<T, E> callback) throws E {
+    return testInstrumenters.runWithClientSpan(spanName, callback);
+  }
+
+  @Override
+  public <T, E extends Throwable> T runWithServerSpan(
+      String spanName, ThrowingSupplier<T, E> callback) throws E {
+    return testInstrumenters.runWithServerSpan(spanName, callback);
+  }
 }

testing-common/src/main/java/io/opentelemetry/instrumentation/testing/InstrumentationTestRunner.java

@@ -6,6 +6,8 @@
 package io.opentelemetry.instrumentation.testing;
 
 import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.instrumentation.testing.util.ThrowingRunnable;
+import io.opentelemetry.instrumentation.testing.util.ThrowingSupplier;
 import io.opentelemetry.sdk.metrics.data.MetricData;
 import io.opentelemetry.sdk.trace.data.SpanData;
 import java.util.List;
@@ -31,4 +33,66 @@ public interface InstrumentationTestRunner {
   List<MetricData> getExportedMetrics();
 
   boolean forceFlushCalled();
+
+  /**
+   * Runs the provided {@code callback} inside the scope of an INTERNAL span with name {@code
+   * spanName}.
+   */
+  default <E extends Exception> void runWithSpan(String spanName, ThrowingRunnable<E> callback)
+      throws E {
+    runWithSpan(
+        spanName,
+        () -> {
+          callback.run();
+          return null;
+        });
+  }
+
+  /**
+   * Runs the provided {@code callback} inside the scope of an INTERNAL span with name {@code
+   * spanName}.
+   */
+  <T, E extends Throwable> T runWithSpan(String spanName, ThrowingSupplier<T, E> callback) throws E;
+
+  /**
+   * Runs the provided {@code callback} inside the scope of an CLIENT span with name {@code
+   * spanName}.
+   */
+  default <E extends Throwable> void runWithClientSpan(
+      String spanName, ThrowingRunnable<E> callback) throws E {
+    runWithClientSpan(
+        spanName,
+        () -> {
+          callback.run();
+          return null;
+        });
+  }
+
+  /**
+   * Runs the provided {@code callback} inside the scope of an CLIENT span with name {@code
+   * spanName}.
+   */
+  <T, E extends Throwable> T runWithClientSpan(String spanName, ThrowingSupplier<T, E> callback)
+      throws E;
+
+  /**
+   * Runs the provided {@code callback} inside the scope of an CLIENT span with name {@code
+   * spanName}.
+   */
+  default <E extends Throwable> void runWithServerSpan(
+      String spanName, ThrowingRunnable<E> callback) throws E {
+    runWithServerSpan(
+        spanName,
+        () -> {
+          callback.run();
+          return null;
+        });
+  }
+
+  /**
+   * Runs the provided {@code callback} inside the scope of an CLIENT span with name {@code
+   * spanName}.
+   */
+  <T, E extends Throwable> T runWithServerSpan(String spanName, ThrowingSupplier<T, E> callback)
+      throws E;
 }

testing-common/src/main/java/io/opentelemetry/instrumentation/testing/LibraryTestRunner.java

@@ -11,6 +11,7 @@ import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.propagation.ContextPropagators;
 import io.opentelemetry.exporter.logging.LoggingSpanExporter;
+import io.opentelemetry.instrumentation.testing.util.ThrowingSupplier;
 import io.opentelemetry.sdk.OpenTelemetrySdk;
 import io.opentelemetry.sdk.common.CompletableResultCode;
 import io.opentelemetry.sdk.metrics.data.MetricData;
@@ -33,7 +34,6 @@ public final class LibraryTestRunner implements InstrumentationTestRunner {
   private static final OpenTelemetrySdk openTelemetry;
   private static final InMemorySpanExporter testExporter;
   private static boolean forceFlushCalled;
-  private static final LibraryTestRunner INSTANCE = new LibraryTestRunner();
 
   static {
     GlobalOpenTelemetry.resetForTest();
@@ -51,10 +51,18 @@ public final class LibraryTestRunner implements InstrumentationTestRunner {
             .buildAndRegisterGlobal();
   }
 
+  private static final LibraryTestRunner INSTANCE = new LibraryTestRunner();
+
   public static LibraryTestRunner instance() {
     return INSTANCE;
   }
 
+  private final TestInstrumenters testInstrumenters;
+
+  private LibraryTestRunner() {
+    testInstrumenters = new TestInstrumenters(openTelemetry);
+  }
+
   @Override
   public void beforeTestClass() {
     // just in case: if there was any test that modified the global instance, reset it
@@ -98,7 +106,23 @@ public final class LibraryTestRunner implements InstrumentationTestRunner {
     return forceFlushCalled;
   }
 
-  private LibraryTestRunner() {}
+  @Override
+  public <T, E extends Throwable> T runWithSpan(String spanName, ThrowingSupplier<T, E> callback)
+      throws E {
+    return testInstrumenters.runWithSpan(spanName, callback);
+  }
+
+  @Override
+  public <T, E extends Throwable> T runWithClientSpan(
+      String spanName, ThrowingSupplier<T, E> callback) throws E {
+    return testInstrumenters.runWithClientSpan(spanName, callback);
+  }
+
+  @Override
+  public <T, E extends Throwable> T runWithServerSpan(
+      String spanName, ThrowingSupplier<T, E> callback) throws E {
+    return testInstrumenters.runWithServerSpan(spanName, callback);
+  }
 
   private static class FlushTrackingSpanProcessor implements SpanProcessor {
     @Override

testing-common/src/main/java/io/opentelemetry/instrumentation/testing/TestInstrumenters.java

@@ -0,0 +1,75 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.instrumentation.testing;
+
+import io.opentelemetry.api.OpenTelemetry;
+import io.opentelemetry.context.Context;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
+import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
+import io.opentelemetry.instrumentation.testing.util.ThrowingSupplier;
+
+/** {@link Instrumenter}s to use when executing test code. */
+final class TestInstrumenters {
+
+  private final Instrumenter<String, Void> testInstrumenter;
+  private final Instrumenter<String, Void> testClientInstrumenter;
+  private final Instrumenter<String, Void> testServerInstrumenter;
+
+  TestInstrumenters(OpenTelemetry openTelemetry) {
+    testInstrumenter =
+        Instrumenter.<String, Void>newBuilder(openTelemetry, "test", name -> name)
+            .newInstrumenter(SpanKindExtractor.alwaysInternal());
+    testClientInstrumenter =
+        Instrumenter.<String, Void>newBuilder(openTelemetry, "test", name -> name)
+            .newInstrumenter(SpanKindExtractor.alwaysClient());
+    testServerInstrumenter =
+        Instrumenter.<String, Void>newBuilder(openTelemetry, "test", name -> name)
+            .newInstrumenter(SpanKindExtractor.alwaysServer());
+  }
+
+  /**
+   * Runs the provided {@code callback} inside the scope of an INTERNAL span with name {@code
+   * spanName}.
+   */
+  <T, E extends Throwable> T runWithSpan(String spanName, ThrowingSupplier<T, E> callback)
+      throws E {
+    return runWithInstrumenter(spanName, testInstrumenter, callback);
+  }
+
+  /**
+   * Runs the provided {@code callback} inside the scope of an CLIENT span with name {@code
+   * spanName}.
+   */
+  <T, E extends Throwable> T runWithClientSpan(String spanName, ThrowingSupplier<T, E> callback)
+      throws E {
+    return runWithInstrumenter(spanName, testClientInstrumenter, callback);
+  }
+
+  /**
+   * Runs the provided {@code callback} inside the scope of an CLIENT span with name {@code
+   * spanName}.
+   */
+  <T, E extends Throwable> T runWithServerSpan(String spanName, ThrowingSupplier<T, E> callback)
+      throws E {
+    return runWithInstrumenter(spanName, testServerInstrumenter, callback);
+  }
+
+  private static <T, E extends Throwable> T runWithInstrumenter(
+      String spanName, Instrumenter<String, Void> instrumenter, ThrowingSupplier<T, E> callback)
+      throws E {
+    Context context = instrumenter.start(Context.current(), spanName);
+    Throwable err = null;
+    try (Scope ignored = context.makeCurrent()) {
+      return callback.get();
+    } catch (Throwable t) {
+      err = t;
+      throw t;
+    } finally {
+      instrumenter.end(context, spanName, null, err);
+    }
+  }
+}

testing-common/src/main/java/io/opentelemetry/instrumentation/testing/junit/InstrumentationExtension.java

@@ -6,11 +6,7 @@
 package io.opentelemetry.instrumentation.testing.junit;
 
 import io.opentelemetry.api.OpenTelemetry;
-import io.opentelemetry.context.Context;
 import io.opentelemetry.context.ContextStorage;
-import io.opentelemetry.context.Scope;
-import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter;
-import io.opentelemetry.instrumentation.api.instrumenter.SpanKindExtractor;
 import io.opentelemetry.instrumentation.testing.InstrumentationTestRunner;
 import io.opentelemetry.instrumentation.testing.util.TelemetryDataUtil;
 import io.opentelemetry.instrumentation.testing.util.ThrowingRunnable;
@@ -32,21 +28,8 @@ public abstract class InstrumentationExtension
 
   private final InstrumentationTestRunner testRunner;
 
-  private final Instrumenter<String, Void> testInstrumenter;
-  private final Instrumenter<String, Void> testClientInstrumenter;
-  private final Instrumenter<String, Void> testServerInstrumenter;
-
   protected InstrumentationExtension(InstrumentationTestRunner testRunner) {
     this.testRunner = testRunner;
-    testInstrumenter =
-        Instrumenter.<String, Void>newBuilder(testRunner.getOpenTelemetry(), "test", name -> name)
-            .newInstrumenter(SpanKindExtractor.alwaysInternal());
-    testClientInstrumenter =
-        Instrumenter.<String, Void>newBuilder(testRunner.getOpenTelemetry(), "test", name -> name)
-            .newInstrumenter(SpanKindExtractor.alwaysClient());
-    testServerInstrumenter =
-        Instrumenter.<String, Void>newBuilder(testRunner.getOpenTelemetry(), "test", name -> name)
-            .newInstrumenter(SpanKindExtractor.alwaysServer());
   }
 
   @Override
@@ -132,12 +115,7 @@ public abstract class InstrumentationExtension
    */
   public <E extends Exception> void runWithSpan(String spanName, ThrowingRunnable<E> callback)
       throws E {
-    runWithSpan(
-        spanName,
-        () -> {
-          callback.run();
-          return null;
-        });
+    testRunner.runWithSpan(spanName, callback);
   }
 
   /**
@@ -146,7 +124,7 @@ public abstract class InstrumentationExtension
    */
   public <T, E extends Throwable> T runWithSpan(String spanName, ThrowingSupplier<T, E> callback)
       throws E {
-    return runWithInstrumenter(spanName, testInstrumenter, callback);
+    return testRunner.runWithSpan(spanName, callback);
   }
 
   /**
@@ -155,12 +133,7 @@ public abstract class InstrumentationExtension
    */
   public <E extends Throwable> void runWithClientSpan(String spanName, ThrowingRunnable<E> callback)
       throws E {
-    runWithClientSpan(
-        spanName,
-        () -> {
-          callback.run();
-          return null;
-        });
+    testRunner.runWithClientSpan(spanName, callback);
   }
 
   /**
@@ -169,7 +142,7 @@ public abstract class InstrumentationExtension
    */
   public <T, E extends Throwable> T runWithClientSpan(
       String spanName, ThrowingSupplier<T, E> callback) throws E {
-    return runWithInstrumenter(spanName, testClientInstrumenter, callback);
+    return testRunner.runWithClientSpan(spanName, callback);
   }
 
   /**
@@ -178,12 +151,7 @@ public abstract class InstrumentationExtension
    */
   public <E extends Throwable> void runWithServerSpan(String spanName, ThrowingRunnable<E> callback)
       throws E {
-    runWithServerSpan(
-        spanName,
-        () -> {
-          callback.run();
-          return null;
-        });
+    testRunner.runWithServerSpan(spanName, callback);
   }
 
   /**
@@ -192,21 +160,6 @@ public abstract class InstrumentationExtension
    */
   public <T, E extends Throwable> T runWithServerSpan(
       String spanName, ThrowingSupplier<T, E> callback) throws E {
-    return runWithInstrumenter(spanName, testServerInstrumenter, callback);
-  }
-
-  private static <T, E extends Throwable> T runWithInstrumenter(
-      String spanName, Instrumenter<String, Void> instrumenter, ThrowingSupplier<T, E> callback)
-      throws E {
-    Context context = instrumenter.start(Context.current(), spanName);
-    Throwable err = null;
-    try (Scope ignored = context.makeCurrent()) {
-      return callback.get();
-    } catch (Throwable t) {
-      err = t;
-      throw t;
-    } finally {
-      instrumenter.end(context, spanName, null, err);
-    }
+    return testRunner.runWithServerSpan(spanName, callback);
   }
 }