/*
 * 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 static org.awaitility.Awaitility.await;

import io.opentelemetry.api.baggage.Baggage;
import io.opentelemetry.context.Scope;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import org.junit.jupiter.api.Test;

class ThreadPoolExecutorTest {

  @Test
  void shouldPassOriginalRunnableToBeforeAfterMethods() throws InterruptedException {
    CountDownLatch latch = new CountDownLatch(1);
    Runnable task = latch::countDown;

    RunnableCheckingThreadPoolExecutor executor = new RunnableCheckingThreadPoolExecutor(task);

    Baggage baggage = Baggage.builder().put("test", "test").build();
    try (Scope ignored = baggage.makeCurrent()) {
      executor.execute(task);
    }
    latch.await(10, TimeUnit.SECONDS);

    assertThat(executor.sameTaskBefore).isTrue();
    await().untilAsserted(() -> assertThat(executor.sameTaskAfter).isTrue());
  }

  // class is configured to be instrumented via otel.instrumentation.executors.include
  static class RunnableCheckingThreadPoolExecutor extends ThreadPoolExecutor {

    final Runnable expectedTask;

    final AtomicBoolean sameTaskBefore = new AtomicBoolean();
    final AtomicBoolean sameTaskAfter = new AtomicBoolean();

    RunnableCheckingThreadPoolExecutor(Runnable expectedTask) {
      super(1, 1, 0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>());
      this.expectedTask = expectedTask;
    }

    @Override
    protected void beforeExecute(Thread t, Runnable r) {
      sameTaskBefore.set(r == expectedTask);
    }

    @Override
    protected void afterExecute(Runnable r, Throwable t) {
      sameTaskAfter.set(r == expectedTask);
    }
  }
}
