/*
 * Copyright The OpenTelemetry Authors
 * SPDX-License-Identifier: Apache-2.0
 */

package io.opentelemetry.extensions.trace;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import io.opentelemetry.trace.Span;
import io.opentelemetry.trace.StatusCode;
import java.util.concurrent.Callable;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;

class CurrentSpanUtilsTest {
  @Mock private Span span;

  @BeforeEach
  void setUp() {
    MockitoAnnotations.initMocks(this);
    when(span.storeInContext(any())).thenCallRealMethod();
  }

  // TODO(bdrutu): When update to junit 4.13 use assertThrows instead of this.
  private void executeRunnableAndExpectError(Runnable runnable, Throwable error) {
    boolean called = false;
    try {
      CurrentSpanUtils.withSpan(span, true, runnable).run();
    } catch (Throwable e) {
      assertThat(e).isEqualTo(error);
      called = true;
    }
    assertThat(called).isTrue();
  }

  // TODO(bdrutu): When update to junit 4.13 use assertThrows instead of this.
  private void executeCallableAndExpectError(Callable<Object> callable, Throwable error) {
    boolean called = false;
    try {
      CurrentSpanUtils.withSpan(span, true, callable).call();
    } catch (Throwable e) {
      assertThat(e).isEqualTo(error);
      called = true;
    }
    assertThat(called).isTrue();
  }

  @Test
  void withSpanRunnable() {
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
    Runnable runnable =
        () -> {
          // When we run the runnable we will have the span in the current Context.
          assertThat(getCurrentSpan()).isSameAs(span);
        };
    CurrentSpanUtils.withSpan(span, false, runnable).run();
    verify(span, never()).end();
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
  }

  @Test
  void withSpanRunnable_EndSpan() {
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
    Runnable runnable =
        () -> {
          // When we run the runnable we will have the span in the current Context.
          assertThat(getCurrentSpan()).isSameAs(span);
        };
    CurrentSpanUtils.withSpan(span, true, runnable).run();
    verify(span).end();
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
  }

  @Test
  void withSpanRunnable_WithError() {
    final AssertionError error = new AssertionError("MyError");
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
    Runnable runnable =
        () -> {
          // When we run the runnable we will have the span in the current Context.
          assertThat(getCurrentSpan()).isSameAs(span);
          throw error;
        };
    executeRunnableAndExpectError(runnable, error);
    verify(span).setStatus(StatusCode.ERROR, "MyError");
    verify(span).end();
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
  }

  @Test
  void withSpanRunnable_WithErrorNoMessage() {
    final AssertionError error = new AssertionError();
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
    Runnable runnable =
        () -> {
          // When we run the runnable we will have the span in the current Context.
          assertThat(getCurrentSpan()).isSameAs(span);
          throw error;
        };
    executeRunnableAndExpectError(runnable, error);
    verify(span).setStatus(StatusCode.ERROR, "AssertionError");
    verify(span).end();
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
  }

  @Test
  void withSpanCallable() throws Exception {
    final Object ret = new Object();
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
    Callable<Object> callable =
        () -> {
          // When we run the runnable we will have the span in the current Context.
          assertThat(getCurrentSpan()).isSameAs(span);
          return ret;
        };
    assertThat(CurrentSpanUtils.withSpan(span, false, callable).call()).isEqualTo(ret);
    verify(span, never()).end();
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
  }

  @Test
  void withSpanCallable_EndSpan() throws Exception {
    final Object ret = new Object();
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
    Callable<Object> callable =
        () -> {
          // When we run the runnable we will have the span in the current Context.
          assertThat(getCurrentSpan()).isSameAs(span);
          return ret;
        };
    assertThat(CurrentSpanUtils.withSpan(span, true, callable).call()).isEqualTo(ret);
    verify(span).end();
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
  }

  @Test
  void withSpanCallable_WithException() {
    final Exception exception = new Exception("MyException");
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
    Callable<Object> callable =
        () -> {
          // When we run the runnable we will have the span in the current Context.
          assertThat(getCurrentSpan()).isSameAs(span);
          throw exception;
        };
    executeCallableAndExpectError(callable, exception);
    verify(span).setStatus(StatusCode.ERROR, "MyException");
    verify(span).end();
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
  }

  @Test
  void withSpanCallable_WithExceptionNoMessage() {
    final Exception exception = new Exception();
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
    Callable<Object> callable =
        () -> {
          // When we run the runnable we will have the span in the current Context.
          assertThat(getCurrentSpan()).isSameAs(span);
          throw exception;
        };
    executeCallableAndExpectError(callable, exception);
    verify(span).setStatus(StatusCode.ERROR, "Exception");
    verify(span).end();
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
  }

  @Test
  void withSpanCallable_WithError() {
    final AssertionError error = new AssertionError("MyError");
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
    Callable<Object> callable =
        () -> {
          // When we run the runnable we will have the span in the current Context.
          assertThat(getCurrentSpan()).isSameAs(span);
          throw error;
        };
    executeCallableAndExpectError(callable, error);
    verify(span).setStatus(StatusCode.ERROR, "MyError");
    verify(span).end();
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
  }

  @Test
  void withSpanCallable_WithErrorNoMessage() {
    final AssertionError error = new AssertionError();
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
    Callable<Object> callable =
        () -> {
          // When we run the runnable we will have the span in the current Context.
          assertThat(getCurrentSpan()).isSameAs(span);
          throw error;
        };
    executeCallableAndExpectError(callable, error);
    verify(span).setStatus(StatusCode.ERROR, "AssertionError");
    verify(span).end();
    assertThat(getCurrentSpan().getSpanContext().isValid()).isFalse();
  }

  private static Span getCurrentSpan() {
    return Span.current();
  }
}

@@ -185,56 +185,8 @@
 at;%0A
-import static org.mockito.ArgumentMatchers.any;%0A
 impo
@@ -267,48 +267,8 @@
 ify;
-%0Aimport static org.mockito.Mockito.when;
 %0A%0Aim
@@ -410,26 +410,20 @@
 ter.api.
-BeforeEach
+Test
 ;%0Aimport
@@ -437,36 +437,52 @@
 nit.jupiter.api.
-Test
+extension.ExtendWith
 ;%0Aimport org.moc
@@ -486,20 +486,19 @@
 mockito.
-Mock
+Spy
 ;%0Aimport
@@ -514,28 +514,76 @@
 ito.
-MockitoAnnotations;%0A
+junit.jupiter.MockitoExtension;%0A%0A@ExtendWith(MockitoExtension.class)
 %0Acla
@@ -611,20 +611,19 @@
 st %7B%0A  @
-Mock
+Spy
  private
@@ -639,143 +639,8 @@
 n;%0A%0A
-  @BeforeEach%0A  void setUp() %7B%0A    MockitoAnnotations.initMocks(this);%0A    when(span.storeInContext(any())).thenCallRealMethod();%0A  %7D%0A%0A
   //

Replace mock with spy in CurrentSpanUtilsTest. (#1876)