content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.javaagent.instrumentation.executors;
  
  import static org.assertj.core.api.Assertions.assertThat;
  
  import io.opentelemetry.api.trace.SpanKind;
  import io.opentelemetry.instrumentation.testing.junit.AgentInstrumentationExtension;
  import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.Executors;
  import java.util.function.Supplier;
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.extension.RegisterExtension;
  
  class CompletableFutureTest {
  
    @RegisterExtension
    static final InstrumentationExtension testing = AgentInstrumentationExtension.create();
  
    @Test
    void multipleCallbacks() {
      ExecutorService executor = Executors.newSingleThreadExecutor();
      ExecutorService executor2 = Executors.newSingleThreadExecutor();
  
      String result;
      try {
        result =
            testing.runWithSpan(
                "parent",
                () ->
                    CompletableFuture.supplyAsync(
                            () -> {
                              testing.runWithSpan("supplier", () -> {});
                              try {
                                Thread.sleep(1);
                              } catch (InterruptedException e) {
                                Thread.currentThread().interrupt();
                                throw new AssertionError(e);
                              }
                              return "a";
                            },
                            executor)
                        .thenCompose(
                            s -> CompletableFuture.supplyAsync(new AppendingSupplier(s), executor2))
                        .thenApply(
                            s -> {
                              testing.runWithSpan("function", () -> {});
                              return s + "c";
                            })
                        .get());
      } catch (Exception e) {
        throw new AssertionError(e);
      }
  
      assertThat(result).isEqualTo("abc");
  
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName("supplier").hasKind(SpanKind.INTERNAL).hasParent(trace.getSpan(0)),
                  span ->
                      span.hasName("appendingSupplier")
                          .hasKind(SpanKind.INTERNAL)
                          .hasParent(trace.getSpan(0)),
                  span ->
                      span.hasName("function")
                          .hasKind(SpanKind.INTERNAL)
                          .hasParent(trace.getSpan(0))));
  
      executor.shutdown();
      executor2.shutdown();
    }
  
    @Test
    void supplyAsync() {
      CompletableFuture<String> future =
          testing.runWithSpan(
              "parent",
              () -> CompletableFuture.supplyAsync(() -> testing.runWithSpan("child", () -> "done")));
  
      assertThat(future.join()).isEqualTo("done");
  
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName("child").hasKind(SpanKind.INTERNAL).hasParent(trace.getSpan(0))));
    }
  
    @Test
    void thenApply() {
      CompletableFuture<String> future =
          testing.runWithSpan(
              "parent",
              () ->
                  CompletableFuture.supplyAsync(() -> "done")
                      .thenApply(result -> testing.runWithSpan("child", () -> result)));
  
      assertThat(future.join()).isEqualTo("done");
  
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName("child").hasKind(SpanKind.INTERNAL).hasParent(trace.getSpan(0))));
    }
  
    @Test
    void thenApplyAsync() {
      CompletableFuture<String> future =
          testing.runWithSpan(
              "parent",
              () ->
                  CompletableFuture.supplyAsync(() -> "done")
                      .thenApplyAsync(result -> testing.runWithSpan("child", () -> result)));
  
      assertThat(future.join()).isEqualTo("done");
  
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName("child").hasKind(SpanKind.INTERNAL).hasParent(trace.getSpan(0))));
    }
  
    @Test
    void thenCompose() {
      CompletableFuture<String> future =
          testing.runWithSpan(
              "parent",
              () ->
                  CompletableFuture.supplyAsync(() -> "done")
                      .thenCompose(
                          result ->
                              CompletableFuture.supplyAsync(
                                  () -> testing.runWithSpan("child", () -> result))));
  
      assertThat(future.join()).isEqualTo("done");
  
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName("child").hasKind(SpanKind.INTERNAL).hasParent(trace.getSpan(0))));
    }
  
    @Test
    void thenComposeAsync() {
      CompletableFuture<String> future =
          testing.runWithSpan(
              "parent",
              () ->
                  CompletableFuture.supplyAsync(() -> "done")
                      .thenComposeAsync(
                          result ->
                              CompletableFuture.supplyAsync(
                                  () -> testing.runWithSpan("child", () -> result))));
  
      assertThat(future.join()).isEqualTo("done");
  
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName("child").hasKind(SpanKind.INTERNAL).hasParent(trace.getSpan(0))));
    }
  
    @Test
    void thenComposeAndApply() {
      CompletableFuture<String> future =
          testing.runWithSpan(
              "parent",
              () ->
                  CompletableFuture.supplyAsync(() -> "do")
                      .thenCompose(result -> CompletableFuture.supplyAsync(() -> result + "ne"))
                      .thenApplyAsync(result -> testing.runWithSpan("child", () -> result)));
  
      assertThat(future.join()).isEqualTo("done");
  
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName("child").hasKind(SpanKind.INTERNAL).hasParent(trace.getSpan(0))));
    }
  
    static final class AppendingSupplier implements Supplier<String> {
  
      private final String letter;
  
      AppendingSupplier(String letter) {
        this.letter = letter;
      }
  
      @Override
      public String get() {
        testing.runWithSpan("appendingSupplier", () -> {});
        return letter + "b";
      }
    }
  }
  