content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.javaagent.instrumentation.opensearch.rest;
  
  import static io.opentelemetry.instrumentation.testing.junit.db.SemconvStabilityUtil.maybeStable;
  import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.equalTo;
  import static io.opentelemetry.semconv.HttpAttributes.HTTP_REQUEST_METHOD;
  import static io.opentelemetry.semconv.HttpAttributes.HTTP_RESPONSE_STATUS_CODE;
  import static io.opentelemetry.semconv.NetworkAttributes.NETWORK_PROTOCOL_VERSION;
  import static io.opentelemetry.semconv.ServerAttributes.SERVER_ADDRESS;
  import static io.opentelemetry.semconv.ServerAttributes.SERVER_PORT;
  import static io.opentelemetry.semconv.UrlAttributes.URL_FULL;
  import static io.opentelemetry.semconv.incubating.DbIncubatingAttributes.DB_OPERATION;
  import static io.opentelemetry.semconv.incubating.DbIncubatingAttributes.DB_STATEMENT;
  import static io.opentelemetry.semconv.incubating.DbIncubatingAttributes.DB_SYSTEM;
  import static org.assertj.core.api.Assertions.assertThat;
  
  import io.opentelemetry.api.trace.SpanKind;
  import io.opentelemetry.instrumentation.testing.junit.InstrumentationExtension;
  import java.io.IOException;
  import java.net.URI;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.atomic.AtomicReference;
  import org.junit.jupiter.api.AfterAll;
  import org.junit.jupiter.api.BeforeAll;
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.TestInstance;
  import org.opensearch.client.Request;
  import org.opensearch.client.Response;
  import org.opensearch.client.ResponseListener;
  import org.opensearch.client.RestClient;
  import org.opensearch.testcontainers.OpensearchContainer;
  import org.testcontainers.utility.DockerImageName;
  
  @SuppressWarnings("deprecation") // using deprecated semconv
  @TestInstance(TestInstance.Lifecycle.PER_CLASS)
  public abstract class AbstractOpenSearchRestTest {
  
    protected OpensearchContainer opensearch;
    protected RestClient client;
    protected URI httpHost;
  
    protected abstract InstrumentationExtension getTesting();
  
    protected abstract RestClient buildRestClient() throws Exception;
  
    protected abstract int getResponseStatus(Response response);
  
    @BeforeAll
    void setUp() throws Exception {
      opensearch =
          new OpensearchContainer(DockerImageName.parse("opensearchproject/opensearch:1.3.6"))
              .withSecurityEnabled();
      // limit memory usage
      opensearch.withEnv("OPENSEARCH_JAVA_OPTS", "-Xmx256m -Xms256m");
      opensearch.start();
      httpHost = URI.create(opensearch.getHttpHostAddress());
  
      client = buildRestClient();
    }
  
    @AfterAll
    void tearDown() {
      opensearch.stop();
    }
  
    @Test
    void shouldGetStatusWithTraces() throws IOException {
      Response response = client.performRequest(new Request("GET", "_cluster/health"));
      assertThat(getResponseStatus(response)).isEqualTo(200);
  
      getTesting()
          .waitAndAssertTraces(
              trace ->
                  trace.hasSpansSatisfyingExactly(
                      span ->
                          span.hasName("GET")
                              .hasKind(SpanKind.CLIENT)
                              .hasAttributesSatisfyingExactly(
                                  equalTo(maybeStable(DB_SYSTEM), "opensearch"),
                                  equalTo(maybeStable(DB_OPERATION), "GET"),
                                  equalTo(maybeStable(DB_STATEMENT), "GET _cluster/health")),
                      span ->
                          span.hasName("GET")
                              .hasKind(SpanKind.CLIENT)
                              .hasParent(trace.getSpan(0))
                              .hasAttributesSatisfyingExactly(
                                  equalTo(NETWORK_PROTOCOL_VERSION, "1.1"),
                                  equalTo(SERVER_ADDRESS, httpHost.getHost()),
                                  equalTo(SERVER_PORT, httpHost.getPort()),
                                  equalTo(HTTP_REQUEST_METHOD, "GET"),
                                  equalTo(URL_FULL, httpHost + "/_cluster/health"),
                                  equalTo(HTTP_RESPONSE_STATUS_CODE, 200L))));
    }
  
    @Test
    void shouldGetStatusAsyncWithTraces() throws Exception {
      AtomicReference<Response> requestResponse = new AtomicReference<>(null);
      AtomicReference<Exception> exception = new AtomicReference<>(null);
      CountDownLatch countDownLatch = new CountDownLatch(1);
  
      ResponseListener responseListener =
          new ResponseListener() {
            @Override
            public void onSuccess(Response response) {
              getTesting()
                  .runWithSpan(
                      "callback",
                      () -> {
                        requestResponse.set(response);
                        countDownLatch.countDown();
                      });
            }
  
            @Override
            public void onFailure(Exception e) {
              getTesting()
                  .runWithSpan(
                      "callback",
                      () -> {
                        exception.set(e);
                        countDownLatch.countDown();
                      });
            }
          };
  
      getTesting()
          .runWithSpan(
              "client",
              () -> {
                client.performRequestAsync(new Request("GET", "_cluster/health"), responseListener);
              });
      countDownLatch.await();
  
      if (exception.get() != null) {
        throw exception.get();
      }
      assertThat(getResponseStatus(requestResponse.get())).isEqualTo(200);
  
      getTesting()
          .waitAndAssertTraces(
              trace ->
                  trace.hasSpansSatisfyingExactly(
                      span -> span.hasName("client").hasKind(SpanKind.INTERNAL),
                      span ->
                          span.hasName("GET")
                              .hasKind(SpanKind.CLIENT)
                              .hasParent(trace.getSpan(0))
                              .hasAttributesSatisfyingExactly(
                                  equalTo(maybeStable(DB_SYSTEM), "opensearch"),
                                  equalTo(maybeStable(DB_OPERATION), "GET"),
                                  equalTo(maybeStable(DB_STATEMENT), "GET _cluster/health")),
                      span ->
                          span.hasName("GET")
                              .hasKind(SpanKind.CLIENT)
                              .hasParent(trace.getSpan(1))
                              .hasAttributesSatisfyingExactly(
                                  equalTo(NETWORK_PROTOCOL_VERSION, "1.1"),
                                  equalTo(SERVER_ADDRESS, httpHost.getHost()),
                                  equalTo(SERVER_PORT, httpHost.getPort()),
                                  equalTo(HTTP_REQUEST_METHOD, "GET"),
                                  equalTo(URL_FULL, httpHost + "/_cluster/health"),
                                  equalTo(HTTP_RESPONSE_STATUS_CODE, 200L)),
                      span ->
                          span.hasName("callback")
                              .hasKind(SpanKind.INTERNAL)
                              .hasParent(trace.getSpan(0))));
    }
  }
  