content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.javaagent.instrumentation.pulsar.v2_8;
  
  import static io.opentelemetry.instrumentation.testing.util.TelemetryDataUtil.orderByRootSpanName;
  
  import io.opentelemetry.api.trace.SpanKind;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.TimeUnit;
  import org.apache.pulsar.client.api.Message;
  import org.apache.pulsar.client.api.MessageId;
  import org.apache.pulsar.client.api.MessageListener;
  import org.apache.pulsar.client.api.Schema;
  import org.apache.pulsar.client.api.SubscriptionInitialPosition;
  import org.junit.jupiter.api.Test;
  
  class PulsarClientSuppressReceiveSpansTest extends AbstractPulsarClientTest {
  
    @Test
    void testConsumeNonPartitionedTopic() throws Exception {
      String topic = "persistent://public/default/testConsumeNonPartitionedTopic";
      CountDownLatch latch = new CountDownLatch(1);
      admin.topics().createNonPartitionedTopic(topic);
      consumer =
          client
              .newConsumer(Schema.STRING)
              .subscriptionName("test_sub")
              .topic(topic)
              .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest)
              .messageListener(
                  (MessageListener<String>)
                      (consumer, msg) -> {
                        acknowledgeMessage(consumer, msg);
                        latch.countDown();
                      })
              .subscribe();
  
      producer = client.newProducer(Schema.STRING).topic(topic).enableBatching(false).create();
  
      String msg = "test";
      MessageId msgId = testing.runWithSpan("parent", () -> producer.send(msg));
  
      latch.await(1, TimeUnit.MINUTES);
  
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName(topic + " publish")
                          .hasKind(SpanKind.PRODUCER)
                          .hasParent(trace.getSpan(0))
                          .hasAttributesSatisfyingExactly(
                              sendAttributes(topic, msgId.toString(), false)),
                  span ->
                      span.hasName(topic + " process")
                          .hasKind(SpanKind.CONSUMER)
                          .hasParent(trace.getSpan(1))
                          .hasAttributesSatisfyingExactly(
                              processAttributes(topic, msgId.toString(), false))));
    }
  
    @Test
    void testConsumeNonPartitionedTopicUsingReceive() throws Exception {
      String topic = "persistent://public/default/testConsumeNonPartitionedTopicCallReceive";
      admin.topics().createNonPartitionedTopic(topic);
      consumer =
          client
              .newConsumer(Schema.STRING)
              .subscriptionName("test_sub")
              .topic(topic)
              .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest)
              .subscribe();
      producer = client.newProducer(Schema.STRING).topic(topic).enableBatching(false).create();
  
      String msg = "test";
      MessageId msgId = testing.runWithSpan("parent", () -> producer.send(msg));
  
      Message<String> receivedMsg = consumer.receive();
      consumer.acknowledge(receivedMsg);
  
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName(topic + " publish")
                          .hasKind(SpanKind.PRODUCER)
                          .hasParent(trace.getSpan(0))
                          .hasAttributesSatisfyingExactly(
                              sendAttributes(topic, msgId.toString(), false)),
                  span ->
                      span.hasName(topic + " receive")
                          .hasKind(SpanKind.CONSUMER)
                          .hasParent(trace.getSpan(1))
                          .hasAttributesSatisfyingExactly(
                              receiveAttributes(topic, msgId.toString(), false))));
    }
  
    @Test
    void testConsumeNonPartitionedTopicUsingReceiveAsync() throws Exception {
      String topic = "persistent://public/default/testConsumeNonPartitionedTopicCallReceiveAsync";
      admin.topics().createNonPartitionedTopic(topic);
      consumer =
          client
              .newConsumer(Schema.STRING)
              .subscriptionName("test_sub")
              .topic(topic)
              .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest)
              .subscribe();
  
      producer = client.newProducer(Schema.STRING).topic(topic).enableBatching(false).create();
  
      CompletableFuture<Message<String>> result =
          consumer
              .receiveAsync()
              .whenComplete(
                  (message, throwable) -> {
                    if (message != null) {
                      testing.runWithSpan("callback", () -> acknowledgeMessage(consumer, message));
                    }
                  });
  
      String msg = "test";
      MessageId msgId = testing.runWithSpan("parent", () -> producer.send(msg));
  
      result.get(1, TimeUnit.MINUTES);
  
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName(topic + " publish")
                          .hasKind(SpanKind.PRODUCER)
                          .hasParent(trace.getSpan(0))
                          .hasAttributesSatisfyingExactly(
                              sendAttributes(topic, msgId.toString(), false)),
                  span ->
                      span.hasName(topic + " receive")
                          .hasKind(SpanKind.CONSUMER)
                          .hasParent(trace.getSpan(1))
                          .hasAttributesSatisfyingExactly(
                              receiveAttributes(topic, msgId.toString(), false)),
                  span ->
                      span.hasName("callback")
                          .hasKind(SpanKind.INTERNAL)
                          .hasParent(trace.getSpan(2))));
    }
  
    @Test
    void testConsumeNonPartitionedTopicUsingReceiveWithTimeout() throws Exception {
      String topic =
          "persistent://public/default/testConsumeNonPartitionedTopicCallReceiveWithTimeout";
      admin.topics().createNonPartitionedTopic(topic);
      consumer =
          client
              .newConsumer(Schema.STRING)
              .subscriptionName("test_sub")
              .topic(topic)
              .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest)
              .subscribe();
      producer = client.newProducer(Schema.STRING).topic(topic).enableBatching(false).create();
  
      String msg = "test";
      MessageId msgId = testing.runWithSpan("parent", () -> producer.send(msg));
  
      Message<String> receivedMsg = consumer.receive(1, TimeUnit.MINUTES);
      consumer.acknowledge(receivedMsg);
  
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName(topic + " publish")
                          .hasKind(SpanKind.PRODUCER)
                          .hasParent(trace.getSpan(0))
                          .hasAttributesSatisfyingExactly(
                              sendAttributes(topic, msgId.toString(), false)),
                  span ->
                      span.hasName(topic + " receive")
                          .hasKind(SpanKind.CONSUMER)
                          .hasParent(trace.getSpan(1))
                          .hasAttributesSatisfyingExactly(
                              receiveAttributes(topic, msgId.toString(), false))));
    }
  
    @Test
    void captureMessageHeaderAsSpanAttribute() throws Exception {
      String topic = "persistent://public/default/testCaptureMessageHeaderTopic";
      CountDownLatch latch = new CountDownLatch(1);
      admin.topics().createNonPartitionedTopic(topic);
      consumer =
          client
              .newConsumer(Schema.STRING)
              .subscriptionName("test_sub")
              .topic(topic)
              .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest)
              .messageListener(
                  (MessageListener<String>)
                      (consumer, msg) -> {
                        acknowledgeMessage(consumer, msg);
                        latch.countDown();
                      })
              .subscribe();
  
      producer = client.newProducer(Schema.STRING).topic(topic).enableBatching(false).create();
  
      String msg = "test";
      MessageId msgId =
          testing.runWithSpan(
              "parent",
              () -> producer.newMessage().value(msg).property("test-message-header", "test").send());
  
      latch.await(1, TimeUnit.MINUTES);
  
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName(topic + " publish")
                          .hasKind(SpanKind.PRODUCER)
                          .hasParent(trace.getSpan(0))
                          .hasAttributesSatisfyingExactly(
                              sendAttributes(topic, msgId.toString(), true)),
                  span ->
                      span.hasName(topic + " process")
                          .hasKind(SpanKind.CONSUMER)
                          .hasParent(trace.getSpan(1))
                          .hasAttributesSatisfyingExactly(
                              processAttributes(topic, msgId.toString(), true))));
    }
  
    @Test
    void testConsumePartitionedTopic() throws Exception {
      String topic = "persistent://public/default/testConsumePartitionedTopic";
      admin.topics().createPartitionedTopic(topic, 1);
      CountDownLatch latch = new CountDownLatch(1);
  
      consumer =
          client
              .newConsumer(Schema.STRING)
              .subscriptionName("test_sub")
              .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest)
              .topic(topic)
              .messageListener(
                  (MessageListener<String>)
                      (consumer, msg) -> {
                        acknowledgeMessage(consumer, msg);
                        latch.countDown();
                      })
              .subscribe();
  
      producer = client.newProducer(Schema.STRING).topic(topic).enableBatching(false).create();
  
      String msg = "test";
      MessageId msgId = testing.runWithSpan("parent", () -> producer.send(msg));
  
      latch.await(1, TimeUnit.MINUTES);
  
      testing.waitAndAssertTraces(
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName(topic + "-partition-0 publish")
                          .hasKind(SpanKind.PRODUCER)
                          .hasParent(trace.getSpan(0))
                          .hasAttributesSatisfyingExactly(
                              sendAttributes(topic + "-partition-0", msgId.toString(), false)),
                  span ->
                      span.hasName(topic + "-partition-0 process")
                          .hasKind(SpanKind.CONSUMER)
                          .hasParent(trace.getSpan(1))
                          .hasAttributesSatisfyingExactly(
                              processAttributes(topic + "-partition-0", msgId.toString(), false))));
    }
  
    @Test
    void testConsumeMultiTopics() throws Exception {
      String topicNamePrefix = "persistent://public/default/testConsumeMulti_";
      String topic1 = topicNamePrefix + "1";
      String topic2 = topicNamePrefix + "2";
      CountDownLatch latch = new CountDownLatch(2);
      producer = client.newProducer(Schema.STRING).topic(topic1).enableBatching(false).create();
      producer2 = client.newProducer(Schema.STRING).topic(topic2).enableBatching(false).create();
  
      MessageId msgId1 = testing.runWithSpan("parent1", () -> producer.send("test1"));
      MessageId msgId2 = testing.runWithSpan("parent2", () -> producer2.send("test2"));
  
      consumer =
          client
              .newConsumer(Schema.STRING)
              .topic(topic2, topic1)
              .subscriptionName("test_sub")
              .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest)
              .messageListener(
                  (MessageListener<String>)
                      (consumer, msg) -> {
                        acknowledgeMessage(consumer, msg);
                        latch.countDown();
                      })
              .subscribe();
  
      latch.await(1, TimeUnit.MINUTES);
  
      testing.waitAndAssertSortedTraces(
          orderByRootSpanName("parent1", "parent2"),
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent1").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName(topic1 + " publish")
                          .hasKind(SpanKind.PRODUCER)
                          .hasParent(trace.getSpan(0))
                          .hasAttributesSatisfyingExactly(
                              sendAttributes(topic1, msgId1.toString(), false)),
                  span ->
                      span.hasName(topic1 + " process")
                          .hasKind(SpanKind.CONSUMER)
                          .hasParent(trace.getSpan(1))
                          .hasAttributesSatisfyingExactly(
                              processAttributes(topic1, msgId1.toString(), false))),
          trace ->
              trace.hasSpansSatisfyingExactly(
                  span -> span.hasName("parent2").hasKind(SpanKind.INTERNAL).hasNoParent(),
                  span ->
                      span.hasName(topic2 + " publish")
                          .hasKind(SpanKind.PRODUCER)
                          .hasParent(trace.getSpan(0))
                          .hasAttributesSatisfyingExactly(
                              sendAttributes(topic2, msgId2.toString(), false)),
                  span ->
                      span.hasName(topic2 + " process")
                          .hasKind(SpanKind.CONSUMER)
                          .hasParent(trace.getSpan(1))
                          .hasAttributesSatisfyingExactly(
                              processAttributes(topic2, msgId2.toString(), false))));
    }
  }
  